diff --git a/lib/parse.js b/lib/parse.js
index 145ea26..f8e35e1 100644
--- a/lib/parse.js
+++ b/lib/parse.js
@@ -304,30 +304,42 @@ const parse = (input, options = {}) => {
     push({ type: 'text', value });
   }
 
+  flattenBlocks(stack)
+  markImbalancedBraces(ast);
+  push({ type: 'eos' });
+
+  return ast;
+};
+
+module.exports = parse;
+
+function markImbalancedBraces({nodes}) {
   // Mark imbalanced braces and brackets as invalid
+  for (const node of nodes) {
+    if (node.nodes || node.invalid)
+      continue;
+
+    if (node.type === 'open') node.isOpen = true;
+    if (node.type === 'close') node.isClose = true;
+    if (!node.nodes) node.type = 'text';
+
+    node.invalid = true;
+    delete node.parent;
+  }
+}
+
+function flattenBlocks(stack) {
+  let block;
   do {
     block = stack.pop();
 
-    if (block.type !== 'root') {
-      block.nodes.forEach(node => {
-        if (!node.nodes) {
-          if (node.type === 'open') node.isOpen = true;
-          if (node.type === 'close') node.isClose = true;
-          if (!node.nodes) node.type = 'text';
-          node.invalid = true;
-        }
-      });
+    if (block.type === 'root')
+      continue;
 
-      // get the location of the block on parent.nodes (block's siblings)
-      let parent = stack[stack.length - 1];
-      let index = parent.nodes.indexOf(block);
-      // replace the (invalid) block with it's nodes
-      parent.nodes.splice(index, 1, ...block.nodes);
-    }
+    // get the location of the block on parent.nodes (block's siblings)
+    let parent = stack.at(-1);
+    let index = parent.nodes.indexOf(block);
+    // replace the (invalid) block with its nodes
+    parent.nodes.splice(index, 1, ...block.nodes);
   } while (stack.length > 0);
-
-  push({ type: 'eos' });
-  return ast;
-};
-
-module.exports = parse;
+}