How to Flatten Your Code

By Luke Simpson

February 27, 2024

Introduction

Deep levels of nested logic make functions harder to reason about. Some code review tools assess the cognitive complexity of a function and will rate it as exponentially more complex for each level of nested logic. In this post, we’re going to look at a few ways to reduce this complexity.

Method 1: Return early

At work, I would often see code that looked like this:

if (variable1 != null) {
  // ... do stuff with variable1

  if (variable2 != null) {
    // ... do stuff with variable2

    if (variable3 != null && variable4 != null) {
      // ... do stuff with variable3 and variable4
    }
  } else {
    // ... handle case if variable2 is null
  }
} else {
  // ... handle case if variable1 is null
}

Here we have three levels of nesting, which may not look so bad at first because there’s no logic in the branches, but imagine if this function was 10 times longer. Now imagine implementing a new feature here. The only thing I foresee happening with a function like this is the nesting gets worse over time, and the cognitive complexity goes up.

This code could be simplified by doing the following:

if (variable1 == null) {
  // ... handle case if variable1 is null
}
if (variable2 == null) {
  // ... handle case if variable2 is null
}
if (variable3 == null) {
  // ... handle case if variable3 is null
}
if (variable4 == null) {
  // ... handle case if variable4 is null
}
// ... do stuff with any variable

I’ve heard this pattern referred to as “return early” or “guard clauses”. More broadly, I look at it as flattening because we have flattened the structure of the code by removing the nesting.

In this example, we no longer have to worry about the null cases in the body of our function because we handle those at the very beginning. For example, if variable3 is null, we could just initialize it to a sensible value, throw an error, or return early. After the validations, we can assume the variables are safe to work with.

This also has the added benefit of not having to scroll down a long function just to see the else cases and tracking which if they belong to.

Method 2: Chaining conditions

One time I was given two, complicated flowcharts to set an affiliate cookie, one for if the cookie was already set and another if it wasn’t. I converted the flowcharts to a bunch of if...else statements that were nested four or five levels deep. I didn’t feel great about it, but the flowcharts were weird and complicated, so I figured that’s just how the code was going to look.

My PR was rejected specifically for the nested ifs. I initially pushed back and showed the flowcharts I was dealing with,

An agitated man waving his arms erratically.

but then I thought of something else I could try.

Here’s an example of what I started with:

if (condition1) {
  if (condition2) {
    if (condition3) {
      affiliateCookie = "123";
    } else {
      affiliateCookie = "12";
    }
  } else {
    if (condition3) {
      affiliateCookie = "13";
    } else {
      affiliateCookie = "none";
    }
  }
} else {
  if (condition2) {
    if (condition3) {
      affiliateCookie = "23";
    } else {
      affiliateCookie = "2";
    }
  } else {
    if (condition3) {
      affiliateCookie = "3";
    } else {
      affiliateCookie = "none";
    }
  }
}

And here is the same thing, but flattened out by chaining the conditions with &&:

if (condition1 && condition2 && condition3) {
  affiliateCookie = "123";
} else if (condition1 && condition2 && !condition3) {
  affiliateCookie = "12";
} else if (condition1 && !condition2 && condition3) {
  affiliateCookie = "13";
} else if (condition1 && !condition2 && !condition3) {
  affiliateCookie = "none";
} else if (!condition1 && condition2 && condition3) {
  affiliateCookie = "23";
} else if (!condition1 && condition2 && !condition3) {
  affiliateCookie = "2";
} else if (!condition1 && !condition2 && condition3) {
  affiliateCookie = "3";
} else if (!condition1 && !condition2 && !condition3) {
  affiliateCookie = "none";
}

So we converted the nested ifs to conditional expressions using short-circuit evaluation.

First, I think flattening the logic like this makes it a little easier to follow because each conditional expression is like taking a full path through the flowchart. But then I noticed something.

} else if (condition1 && !condition2 && !condition3) {
  affiliateCookie = "none";
}
...
} else if (!condition1 && !condition2 && !condition3) {
  affiliateCookie = "none";
}

For those two cases, the value is the same whether condition1 is true or not. Only condition2 and condition3 are relevant, so we can remove condition1. When we do that, we’re left with a redundant branch.

} else if (!condition2 && !condition3) {
  affiliateCookie = "none";
}
...
} else if (!condition2 && !condition3) {
  affiliateCookie = "none";
}

After removing that, here is what it looks like:

if (condition1 && condition2 && condition3) {
  affiliateCookie = "123";
} else if (condition1 && condition2 && !condition3) {
  affiliateCookie = "12";
} else if (condition1 && !condition2 && condition3) {
  affiliateCookie = "13";
} else if (!condition1 && condition2 && condition3) {
  affiliateCookie = "23";
} else if (!condition1 && condition2 && !condition3) {
  affiliateCookie = "2";
} else if (!condition1 && !condition2 && condition3) {
  affiliateCookie = "3";
} else if (!condition2 && !condition3) {
  affiliateCookie = "none";
}

In this case, we could simply use an else at the end,

} else {
  affiliateCookie = "none";
}

but I wanted to show the full process I used.

I ran tests as I was refactoring as well, so I could be sure my improvements worked at every step.

After this change, my PR was accepted. The real-world example was even more striking because the flowcharts were nearly identical except for one, small difference. After going through the simplification process, they nearly cancelled each other out. What started as four or five levels of nested ifs that were a page and a half long eventually turned into three, relatively short if...else blocks.

Method 3: Remove loops by using built-in array methods

For example, here is a simple case of pushing items onto an array using a loop.

let arr1 = ["red", "green"];
let arr2 = ["blue"];
for (item of arr1) {
  arr2.push(item);
}

But the same thing can be accomplished without the loop by using concat:

let arr1 = ["red", "green"];
let arr2 = ["blue"];
arr2 = arr2.concat(arr1);

Or what about filtering based on some condition? Here’s the loop version:

let arr1 = ["red", "green"];
let arr2 = [];
for (item of arr1) {
  if (item === "red") {
    arr2.push(item);
  }
}

And the built-in filter method:

let arr1 = ["red", "green"];
let arr2 = arr1.filter((item) => item === "red");

Conclusion

Flattening out your logic will simplify your code and make it more readable and maintainable.

To recap:

Validate variables at the beginning of the function and return early.
Try converting deep levels of nesting into chained conditions, then see if any simplifications can be made.
Use built-in methods where you can for array types.
Write good tests so that you don’t accidentally introduce a regression.

Support the blog! Buy a t-shirt or a mug!

Tags: