diff --git a/packages/forgetti/src/core/expand-expressions.ts b/packages/forgetti/src/core/expand-expressions.ts
index 78e3199..2a7026d 100644
--- a/packages/forgetti/src/core/expand-expressions.ts
+++ b/packages/forgetti/src/core/expand-expressions.ts
@@ -1,9 +1,9 @@
-import * as t from '@babel/types';
 import type * as babel from '@babel/core';
+import * as t from '@babel/types';
+import { isPathValid } from './utils/checks';
+import { getHookCallType } from './utils/get-hook-call-type';
 import type { ComponentNode, StateContext } from './types';
-import { getHookCallType } from './get-hook-call-type';
-import { isPathValid } from './checks';
-import unwrapNode from './unwrap-node';
+import unwrapNode from './utils/unwrap-node';
 
 function isStatementValid(path: babel.NodePath): boolean {
   if (path) {
@@ -170,8 +170,10 @@ export function expandExpressions(
         !isPathValid(p.parentPath, t.isVariableDeclarator)
       ) {
         const id = p.scope.generateUidIdentifier('hoisted');
-        statement.insertBefore(
-          t.variableDeclaration('let', [t.variableDeclarator(id, p.node)]),
+        statement.scope.registerDeclaration(
+          statement.insertBefore(
+            t.variableDeclaration('let', [t.variableDeclarator(id, p.node)]),
+          )[0],
         );
         p.replaceWith(id);
       }
@@ -189,14 +191,14 @@ export function expandExpressions(
         const hookType = getHookCallType(ctx, p);
         if (hookType === 'custom' || hookType === 'effect') {
           const id = p.scope.generateUidIdentifier('hoisted');
-          statement.insertBefore(
-            t.variableDeclaration('let', [t.variableDeclarator(id, p.node)]),
+          statement.scope.registerDeclaration(
+            statement.insertBefore(
+              t.variableDeclaration('let', [t.variableDeclarator(id, p.node)]),
+            )[0],
           );
           p.replaceWith(id);
         }
       }
     },
   });
-
-  path.scope.crawl();
 }
diff --git a/packages/forgetti/src/core/inline-expressions.ts b/packages/forgetti/src/core/inline-expressions.ts
index 91b4a7b..dfd146b 100644
--- a/packages/forgetti/src/core/inline-expressions.ts
+++ b/packages/forgetti/src/core/inline-expressions.ts
@@ -1,7 +1,7 @@
 import type * as babel from '@babel/core';
 import * as t from '@babel/types';
+import { isPathValid } from './utils/checks';
 import type { ComponentNode } from './types';
-import { isPathValid } from './checks';
 
 function isInValidExpression(path: babel.NodePath): boolean {
   let current = path.parentPath;
diff --git a/packages/forgetti/src/core/is-constant.ts b/packages/forgetti/src/core/is-constant.ts
index 644850f..86bbd1a 100644
--- a/packages/forgetti/src/core/is-constant.ts
+++ b/packages/forgetti/src/core/is-constant.ts
@@ -1,10 +1,17 @@
 import type * as babel from '@babel/core';
 import * as t from '@babel/types';
-import { isNestedExpression, isPathValid, shouldSkipNode } from './checks';
-import getForeignBindings, { isForeignBinding } from './get-foreign-bindings';
-import { getHookCallType } from './get-hook-call-type';
 import type OptimizerScope from './optimizer-scope';
 import type { ComponentNode, StateContext } from './types';
+import {
+  isNestedExpression,
+  isPathValid,
+  shouldSkipNode,
+} from './utils/checks';
+import {
+  getForeignBindings,
+  isForeignBinding,
+} from './utils/get-foreign-bindings';
+import { getHookCallType } from './utils/get-hook-call-type';
 
 interface OptimizerInstance {
   ctx: StateContext;
diff --git a/packages/forgetti/src/core/optimize-jsx.ts b/packages/forgetti/src/core/optimize-jsx.ts
index 387cb6f..b29f349 100644
--- a/packages/forgetti/src/core/optimize-jsx.ts
+++ b/packages/forgetti/src/core/optimize-jsx.ts
@@ -1,11 +1,12 @@
 import type * as babel from '@babel/core';
 import * as t from '@babel/types';
-import type { ComponentNode, StateContext } from './types';
-import getImportIdentifier from './get-import-identifier';
+import { isPathValid, shouldSkipJSX } from './utils/checks';
+import { getDescriptiveName } from './utils/get-descriptive-name';
+import { getImportIdentifier } from './utils/get-import-identifier';
+import { getRootStatementPath } from './utils/get-root-statement-path';
 import { RUNTIME_MEMO } from './imports';
-import { shouldSkipJSX, isPathValid } from './checks';
 import type { ImportDefinition } from './presets';
-import { getDescriptiveName } from './get-descriptive-name';
+import type { ComponentNode, StateContext } from './types';
 
 interface JSXReplacement {
   id: t.Identifier;
@@ -46,19 +47,7 @@ function extractJSXExpressionFromNormalAttribute(
     extractJSXExpressions(value, state, false);
   }
   if (isPathValid(value, t.isJSXExpressionContainer)) {
-    const expr = value.get('expression');
-    if (
-      isPathValid(expr, t.isJSXElement) ||
-      isPathValid(expr, t.isJSXFragment)
-    ) {
-      extractJSXExpressions(expr, state, false);
-    } else if (isPathValid(expr, t.isExpression)) {
-      const id = state.expressions.length;
-      state.expressions.push(expr.node);
-      expr.replaceWith(
-        t.memberExpression(state.source, t.numericLiteral(id), true),
-      );
-    }
+    extractJSXExpressionsFromJSXExpressionContainer(state, value);
   }
 }
 
@@ -240,15 +229,23 @@ function transformJSX(
   } else {
     body = path.node;
   }
-  path.scope.getProgramParent().push({
-    kind: 'const',
-    id: memoComponent,
-    init: t.callExpression(getImportIdentifier(ctx, path, RUNTIME_MEMO), [
-      getImportIdentifier(ctx, path, memoDefinition),
-      t.stringLiteral(memoComponent.name),
-      t.arrowFunctionExpression([state.source], body),
-    ]),
-  });
+
+  const root = getRootStatementPath(path);
+
+  root.scope.registerDeclaration(
+    root.insertBefore(
+      t.variableDeclaration('const', [
+        t.variableDeclarator(
+          memoComponent,
+          t.callExpression(getImportIdentifier(ctx, path, RUNTIME_MEMO), [
+            getImportIdentifier(ctx, path, memoDefinition),
+            t.stringLiteral(memoComponent.name),
+            t.arrowFunctionExpression([state.source], body),
+          ]),
+        ),
+      ]),
+    )[0],
+  );
 
   const attrs = [];
 
@@ -297,6 +294,5 @@ export default function optimizeJSX(
         transformJSX(ctx, p, memoDefinition);
       },
     });
-    path.scope.crawl();
   }
 }
diff --git a/packages/forgetti/src/core/optimizer-scope.ts b/packages/forgetti/src/core/optimizer-scope.ts
index 1d62d99..71e57c0 100644
--- a/packages/forgetti/src/core/optimizer-scope.ts
+++ b/packages/forgetti/src/core/optimizer-scope.ts
@@ -1,8 +1,10 @@
 import type * as babel from '@babel/core';
 import * as t from '@babel/types';
-import type { OptimizedExpression, StateContext } from './types';
-import getImportIdentifier from './get-import-identifier';
 import { RUNTIME_BRANCH, RUNTIME_CACHE, RUNTIME_REF } from './imports';
+import type { OptimizedExpression, StateContext } from './types';
+import { getImportIdentifier } from './utils/get-import-identifier';
+
+const ArrayPrototypePush = Array.prototype.push;
 
 function mergeVariableDeclaration(statements: t.Statement[]): t.Statement[] {
   let stack: t.VariableDeclarator[] = [];
@@ -10,8 +12,8 @@ function mergeVariableDeclaration(statements: t.Statement[]): t.Statement[] {
   let value: t.Statement;
   for (let i = 0, len = statements.length; i < len; i++) {
     value = statements[i];
-    if (value.type === 'VariableDeclaration' && value.kind === 'let') {
-      stack.push(...value.declarations);
+    if (t.isVariableDeclaration(value) && value.kind === 'let') {
+      ArrayPrototypePush.apply(stack, value.declarations);
     } else {
       if (stack.length) {
         newStatements.push(t.variableDeclaration('let', stack));
@@ -24,12 +26,16 @@ function mergeVariableDeclaration(statements: t.Statement[]): t.Statement[] {
 }
 
 export default class OptimizerScope {
+  // Reference to the root memo
   memo: t.Identifier | undefined;
 
+  // Reference to the root ref
   ref: t.Identifier | undefined;
 
+  // Size of memo
   indecesMemo = 0;
 
+  // Size of ref
   indecesRef = 0;
 
   ctx: StateContext;
diff --git a/packages/forgetti/src/core/optimizer.ts b/packages/forgetti/src/core/optimizer.ts
index a5c26ce..755d79d 100644
--- a/packages/forgetti/src/core/optimizer.ts
+++ b/packages/forgetti/src/core/optimizer.ts
@@ -1,13 +1,20 @@
 import type * as babel from '@babel/core';
 import * as t from '@babel/types';
-import { isNestedExpression, shouldSkipNode, isPathValid } from './checks';
-import getForeignBindings, { isForeignBinding } from './get-foreign-bindings';
-import getImportIdentifier from './get-import-identifier';
 import { RUNTIME_EQUALS } from './imports';
+import isConstant from './is-constant';
 import OptimizerScope from './optimizer-scope';
 import type { ComponentNode, OptimizedExpression, StateContext } from './types';
-import isConstant from './is-constant';
-import { getHookCallType } from './get-hook-call-type';
+import {
+  isNestedExpression,
+  isPathValid,
+  shouldSkipNode,
+} from './utils/checks';
+import {
+  getForeignBindings,
+  isForeignBinding,
+} from './utils/get-foreign-bindings';
+import { getHookCallType } from './utils/get-hook-call-type';
+import { getImportIdentifier } from './utils/get-import-identifier';
 
 function optimizedExpr(
   expr: t.Expression,
@@ -985,8 +992,9 @@ export default class Optimizer {
   optimizeExpressionStatement(
     path: babel.NodePath<t.ExpressionStatement>,
   ): void {
-    const optimized = this.optimizeExpression(path.get('expression'));
-    this.scope.push(t.expressionStatement(optimized.expr));
+    const expr = path.get('expression');
+    const optimized = this.optimizeExpression(expr);
+    expr.replaceWith(optimized.expr);
   }
 
   optimizeVariableDeclaration(
@@ -1011,18 +1019,16 @@ export default class Optimizer {
 
   optimizeReturnStatement(path: babel.NodePath<t.ReturnStatement>): void {
     if (path.node.argument) {
-      const optimized = this.optimizeExpression(
-        path.get('argument') as babel.NodePath<t.Expression>,
-      );
-      path.node.argument = optimized.expr;
+      const argument = path.get('argument') as babel.NodePath<t.Expression>;
+      const optimized = this.optimizeExpression(argument);
+      argument.replaceWith(optimized.expr);
     }
-    this.scope.push(path.node);
   }
 
   optimizeThrowStatement(path: babel.NodePath<t.ThrowStatement>): void {
-    const optimized = this.optimizeExpression(path.get('argument'));
-    path.node.argument = optimized.expr;
-    this.scope.push(path.node);
+    const argument = path.get('argument');
+    const optimized = this.optimizeExpression(argument);
+    argument.replaceWith(optimized.expr);
   }
 
   private optimizeBlock(path: babel.NodePath<t.BlockStatement>): void {
@@ -1031,7 +1037,7 @@ export default class Optimizer {
     }
     const statements = path.get('body');
     for (let i = 0, len = statements.length; i < len; i++) {
-      this.optimizeStatement(statements[i]);
+      this.optimizeStatement(statements[i], false);
     }
   }
 
@@ -1077,10 +1083,11 @@ export default class Optimizer {
   }
 
   optimizeLoopStatement(path: babel.NodePath<t.Loop>): void {
+    const body = path.get('body');
     const parent = this.scope;
     const loop = new OptimizerScope(this.ctx, path, parent, true);
     this.scope = loop;
-    this.optimizeStatement(path.get('body'), true);
+    this.optimizeStatement(body, true);
     this.scope = parent;
 
     const statements = loop.getStatements();
@@ -1090,19 +1097,20 @@ export default class Optimizer {
       this.scope.push(memoDeclaration);
     }
 
-    path.node.body = t.blockStatement(statements);
-    this.scope.push(path.node);
+    body.replaceWith(t.blockStatement(statements));
   }
 
   optimizeForXStatement(path: babel.NodePath<t.ForXStatement>): void {
-    const optimized = this.optimizeExpression(path.get('right'));
-    path.node.right = optimized.expr;
+    const right = path.get('right');
+    const optimized = this.optimizeExpression(right);
+    right.replaceWith(optimized.expr);
     this.optimizeLoopStatement(path);
   }
 
   optimizeSwitchStatement(path: babel.NodePath<t.SwitchStatement>): void {
-    const discriminant = this.optimizeExpression(path.get('discriminant'));
-    path.node.discriminant = discriminant.expr;
+    const discriminant = path.get('discriminant');
+    const newDiscriminant = this.optimizeExpression(discriminant);
+    discriminant.replaceWith(newDiscriminant.expr);
 
     const parent = this.scope;
     const cases = path.get('cases');
@@ -1113,31 +1121,31 @@ export default class Optimizer {
       this.scope = scope;
       const consequents = current.get('consequent');
       for (let k = 0, klen = consequents.length; k < klen; k++) {
-        this.optimizeStatement(consequents[k]);
+        this.optimizeStatement(consequents[k], false);
       }
       this.scope = parent;
       current.node.consequent = scope.getStatements();
     }
     this.scope = parent;
-
-    this.scope.push(path.node);
   }
 
   optimizeTryStatement(path: babel.NodePath<t.TryStatement>): void {
+    const block = path.get('block');
     const parent = this.scope;
     const tryBlock = new OptimizerScope(this.ctx, path, parent);
     this.scope = tryBlock;
-    this.optimizeBlock(path.get('block'));
+    this.optimizeBlock(block);
     this.scope = parent;
-    path.node.block = t.blockStatement(tryBlock.getStatements());
+    block.replaceWith(t.blockStatement(tryBlock.getStatements()));
 
     if (path.node.handler) {
       const handler = path.get('handler') as babel.NodePath<t.CatchClause>;
+      const body = handler.get('body');
       const catchClause = new OptimizerScope(this.ctx, handler, parent);
       this.scope = catchClause;
-      this.optimizeBlock(handler.get('body'));
+      this.optimizeBlock(body);
       this.scope = parent;
-      handler.node.body = t.blockStatement(catchClause.getStatements());
+      body.replaceWith(t.blockStatement(catchClause.getStatements()));
     }
     if (path.node.finalizer) {
       const handler = path.get('finalizer') as babel.NodePath<t.BlockStatement>;
@@ -1145,88 +1153,54 @@ export default class Optimizer {
       this.scope = finalizerBlock;
       this.optimizeBlock(handler);
       this.scope = parent;
-      path.node.finalizer = t.blockStatement(finalizerBlock.getStatements());
+      handler.replaceWith(t.blockStatement(finalizerBlock.getStatements()));
     }
-
-    this.scope.push(path.node);
   }
 
   optimizeLabeledStatement(path: babel.NodePath<t.LabeledStatement>): void {
+    const body = path.get('body');
     const parent = this.scope;
     const block = new OptimizerScope(this.ctx, path, parent);
     this.scope = block;
-    this.optimizeStatement(path.get('body'));
+    this.optimizeStatement(body, false);
     this.scope = parent;
-    path.node.body = t.blockStatement(block.getStatements());
-    this.scope.push(path.node);
+    body.replaceWith(t.blockStatement(block.getStatements()));
   }
 
-  optimizeStatement(path: babel.NodePath<t.Statement>, topBlock = false): void {
+  optimizeStatement(
+    path: babel.NodePath<t.Statement>,
+    topBlock: boolean,
+  ): void {
     if (shouldSkipNode(path.node)) {
       return;
     }
-    switch (path.type) {
-      case 'ExpressionStatement': {
-        this.optimizeExpressionStatement(
-          path as babel.NodePath<t.ExpressionStatement>,
-        );
-        break;
-      }
-      case 'VariableDeclaration': {
-        this.optimizeVariableDeclaration(
-          path as babel.NodePath<t.VariableDeclaration>,
-        );
-        break;
-      }
-      case 'ReturnStatement': {
-        this.optimizeReturnStatement(path as babel.NodePath<t.ReturnStatement>);
-        break;
-      }
-      case 'ThrowStatement': {
-        this.optimizeThrowStatement(path as babel.NodePath<t.ThrowStatement>);
-        break;
-      }
-      case 'BlockStatement': {
-        this.optimizeBlockStatement(
-          path as babel.NodePath<t.BlockStatement>,
-          topBlock,
-        );
-        break;
-      }
-      case 'IfStatement': {
-        this.optimizeIfStatement(path as babel.NodePath<t.IfStatement>);
-        break;
-      }
-      case 'ForInStatement':
-      case 'ForOfStatement': {
-        this.optimizeForXStatement(path as babel.NodePath<t.ForXStatement>);
-        break;
-      }
-      case 'ForStatement':
-      case 'WhileStatement':
-      case 'DoWhileStatement': {
-        this.optimizeLoopStatement(path as babel.NodePath<t.Loop>);
-        break;
-      }
-      case 'SwitchStatement': {
-        this.optimizeSwitchStatement(path as babel.NodePath<t.SwitchStatement>);
-        break;
-      }
-      case 'TryStatement': {
-        this.optimizeTryStatement(path as babel.NodePath<t.TryStatement>);
-        break;
-      }
-      case 'LabeledStatement': {
-        this.optimizeLabeledStatement(
-          path as babel.NodePath<t.LabeledStatement>,
-        );
-        break;
-      }
-      default: {
-        this.scope.push(path.node);
-        break;
-      }
+    if (isPathValid(path, t.isExpressionStatement)) {
+      this.optimizeExpressionStatement(path);
+    } else if (isPathValid(path, t.isVariableDeclaration)) {
+      this.optimizeVariableDeclaration(path);
+      return;
+    } else if (isPathValid(path, t.isReturnStatement)) {
+      this.optimizeReturnStatement(path);
+    } else if (isPathValid(path, t.isThrowStatement)) {
+      this.optimizeThrowStatement(path);
+    } else if (isPathValid(path, t.isBlockStatement)) {
+      this.optimizeBlockStatement(path, topBlock);
+      return;
+    } else if (isPathValid(path, t.isIfStatement)) {
+      this.optimizeIfStatement(path);
+      return;
+    } else if (isPathValid(path, t.isForXStatement)) {
+      this.optimizeForXStatement(path);
+    } else if (isPathValid(path, t.isLoop)) {
+      this.optimizeLoopStatement(path);
+    } else if (isPathValid(path, t.isSwitchStatement)) {
+      this.optimizeSwitchStatement(path);
+    } else if (isPathValid(path, t.isTryStatement)) {
+      this.optimizeTryStatement(path);
+    } else if (isPathValid(path, t.isLabeledStatement)) {
+      this.optimizeLabeledStatement(path);
     }
+    this.scope.push(path.node);
   }
 
   optimizeArrowComponent(
@@ -1235,38 +1209,29 @@ export default class Optimizer {
     path.node.body = t.isStatement(path.node.body)
       ? path.node.body
       : t.blockStatement([t.returnStatement(path.node.body)]);
-    this.optimizeBlock(path.get('body') as babel.NodePath<t.BlockStatement>);
-    path.node.body = t.blockStatement(this.scope.getStatements());
-    path.scope.crawl();
+    const body = path.get('body');
+    this.optimizeBlock(body as babel.NodePath<t.BlockStatement>);
+    body.replaceWith(t.blockStatement(this.scope.getStatements()));
   }
 
   optimizeFunctionComponent(
     path: babel.NodePath<t.FunctionExpression | t.FunctionDeclaration>,
   ): void {
-    this.optimizeBlock(path.get('body'));
-    path.node.body = t.blockStatement(this.scope.getStatements());
-    path.scope.crawl();
+    const body = path.get('body');
+    this.optimizeBlock(body);
+    body.replaceWith(t.blockStatement(this.scope.getStatements()));
   }
 
   optimize(): void {
-    switch (this.path.type) {
-      case 'ArrowFunctionExpression': {
-        this.optimizeArrowComponent(
-          this.path as babel.NodePath<t.ArrowFunctionExpression>,
-        );
-        break;
-      }
-      case 'FunctionDeclaration':
-      case 'FunctionExpression': {
-        this.optimizeFunctionComponent(
-          this.path as babel.NodePath<
-            t.FunctionExpression | t.FunctionDeclaration
-          >,
-        );
-        break;
-      }
-      default:
-        break;
+    if (isPathValid(this.path, t.isArrowFunctionExpression)) {
+      this.optimizeArrowComponent(this.path);
+    } else {
+      this.optimizeFunctionComponent(
+        this.path as babel.NodePath<
+          t.FunctionExpression | t.FunctionDeclaration
+        >,
+      );
     }
+    this.path.scope.crawl();
   }
 }
diff --git a/packages/forgetti/src/core/simplify-expressions.ts b/packages/forgetti/src/core/simplify-expressions.ts
index 79c06d5..3db2f47 100644
--- a/packages/forgetti/src/core/simplify-expressions.ts
+++ b/packages/forgetti/src/core/simplify-expressions.ts
@@ -119,5 +119,4 @@ export function simplifyExpressions(path: babel.NodePath<ComponentNode>): void {
       },
     },
   });
-  path.scope.crawl();
 }
diff --git a/packages/forgetti/src/core/checks.ts b/packages/forgetti/src/core/utils/checks.ts
similarity index 94%
rename from packages/forgetti/src/core/checks.ts
rename to packages/forgetti/src/core/utils/checks.ts
index 7c654c1..716983f 100644
--- a/packages/forgetti/src/core/checks.ts
+++ b/packages/forgetti/src/core/utils/checks.ts
@@ -1,9 +1,9 @@
-import type * as t from '@babel/types';
 import type * as babel from '@babel/core';
-import type { ComponentNode, StateContext } from './types';
+import * as t from '@babel/types';
+import type { ComponentNode, StateContext } from '../types';
 
 export function getImportSpecifierName(specifier: t.ImportSpecifier): string {
-  if (specifier.imported.type === 'Identifier') {
+  if (t.isIdentifier(specifier.imported)) {
     return specifier.imported.name;
   }
   return specifier.imported.value;
diff --git a/packages/forgetti/src/core/get-descriptive-name.ts b/packages/forgetti/src/core/utils/get-descriptive-name.ts
similarity index 100%
rename from packages/forgetti/src/core/get-descriptive-name.ts
rename to packages/forgetti/src/core/utils/get-descriptive-name.ts
diff --git a/packages/forgetti/src/core/get-foreign-bindings.ts b/packages/forgetti/src/core/utils/get-foreign-bindings.ts
similarity index 94%
rename from packages/forgetti/src/core/get-foreign-bindings.ts
rename to packages/forgetti/src/core/utils/get-foreign-bindings.ts
index 20fcb38..e35a1e9 100644
--- a/packages/forgetti/src/core/get-foreign-bindings.ts
+++ b/packages/forgetti/src/core/utils/get-foreign-bindings.ts
@@ -46,9 +46,7 @@ function getForeignBindingsFromExpression(
   }
 }
 
-export default function getForeignBindings(
-  path: babel.NodePath,
-): t.Identifier[] {
+export function getForeignBindings(path: babel.NodePath): t.Identifier[] {
   const identifiers = new Set<string>();
   path.traverse({
     ReferencedIdentifier(p) {
diff --git a/packages/forgetti/src/core/get-hook-call-type.ts b/packages/forgetti/src/core/utils/get-hook-call-type.ts
similarity index 96%
rename from packages/forgetti/src/core/get-hook-call-type.ts
rename to packages/forgetti/src/core/utils/get-hook-call-type.ts
index 81302d1..0740212 100644
--- a/packages/forgetti/src/core/get-hook-call-type.ts
+++ b/packages/forgetti/src/core/utils/get-hook-call-type.ts
@@ -1,9 +1,9 @@
 import type * as babel from '@babel/core';
 import * as t from '@babel/types';
 import { isHookName, isPathValid } from './checks';
-import type { HookIdentity } from './presets';
+import type { HookIdentity } from '../presets';
 import unwrapNode from './unwrap-node';
-import type { StateContext } from './types';
+import type { StateContext } from '../types';
 
 type HookCallType = HookIdentity | 'custom' | 'none';
 
diff --git a/packages/forgetti/src/core/get-import-identifier.ts b/packages/forgetti/src/core/utils/get-import-identifier.ts
similarity index 91%
rename from packages/forgetti/src/core/get-import-identifier.ts
rename to packages/forgetti/src/core/utils/get-import-identifier.ts
index 3ea035e..1840294 100644
--- a/packages/forgetti/src/core/get-import-identifier.ts
+++ b/packages/forgetti/src/core/utils/get-import-identifier.ts
@@ -1,7 +1,7 @@
 import type * as babel from '@babel/core';
 import * as t from '@babel/types';
-import type { ImportDefinition } from './presets';
-import type { StateContext } from './types';
+import type { ImportDefinition } from '../presets';
+import type { StateContext } from '../types';
 
 export function getImportIdentifier(
   state: StateContext,
diff --git a/packages/forgetti/src/core/utils/get-root-statement-path.ts b/packages/forgetti/src/core/utils/get-root-statement-path.ts
new file mode 100644
index 0000000..304a7ef
--- /dev/null
+++ b/packages/forgetti/src/core/utils/get-root-statement-path.ts
@@ -0,0 +1,14 @@
+import type * as babel from '@babel/core';
+import * as t from '@babel/types';
+
+export function getRootStatementPath(path: babel.NodePath): babel.NodePath {
+  let current = path.parentPath;
+  while (current) {
+    const next = current.parentPath;
+    if (next && t.isProgram(next.node)) {
+      return current;
+    }
+    current = next;
+  }
+  return path;
+}
diff --git a/packages/forgetti/src/core/unwrap-node.ts b/packages/forgetti/src/core/utils/unwrap-node.ts
similarity index 100%
rename from packages/forgetti/src/core/unwrap-node.ts
rename to packages/forgetti/src/core/utils/unwrap-node.ts
diff --git a/packages/forgetti/src/core/unwrap-path.ts b/packages/forgetti/src/core/utils/unwrap-path.ts
similarity index 100%
rename from packages/forgetti/src/core/unwrap-path.ts
rename to packages/forgetti/src/core/utils/unwrap-path.ts
diff --git a/packages/forgetti/src/index.ts b/packages/forgetti/src/index.ts
index 8daeb1c..245e1d8 100644
--- a/packages/forgetti/src/index.ts
+++ b/packages/forgetti/src/index.ts
@@ -6,7 +6,7 @@ import {
   getImportSpecifierName,
   isHookOrComponentName,
   shouldSkipNode,
-} from './core/checks';
+} from './core/utils/checks';
 import Optimizer from './core/optimizer';
 import type {
   HookRegistration,
@@ -15,8 +15,8 @@ import type {
 } from './core/presets';
 import { PRESETS } from './core/presets';
 import type { StateContext, State } from './core/types';
-import unwrapNode from './core/unwrap-node';
-import unwrapPath from './core/unwrap-path';
+import unwrapNode from './core/utils/unwrap-node';
+import unwrapPath from './core/utils/unwrap-path';
 import { expandExpressions } from './core/expand-expressions';
 import { inlineExpressions } from './core/inline-expressions';
 import { simplifyExpressions } from './core/simplify-expressions';
diff --git a/packages/forgetti/test/__snapshots__/expressions.test.ts.snap b/packages/forgetti/test/__snapshots__/expressions.test.ts.snap
index d66c452..0ba53c3 100644
--- a/packages/forgetti/test/__snapshots__/expressions.test.ts.snap
+++ b/packages/forgetti/test/__snapshots__/expressions.test.ts.snap
@@ -26,7 +26,7 @@ function Example(props) {
     _value8 = _equals2 && _equals3 && _equals4 ? _cache[7] : _cache[7] = [_value3, _value5, _value7],
     _equals6 = _$$equals(_cache, 8, _value8),
     _value9 = _equals6 ? _cache[8] : _cache[8] = _value8;
-  return _equals6 ? _cache[9] : _cache[9] = /*@forgetti jsx*/<_Example v={_value9} />;
+  return /*@forgetti jsx*/_equals6 ? _cache[9] : _cache[9] = <_Example v={_value9} />;
 }"
 `;
 
@@ -56,7 +56,7 @@ function Example(props) {
     _value8 = _equals2 && _equals3 && _equals4 ? _cache[7] : _cache[7] = [_value3, _value5, _value7],
     _equals6 = _$$equals(_cache, 8, _value8),
     _value9 = _equals6 ? _cache[8] : _cache[8] = _value8;
-  return _equals6 ? _cache[9] : _cache[9] = /*@forgetti jsx*/<_Example v={_value9} />;
+  return /*@forgetti jsx*/_equals6 ? _cache[9] : _cache[9] = <_Example v={_value9} />;
 }"
 `;
 
diff --git a/packages/forgetti/test/__snapshots__/forgetti-skip.test.ts.snap b/packages/forgetti/test/__snapshots__/forgetti-skip.test.ts.snap
index 437de8b..83f611a 100644
--- a/packages/forgetti/test/__snapshots__/forgetti-skip.test.ts.snap
+++ b/packages/forgetti/test/__snapshots__/forgetti-skip.test.ts.snap
@@ -20,7 +20,7 @@ function Example(props) {
     _value6 = _equals2 && _equals3 ? _cache[5] : _cache[5] = [_value3, _value5],
     _equals5 = _$$equals(_cache, 6, _value6),
     _value7 = _equals5 ? _cache[6] : _cache[6] = _value6;
-  return _equals5 ? _cache[7] : _cache[7] = /*@forgetti jsx*/<_Example v={_value7} />;
+  return /*@forgetti jsx*/_equals5 ? _cache[7] : _cache[7] = <_Example v={_value7} />;
 }"
 `;
 
@@ -44,7 +44,7 @@ const Example = function (props) {
     _value6 = _equals2 && _equals3 ? _cache[5] : _cache[5] = [_value3, _value5],
     _equals5 = _$$equals(_cache, 6, _value6),
     _value7 = _equals5 ? _cache[6] : _cache[6] = _value6;
-  return _equals5 ? _cache[7] : _cache[7] = /*@forgetti jsx*/<_Example v={_value7} />;
+  return /*@forgetti jsx*/_equals5 ? _cache[7] : _cache[7] = <_Example v={_value7} />;
 };"
 `;
 
@@ -68,7 +68,7 @@ const Example = props => {
     _value6 = _equals2 && _equals3 ? _cache[5] : _cache[5] = [_value3, _value5],
     _equals5 = _$$equals(_cache, 6, _value6),
     _value7 = _equals5 ? _cache[6] : _cache[6] = _value6;
-  return _equals5 ? _cache[7] : _cache[7] = /*@forgetti jsx*/<_Example v={_value7} />;
+  return /*@forgetti jsx*/_equals5 ? _cache[7] : _cache[7] = <_Example v={_value7} />;
 };"
 `;
 
diff --git a/packages/forgetti/test/__snapshots__/hooks.test.ts.snap b/packages/forgetti/test/__snapshots__/hooks.test.ts.snap
index 0088a0f..0c301cd 100644
--- a/packages/forgetti/test/__snapshots__/hooks.test.ts.snap
+++ b/packages/forgetti/test/__snapshots__/hooks.test.ts.snap
@@ -54,9 +54,9 @@ import { $$cache as _$$cache } from \\"forgetti/runtime\\";
 import { $$equals as _$$equals } from \\"forgetti/runtime\\";
 import { memo as _memo } from \\"react\\";
 import { $$memo as _$$memo } from \\"forgetti/runtime\\";
-const _Example = _$$memo(_memo, \\"_Example\\", _values => <div>{_values[0]}</div>),
-  _Example2 = _$$memo(_memo, \\"_Example2\\", _values2 => <>{_values2[0]}{_values2[1]}</>);
 import { useA, useB, useC } from 'whatever';
+const _Example = _$$memo(_memo, \\"_Example\\", _values => <div>{_values[0]}</div>);
+const _Example2 = _$$memo(_memo, \\"_Example2\\", _values2 => <>{_values2[0]}{_values2[1]}</>);
 function Example(props) {
   let _cache = _$$cache(_useMemo, 24),
     a = null,