WIP

ROADMAP.md

@@ -14,6 +14,7 @@ Type: `Feature` - `Bug` - `Enhancement` - `Refactor` - `Unknown`
 - [ ] `Enhancement` Skip memoization of useState setter functions
 - [ ] `Enhancement` When unwrapping array and object patterns, optimize the unwrapping by limiting it to component variables
 - [ ] `Enhancement` Support React.createElement calls
+- [ ] `Enhancement` Hot module reloading improvement utilizing a checksum to invalidate cache
 - [ ] `Feature` Memoize array map items
 
 ## `Feature` ESLint Plugin
@@ -30,3 +31,4 @@ The following are unknowns that need to be researched and tested.
 - [ ] Assignment expressions in JSX `<div>{(a = 1)} {a}</div>` - is the order of execution guaranteed?
 - [ ] Source maps - is it possible to generate source maps for the transformed code?
 - [ ] Hot reloading - how does hot reloading work with the transformed code?
+- [ ] Memoization of values declared after the first return statement

packages/compiler/src/ast-factories/convert-statement-to-segment-callable.ts

@@ -0,0 +1,80 @@
+import * as t from "@babel/types";
+import { DEFAULT_SEGMENT_CALLABLE_VARIABLE_NAME } from "~/utils/constants";
+
+export function convertStatementToSegmentCallable(
+  statement: babel.NodePath<babel.types.Statement>,
+  {
+    initialValue,
+    performReplacement = true,
+    cacheNullValue,
+    segmentCallableId = statement.scope.generateUidIdentifier(
+      DEFAULT_SEGMENT_CALLABLE_VARIABLE_NAME
+    ),
+  }: {
+    initialValue?: t.Expression;
+    performReplacement?: boolean;
+    cacheNullValue?: t.Expression;
+    segmentCallableId?: t.Identifier;
+  }
+) {
+  const parentDeclaration = statement.find((p) =>
+    p.isVariableDeclaration()
+  ) as babel.NodePath<babel.types.VariableDeclaration> | null;
+
+  const makeSegmentCallable = (statements: t.Statement[]) => {
+    return t.variableDeclaration("const", [
+      t.variableDeclarator(
+        segmentCallableId,
+        t.arrowFunctionExpression(
+          [],
+          t.blockStatement(
+            statements.concat(
+              cacheNullValue ? [t.returnStatement(cacheNullValue)] : []
+            )
+          )
+        )
+      ),
+    ]);
+  };
+
+  let replacements: t.Node[] | null = null;
+  if (parentDeclaration) {
+    const newKind = parentDeclaration.node.kind === "var" ? "var" : "let";
+    const newDeclaration = t.variableDeclaration(
+      newKind,
+      parentDeclaration.node.declarations.map((declaration) => {
+        return t.variableDeclarator(declaration.id, initialValue);
+      })
+    );
+
+    const assignmentExpressionStatements = parentDeclaration.node.declarations
+      .map((declarator) => {
+        return declarator.init
+          ? t.expressionStatement(
+              t.assignmentExpression("=", declarator.id, declarator.init)
+            )
+          : null;
+      })
+      .filter((v): v is t.ExpressionStatement => Boolean(v));
+
+    if (performReplacement) {
+      replacements = [
+        newDeclaration,
+        makeSegmentCallable(assignmentExpressionStatements),
+      ];
+    }
+  } else {
+    replacements = [makeSegmentCallable([statement.node])];
+  }
+
+  const newPaths =
+    performReplacement && replacements
+      ? statement.replaceWithMultiple(replacements)
+      : null;
+
+  return {
+    segmentCallableId,
+    replacements,
+    newPaths,
+  };
+}

packages/compiler/src/ast-factories/make-unwrapped-declarations.ts

@@ -10,9 +10,14 @@ export function makeUnwrappedDeclarations(
   const unwrappedEntries = unwrapPatternAssignment(id, tempVariableId);
 
   const unwrappedDeclarations = unwrappedEntries.map((entry) => {
-    return t.variableDeclaration(kind, [
-      t.variableDeclarator(entry.id, entry.value),
-    ]);
+    const binding = id.scope.getBinding(entry.name);
+
+    return [
+      t.variableDeclaration(kind, [
+        t.variableDeclarator(entry.id, entry.value),
+      ]),
+      binding,
+    ] as const;
   });
 
   return {

packages/compiler/src/classes/Component.ts

@@ -3,6 +3,7 @@ import { Binding } from "@babel/traverse";
 import * as t from "@babel/types";
 import {
   DEFAULT_CACHE_COMMIT_VARIABLE_NAME,
+  DEFAULT_CACHE_NULL_VARIABLE_NAME,
   DEFAULT_CACHE_VARIABLE_NAME,
   RUNTIME_MODULE_CREATE_CACHE_HOOK_NAME,
 } from "~/utils/constants";
@@ -13,6 +14,12 @@ import { unwrapJsxExpressions } from "~/utils/unwrap-jsx-expressions";
 import { getFunctionParent } from "../utils/get-function-parent";
 import { ComponentSideEffect } from "./ComponentSideEffect";
 import { ComponentVariable } from "./ComponentVariable";
+import {
+  ComponentMutableSegment,
+  SegmentTransformationResult,
+} from "./ComponentMutableSegment";
+import { hasHookCall } from "~/utils/is-hook-call";
+import { convertStatementToSegmentCallable } from "~/ast-factories/convert-statement-to-segment-callable";
 
 export class Component {
   private sideEffects = new Map<
@@ -26,6 +33,11 @@ export class Component {
   >();
   private cacheValueIdentifier: t.Identifier;
   private cacheCommitIdentifier: t.Identifier;
+  private cacheNullIdentifier: t.Identifier;
+  private mapOfReturnStatementToReferencedBindings = new Map<
+    babel.NodePath<babel.types.ReturnStatement>,
+    Binding[]
+  >();
 
   constructor(public path: babel.NodePath<babel.types.Function>) {
     path.assertFunction();
@@ -37,12 +49,21 @@ export class Component {
     this.cacheCommitIdentifier = path.scope.generateUidIdentifier(
       DEFAULT_CACHE_COMMIT_VARIABLE_NAME
     );
+
+    this.cacheNullIdentifier = path.scope.generateUidIdentifier(
+      DEFAULT_CACHE_NULL_VARIABLE_NAME
+    );
   }
 
   computeComponentVariables() {
     this.prepareComponentBody();
     getReturnsOfFunction(this.path).forEach((returnPath) => {
-      const bindings = getReferencedVariablesInside(returnPath);
+      const bindings = getReferencedVariablesInside(returnPath).filter(
+        (binding) => this.isBindingInComponentScope(binding)
+      );
+
+      this.mapOfReturnStatementToReferencedBindings.set(returnPath, bindings);
+
       bindings.forEach((binding) => {
         this.addComponentVariable(binding);
       });
@@ -81,7 +102,8 @@ export class Component {
     const name = binding.identifier.name;
 
     if (this.hasComponentVariable(name)) {
-      return this.getComponentVariable(name);
+      const componentVariable = this.getComponentVariable(name)!;
+      return componentVariable;
     }
 
     const componentVariable = new ComponentVariable(
@@ -91,6 +113,7 @@ export class Component {
     );
 
     this.componentVariables.set(name, componentVariable);
+
     componentVariable.unwrapAssignmentPatterns();
     componentVariable.computeDependencyGraph();
 
@@ -128,30 +151,185 @@ export class Component {
     return t.cloneNode(this.cacheCommitIdentifier);
   }
 
-  applyModification() {
+  applyTransformation() {
     const cacheVariableDeclaration = this.makeCacheVariableDeclaration();
 
     const body = this.path.get("body");
 
-    this.componentVariables.forEach((componentVariable) => {
-      componentVariable.applyModification();
-    });
+    if (!body.isBlockStatement()) {
+      return;
+    }
+
+    const statementsToMutableSegmentMap = new Map<
+      babel.NodePath<babel.types.Statement>,
+      ComponentMutableSegment
+    >();
 
-    if (body.isBlockStatement()) {
-      body.unshiftContainer("body", cacheVariableDeclaration);
+    const statementsMapSet = (segment: ComponentMutableSegment) => {
+      const parent = segment.getParentStatement();
+      if (parent) {
+        statementsToMutableSegmentMap.set(parent, segment);
+      }
+    };
+
+    this.componentVariables.forEach(statementsMapSet);
+    this.sideEffects.forEach(statementsMapSet);
+
+    const statements = body.get("body");
+
+    let lastStatementWithHookCallIdx = -1;
+
+    for (let i = statements.length - 1; i >= 0; i--) {
+      const currentStatement = statements[i]!;
+      if (hasHookCall(currentStatement, this.path)) {
+        lastStatementWithHookCallIdx = i;
+        break;
+      }
     }
 
-    this.sideEffects.forEach((sideEffect) => {
-      sideEffect.applyModification();
-    });
+    const clonedStatements = statements.slice();
+
+    const getLastComponentVariableStatementIdx = (bindings: Binding[]) => {
+      return bindings.reduce((lastStatement, binding) => {
+        const componentVariable = this.getComponentVariable(
+          binding.identifier.name
+        );
+        if (componentVariable) {
+          const statement = componentVariable.getParentStatement();
+          if (statement) {
+            const index = clonedStatements.indexOf(statement);
+            if (index > lastStatement) {
+              return index;
+            }
+          }
+        }
+
+        return lastStatement;
+      }, -1);
+    };
+
+    const transformStatement = (
+      statement: babel.NodePath<babel.types.Statement>
+    ) => {
+      if (statement.isReturnStatement()) {
+        return null;
+      }
 
-    const returns = getReturnsOfFunction(this.path);
+      const returnDescendant = Array.from(
+        this.mapOfReturnStatementToReferencedBindings.keys()
+      ).find((returnStatement) => returnStatement.isDescendant(statement));
+      const shouldPerformReplacement = !returnDescendant;
+
+      const segment = shouldPerformReplacement
+        ? statementsToMutableSegmentMap.get(statement)
+        : null;
+
+      if (segment) {
+        const transformationResult = segment.applyTransformation(
+          shouldPerformReplacement
+        );
+        if (transformationResult) {
+          return { ...transformationResult, returnDescendant };
+        } else {
+          return null;
+        }
+      }
 
-    returns.forEach((returnPath) => {
-      returnPath.insertBefore(
-        t.expressionStatement(t.callExpression(this.cacheCommitIdentifier, []))
-      );
+      const hadHookCall = hasHookCall(statement, this.path);
+
+      this.componentVariables.forEach((componentVariable) => {
+        if (
+          componentVariable.path.isDescendant(statement) &&
+          statement !== componentVariable.path
+        ) {
+          const transformation = componentVariable.applyTransformation();
+          const callStatement = this.makeSegmentCallStatement(transformation);
+          if (callStatement && transformation) {
+            const lastPath = transformation.newPaths?.pop();
+            lastPath?.insertAfter(callStatement);
+          }
+        }
+      });
+
+      const { segmentCallableId, newPaths, replacements } =
+        convertStatementToSegmentCallable(statement, {
+          cacheNullValue: returnDescendant
+            ? this.cacheNullIdentifier
+            : undefined,
+        });
+
+      return {
+        dependencyConditions: null,
+        replacements,
+        segmentCallableId,
+        newPaths,
+        hasHookCall: hadHookCall,
+        returnDescendant,
+      };
+    };
+
+    let currentIndex = -1;
+    let currentScanIndex = 0;
+
+    const returnsSize = this.mapOfReturnStatementToReferencedBindings.size;
+    this.mapOfReturnStatementToReferencedBindings.forEach(
+      (bindings, returnPath) => {
+        currentIndex++;
+        let maxIndex = -1;
+
+        if (currentIndex === 0) {
+          maxIndex = Math.max(lastStatementWithHookCallIdx, maxIndex);
+        }
+
+        maxIndex = Math.max(
+          getLastComponentVariableStatementIdx(bindings),
+          maxIndex
+        );
+
+        const statementsToTransform: babel.NodePath<babel.types.Statement>[] =
+          [];
+
+        if (currentIndex === returnsSize - 1) {
+          statementsToTransform.push(...statements);
+        } else {
+          for (; currentScanIndex <= maxIndex; currentScanIndex++) {
+            const statement = statements.shift();
+            if (statement) {
+              statementsToTransform.push(statement);
+            }
+          }
+        }
+
+        const returnStatement = returnPath.find(
+          (p) => p.isStatement() && p.parentPath === body
+        );
+
+        statementsToTransform.forEach((statement) => {
+          const transformation = transformStatement(statement);
+          const callStatement = this.makeSegmentCallStatement(transformation);
+
+          if (transformation && callStatement) {
+            if (transformation.returnDescendant) {
+              transformation.newPaths?.[0]?.insertAfter(callStatement);
+            } else {
+              returnStatement?.insertBefore(callStatement);
+            }
+          }
+        });
+      }
+    );
+
+    this.componentVariables.forEach((componentVariable) => {
+      const transformation = componentVariable.applyTransformation();
+      const callStatement = this.makeSegmentCallStatement(transformation);
+
+      if (callStatement && transformation) {
+        const lastNewPath = transformation.newPaths?.pop();
+        lastNewPath?.insertAfter(callStatement);
+      }
     });
+
+    body.unshiftContainer("body", cacheVariableDeclaration);
   }
 
   getFunctionBlockStatement(): babel.NodePath<babel.types.BlockStatement> | null {
@@ -176,7 +354,11 @@ export class Component {
     const sizeNumber = t.numericLiteral(this.componentVariables.size);
     const declaration = t.variableDeclaration("const", [
       t.variableDeclarator(
-        t.arrayPattern([this.cacheValueIdentifier, this.cacheCommitIdentifier]),
+        t.arrayPattern([
+          this.cacheValueIdentifier,
+          this.cacheCommitIdentifier,
+          this.cacheNullIdentifier,
+        ]),
         t.callExpression(t.identifier(RUNTIME_MODULE_CREATE_CACHE_HOOK_NAME), [
           sizeNumber,
         ])
@@ -202,4 +384,66 @@ export class Component {
   __debug_getComponentVariables() {
     return this.componentVariables;
   }
+
+  isBindingInComponentScope(binding: Binding) {
+    return this.path.scope.getBinding(binding.identifier.name) === binding;
+  }
+
+  private makeSegmentCallStatement(
+    transformation: SegmentTransformationResult
+  ) {
+    if (!transformation) {
+      return null;
+    }
+
+    const {
+      dependencyConditions,
+      segmentCallableId,
+      hasHookCall,
+      updateCache,
+      returnDescendant,
+    } = transformation;
+
+    const callSegmentCallable = t.callExpression(segmentCallableId, []);
+    const updateStatements: t.Statement[] = [];
+
+    if (returnDescendant) {
+      const customCallVariable = this.path.scope.generateUidIdentifier();
+      updateStatements.push(
+        t.variableDeclaration("const", [
+          t.variableDeclarator(customCallVariable, callSegmentCallable),
+        ])
+      );
+
+      updateStatements.push(
+        // if customCallVariable not equal to null, return it
+        t.ifStatement(
+          t.binaryExpression(
+            "!==",
+            customCallVariable,
+            this.cacheNullIdentifier
+          ),
+          t.blockStatement([t.returnStatement(customCallVariable)])
+        )
+      );
+    } else {
+      updateStatements.push(t.expressionStatement(callSegmentCallable));
+    }
+
+    if (updateCache) {
+      updateStatements.push(updateCache);
+    }
+
+    const callStatementWithCondition =
+      dependencyConditions && !hasHookCall
+        ? [
+            t.ifStatement(
+              dependencyConditions,
+              t.blockStatement(updateStatements)
+            ),
+          ]
+        : updateStatements;
+
+    return callStatementWithCondition;
+  }
 }

packages/compiler/src/classes/ComponentMutableSegment.ts

@@ -1,6 +1,6 @@
 import type * as babel from "@babel/core";
 import { makeDependencyCondition } from "~/ast-factories/make-dependency-condition";
-import { isHookCall } from "~/utils/is-hook-call";
+import { hasHookCall } from "~/utils/is-hook-call";
 import { Component } from "./Component";
 import type { ComponentSideEffect } from "./ComponentSideEffect";
 import type { ComponentVariable } from "./ComponentVariable";
@@ -16,7 +16,19 @@ type ComponentMutableSegmentType =
   | typeof COMPONENT_MUTABLE_SEGMENT_COMPONENT_VARIABLE_TYPE
   | typeof COMPONENT_MUTABLE_SEGMENT_COMPONENT_SIDE_EFFECT_TYPE;
 
+export type SegmentTransformationResult = {
+  segmentCallableId: babel.types.Identifier;
+  dependencyConditions: babel.types.Expression | null;
+  newPaths: babel.NodePath<babel.types.Node>[] | null;
+  hasHookCall: boolean;
+  returnDescendant?: babel.NodePath<babel.types.Node> | null;
+  updateCache?: babel.types.Statement | null;
+  replacements: babel.types.Node[] | null;
+} | null;
+
 export abstract class ComponentMutableSegment {
+  protected appliedTransformation = false;
+
   // The side effects of this segment
   protected sideEffects = new Map<
     babel.NodePath<babel.types.Statement>,
@@ -39,10 +51,6 @@ export abstract class ComponentMutableSegment {
     this.dependencies.set(componentVariable.name, componentVariable);
   }
 
-  protected makeDependencyCondition() {
-    return makeDependencyCondition(this);
-  }
-
   hasDependencies() {
     return this.dependencies.size > 0;
   }
@@ -51,8 +59,16 @@ export abstract class ComponentMutableSegment {
     return this.dependencies;
   }
 
-  computeDependencyGraph() {
-    throw new Error("Not implemented");
+  abstract computeDependencyGraph(): void;
+
+  abstract get path(): babel.NodePath<babel.types.Node>;
+
+  abstract applyTransformation(
+    performReplacement?: boolean
+  ): SegmentTransformationResult;
+
+  protected makeDependencyCondition() {
+    return makeDependencyCondition(this);
   }
 
   getSideEffectDependencies() {
@@ -64,32 +80,15 @@ export abstract class ComponentMutableSegment {
   }
 
   hasHookCall() {
-    if (this.hasDependencies()) {
-      return false;
-    }
-
-    let hasHookCall = false;
-    this.path.traverse({
-      CallExpression: (innerPath) => {
-        if (
-          isHookCall(innerPath) &&
-          this.component.isTheFunctionParentOf(innerPath)
-        ) {
-          hasHookCall = true;
-          return;
-        }
-      },
-    });
-
-    return hasHookCall;
+    return hasHookCall(this.path, this.component.path);
   }
 
   getParentStatement() {
-    return this.path.getStatementParent();
-  }
+    const parentStatement = this.path.find(
+      (p) => p.isStatement() && p.parentPath === this.component.path.get("body")
+    ) as babel.NodePath<babel.types.Statement> | null;
 
-  get path(): babel.NodePath<babel.types.Node> {
-    throw new Error("Not implemented");
+    return parentStatement;
   }
 
   isComponentVariable(): this is ComponentVariable {

packages/compiler/src/classes/ComponentSideEffect.ts

@@ -3,10 +3,9 @@ import * as t from "@babel/types";
 import { Component } from "./Component";
 import { ComponentMutableSegment } from "./ComponentMutableSegment";
 import { getReferencedVariablesInside } from "~/utils/get-referenced-variables-inside";
+import { convertStatementToSegmentCallable } from "~/ast-factories/convert-statement-to-segment-callable";
 
 export class ComponentSideEffect extends ComponentMutableSegment {
-  private appliedModification = false;
-
   constructor(
     component: Component,
     private _path: babel.NodePath<babel.types.Statement>
@@ -30,23 +29,24 @@ export class ComponentSideEffect extends ComponentMutableSegment {
     });
   }
 
-  applyModification() {
-    if (this.appliedModification) {
-      return;
+  applyTransformation(performReplacement = true) {
+    if (this.appliedTransformation) {
+      return null;
     }
 
-    const dependenciesCondition = this.makeDependencyCondition();
+    const hasHookCall = this.hasHookCall();
 
-    if (dependenciesCondition) {
-      this.path.replaceWithMultiple([
-        t.ifStatement(
-          dependenciesCondition,
-          t.blockStatement([this.path.node])
-        ),
-      ]);
-    }
+    const { segmentCallableId, newPaths, replacements } =
+      convertStatementToSegmentCallable(this.path, { performReplacement });
+    this.appliedTransformation = true;
 
-    this.appliedModification = true;
+    return {
+      replacements,
+      segmentCallableId,
+      newPaths,
+      hasHookCall,
+      dependencyConditions: this.makeDependencyCondition(),
+    };
   }
 
   get path() {

packages/compiler/src/classes/ComponentVariable.ts

@@ -1,10 +1,12 @@
 import type * as babel from "@babel/core";
 import { Binding } from "@babel/traverse";
 import * as t from "@babel/types";
+import { convertStatementToSegmentCallable } from "~/ast-factories/convert-statement-to-segment-callable";
 import { makeCacheEnqueueCallStatement } from "~/ast-factories/make-cache-enqueue-call-statement";
 import { makeUnwrappedDeclarations } from "~/ast-factories/make-unwrapped-declarations";
-import { variableDeclarationToAssignment } from "~/ast-factories/variable-declaration-to-assignment";
 import {
+  DEFAULT_SEGMENT_CALLABLE_VARIABLE_NAME,
+  DEFAULT_UNWRAPPED_PROPS_VARIABLE_NAME,
   DEFAULT_UNWRAPPED_VARIABLE_NAME,
   RUNTIME_MODULE_CACHE_IS_NOT_SET_PROP_NAME,
   RUNTIME_MODULE_CACHE_VALUE_PROP_NAME,
@@ -16,7 +18,9 @@ import { Component } from "./Component";
 import { ComponentMutableSegment } from "./ComponentMutableSegment";
 
 export class ComponentVariable extends ComponentMutableSegment {
-  private appliedModification = false;
+  private variableSideEffects = new Set<ComponentVariable>();
+
+  private segmentCallableId: t.Identifier | null = null;
 
   constructor(
     public binding: Binding,
@@ -37,7 +41,7 @@ export class ComponentVariable extends ComponentMutableSegment {
   unwrapAssignmentPatterns() {
     const { path, scope } = this.binding;
 
-    let newPaths: babel.NodePath<babel.types.Node>[] = [];
+    const newPaths: babel.NodePath<babel.types.Node>[] = [];
     let unwrappedEntries: UnwrappedAssignmentEntry[] = [];
     let unwrapVariableId: t.Identifier | null = null;
 
@@ -73,13 +77,20 @@ export class ComponentVariable extends ComponentMutableSegment {
         tempVariableDeclarator,
       ]);
 
-      const allNewPaths = parentPath.replaceWithMultiple([
-        initDeclaration,
-        ...unwrapResult.unwrappedDeclarations,
-      ]);
+      [initPath] = parentPath.replaceWith(initDeclaration);
+
+      unwrapResult.unwrappedDeclarations.forEach(([newNode, binding]) => {
+        if (!initPath) {
+          return;
+        }
 
-      initPath = allNewPaths[0];
-      newPaths = allNewPaths.slice(1);
+        const [newPath] = initPath.insertAfter(newNode);
+        newPaths.push(newPath);
+        if (binding) {
+          binding.kind = newNode.kind as "const" | "let" | "var";
+          binding.path = newPath;
+        }
+      });
 
       scope.registerDeclaration(initPath);
 
@@ -93,7 +104,9 @@ export class ComponentVariable extends ComponentMutableSegment {
         return;
       }
 
-      unwrapVariableId = scope.generateUidIdentifier("props");
+      unwrapVariableId = scope.generateUidIdentifier(
+        DEFAULT_UNWRAPPED_PROPS_VARIABLE_NAME
+      );
 
       const unwrapResult = makeUnwrappedDeclarations(
         path as babel.NodePath<babel.types.LVal>,
@@ -105,10 +118,14 @@ export class ComponentVariable extends ComponentMutableSegment {
 
       scope.registerBinding("param", initPath);
 
-      newPaths = componentBlock.unshiftContainer(
-        "body",
-        unwrapResult.unwrappedDeclarations
-      );
+      unwrapResult.unwrappedDeclarations.forEach(([newNode, binding]) => {
+        const [newPath] = componentBlock.unshiftContainer("body", newNode);
+        newPaths.push(newPath);
+        if (binding) {
+          binding.kind = newNode.kind as "const" | "let" | "var";
+          binding.path = newPath;
+        }
+      });
 
       this.binding.kind = "let";
 
@@ -188,28 +205,6 @@ export class ComponentVariable extends ComponentMutableSegment {
 
         visitDependencies(ids);
       } else {
-        // TODO: Side effect calculation
-        const parentStatement = referencePath.find(
-          (p) =>
-            p.isStatement() && p.parentPath === this.component.path.get("body")
-        ) as babel.NodePath<babel.types.Statement> | null;
-
-        if (
-          !parentStatement ||
-          // Skip variable declarations for side effects
-          parentStatement.isVariableDeclaration() ||
-          // Skip return statements for side effects
-          parentStatement.isReturnStatement()
-        ) {
-          return;
-        }
-
-        const sideEffect = this.component.addSideEffect(parentStatement);
-
-        if (sideEffect) {
-          this.addSideEffect(sideEffect);
-          sideEffect.addDependency(this);
-        }
       }
     });
 
@@ -222,60 +217,82 @@ export class ComponentVariable extends ComponentMutableSegment {
     });
   }
 
-  applyModification() {
-    if (this.appliedModification) {
-      return;
-    }
-
-    const cacheUpdateEnqueueStatement = makeCacheEnqueueCallStatement(
+  getCacheUpdateEnqueueStatement() {
+    return makeCacheEnqueueCallStatement(
       this.getCacheAccessorExpression(),
       this.name
     );
+  }
 
-    const valueDeclarationWithCache = t.variableDeclaration("let", [
-      t.variableDeclarator(
-        t.identifier(this.name),
-        this.getCacheValueAccessExpression()
-      ),
-    ]);
+  getSegmentCallableId() {
+    if (!this.segmentCallableId) {
+      this.segmentCallableId = this.component.path.scope.generateUidIdentifier(
+        DEFAULT_SEGMENT_CALLABLE_VARIABLE_NAME
+      );
+    }
+
+    return this.segmentCallableId;
+  }
+
+  applyTransformation(performReplacement = true) {
+    if (this.appliedTransformation) {
+      return null;
+    }
+
+    const cacheUpdateEnqueueStatement = this.getCacheUpdateEnqueueStatement();
 
     switch (this.binding.kind) {
       case "const":
       case "let":
       case "var": {
-        const variableDeclaration =
-          this.getParentStatement() as babel.NodePath<babel.types.VariableDeclaration>;
+        const hasHookCall = this.hasHookCall();
+        const variableDeclaration = this.path.find((p) =>
+          p.isVariableDeclaration()
+        ) as babel.NodePath<babel.types.VariableDeclaration>;
 
-        if (this.hasHookCall()) {
-          variableDeclaration.insertAfter(cacheUpdateEnqueueStatement);
-          return;
-        }
+        const cacheValueAccessExpression = this.getCacheValueAccessExpression();
 
         const dependencyConditions = this.makeDependencyCondition();
-        if (dependencyConditions) {
-          variableDeclaration.insertBefore(valueDeclarationWithCache);
-          variableDeclaration.replaceWith(
-            t.ifStatement(
-              dependencyConditions,
-              t.blockStatement([
-                ...variableDeclarationToAssignment(variableDeclaration),
-                cacheUpdateEnqueueStatement,
-              ])
-            )
-          );
+
+        const { newPaths, segmentCallableId, replacements } =
+          convertStatementToSegmentCallable(variableDeclaration, {
+            initialValue: cacheValueAccessExpression,
+            performReplacement,
+            segmentCallableId: this.getSegmentCallableId(),
+          });
+
+        this.appliedTransformation = true;
+
+        const newId = newPaths?.[0]?.get(
+          "declarations.0.id"
+        ) as babel.NodePath<babel.types.LVal>;
+
+        if (newId) {
+          this.binding.path = newId;
         }
-        break;
+
+        return {
+          replacements,
+          segmentCallableId,
+          dependencyConditions,
+          newPaths,
+          hasHookCall,
+          updateCache: cacheUpdateEnqueueStatement,
+        };
       }
 
       case "param": {
         this.component
           .getFunctionBlockStatement()
           ?.unshiftContainer("body", cacheUpdateEnqueueStatement);
+        this.appliedTransformation = true;
         break;
       }
     }
 
-    this.appliedModification = true;
+    this.appliedTransformation = true;
+
+    return null;
   }
 
   private getCacheAccessorExpression() {
@@ -300,6 +317,14 @@ export class ComponentVariable extends ComponentMutableSegment {
     );
   }
 
+  addVariableSideEffect(componentVariable: ComponentVariable) {
+    this.variableSideEffects.add(componentVariable);
+  }
+
+  getVariableSideEffects() {
+    return this.variableSideEffects;
+  }
+
   get path() {
     return this.binding.path;
   }

packages/compiler/src/classes/tests/Component-fixtures.test.ts

@@ -3,23 +3,32 @@ import { parseFixture } from "~/utils/testing";
 
 const parseCodeAndRun = (fixtureName: string) => {
   const programPath = parseFixture(fixtureName);
-  const [component] = findComponents(programPath);
+  const components = findComponents(programPath);
 
-  return [component!, programPath] as const;
+  components.forEach((component) => {
+    component.computeComponentVariables();
+    component.applyTransformation();
+  });
+
+  return [components, programPath] as const;
 };
 
 describe("Component fixtures", () => {
   describe.only("applyModification", () => {
     it.each([
       // "fixture_1", "fixture_2", "fixture_4", "fixture_3",
+      // "fixture_10",
       "fixture_5",
+      // "fixture_7",
+      // "fixture_8",
+      // "fixture_9",
     ])("%s", (fixtureName) => {
-      const [component, program] = parseCodeAndRun(fixtureName);
-      component.computeComponentVariables();
-      component.applyModification();
+      const [, program] = parseCodeAndRun(fixtureName);
 
       const codeAfter = program.toString();
 
+      console.log(codeAfter);
+
       expect(codeAfter).toMatchSnapshot();
     });
   });

packages/compiler/src/classes/tests/__snapshots__/Component-fixtures.test.ts.snap

@@ -1,83 +1,164 @@
 // Jest Snapshot v1, https://goo.gl/fbAQLP
 
 exports[`Component fixtures applyModification fixture_5 1`] = `
-"const useMovies = () => {
-  return [{
-    title: "The Shawshank Redemption",
-    year: 1994
-  }, {
-    title: "The Godfather",
-    year: 1972
-  }, {
-    title: "The Dark Knight",
-    year: 2008
-  }];
+"const words = ["Shawshank", "Redemption", "Godfather", "Dark", "Knight", "Forest", "Gump", "Private", "Ryan"];
+const initalMovies = [{
+  title: "The Shawshank Redemption",
+  year: 1994
+}, {
+  title: "The Godfather",
+  year: 1972
+}, {
+  title: "The Dark Knight",
+  year: 2008
+}];
+const useMovies = () => {
+  const [movies, setMovies] = useState(initalMovies);
+  const addRandomMovie = () => {
+    setMovies(p => [...p, {
+      title: "The " + words[Math.floor(Math.random() * words.length)] + " " + words[Math.floor(Math.random() * words.length)],
+      year: 1990 + Math.floor(Math.random() * 20)
+    }]);
+  };
+  return [movies, addRandomMovie];
 };
+const useKey = () => "key1";
 const object = {
   key1: "value1",
   key2: "value2"
 };
 export function MyComponent() {
-  const [_$unforgetCache, _$unforgetCommit] = useCreateCache$unforget(
+  const [_$unforgetCache, _$unforgetCommit, _$unforgetNull] = useCreateCache$unforget(
   /*
-  0 => _unwrappedJsxEl2
-  1 => _unwrappedJsxExp
-  2 => filteredMovies
-  3 => _unwrappedJsxEl
-  4 => _unwrappedJsxExp2
-  5 => _unwrappedJsxExp3
-  6 => movies
+  0 => _unwrappedJsxEl3
+  1 => _unwrappedJsxEl
+  2 => addRandomMovie
+  3 => _unwrapped
+  4 => movies
+  5 => _unwrappedJsxEl2
+  6 => _unwrappedJsxExp
+  7 => _unwrappedJsxExp2
+  8 => filteredMovies
+  9 => value
+  10 => key
   */
-  7);
-  const movies = useMovies();
-
-  // const key = useKey();
-
-  // const value = object[key];
-  _$unforgetCache[6].e(movies);
-  let filteredMovies = _$unforgetCache[2].v;
-  if (_$unforgetCache[2].n) {
+  11);
+  let _unwrapped = _$unforgetCache[3].v;
+  const _segment = () => {
+    _unwrapped = useMovies();
+  };
+  let addRandomMovie = _$unforgetCache[2].v;
+  const _segment2 = () => {
+    addRandomMovie = _unwrapped[1];
+  };
+  let movies = _$unforgetCache[4].v;
+  const _segment3 = () => {
+    movies = _unwrapped[0];
+  };
+  let key = _$unforgetCache[10].v;
+  const _segment4 = () => {
+    key = useKey();
+  };
+  let value = _$unforgetCache[9].v;
+  const _segment5 = () => {
+    value = object[key];
+  };
+  let filteredMovies = _$unforgetCache[8].v;
+  const _segment6 = () => {
     filteredMovies = [];
-    _$unforgetCache[2].e(filteredMovies);
-  } // const i = 5;
-  if (movies.length !== _$unforgetCache[6].v.length || movies !== _$unforgetCache[6].v || filteredMovies !== _$unforgetCache[2].v) {
+  };
+  let i;
+  const _segment7 = () => {
+    i = 5;
+  };
+  const _segment8 = () => {
     for (let i = 0; i < movies.length; i++) {
       if (movies[i].year > 2000) {
         filteredMovies.push(movies[i]);
       }
     }
+  };
+  const _segment9 = () => {
+    if (filteredMovies.length > 0) {
+      // Just some side effect
+      console.log("Movies after 2000: ", filteredMovies);
+    }
+  };
+  let _unwrappedJsxExp = _$unforgetCache[6].v;
+  const _segment10 = () => {
+    _unwrappedJsxExp = movies.length;
+  };
+  let _unwrappedJsxExp2 = _$unforgetCache[7].v;
+  const _segment11 = () => {
+    _unwrappedJsxExp2 = filteredMovies.map(movie => <div key={movie.title}>
+          {movie.title}
+          {value}
+        </div>);
+  };
+  let _unwrappedJsxEl = _$unforgetCache[1].v;
+  const _segment12 = () => {
+    _unwrappedJsxEl = <button onClick={addRandomMovie}>Add random movie</button>;
+  };
+  let _unwrappedJsxEl2 = _$unforgetCache[5].v;
+  const _segment13 = () => {
+    _unwrappedJsxEl2 = <div>total number of movies: {_unwrappedJsxExp}</div>;
+  };
+  let _unwrappedJsxEl3 = _$unforgetCache[0].v;
+  const _segment14 = () => {
+    _unwrappedJsxEl3 = <div>
+      {_unwrappedJsxEl}
+      {_unwrappedJsxEl2}
+      {_unwrappedJsxExp2}
+    </div>;
+  };
+  _segment();
+  _$unforgetCache[3].e(_unwrapped);
+  if (_$unforgetCache[2].n || _unwrapped[1] !== _$unforgetCache[3].v[1]) {
+    _segment2();
+    _$unforgetCache[2].e(addRandomMovie);
   }
-  let _unwrappedJsxExp2 = _$unforgetCache[4].v;
-  if (_$unforgetCache[4].n) {
-    _unwrappedJsxExp2 = movie.title;
-    _$unforgetCache[4].e(_unwrappedJsxExp2);
+  if (_$unforgetCache[4].n || _unwrapped[0] !== _$unforgetCache[3].v[0]) {
+    _segment3();
+    _$unforgetCache[4].e(movies);
   }
-  let _unwrappedJsxExp3 = _$unforgetCache[5].v;
-  if (_$unforgetCache[5].n) {
-    _unwrappedJsxExp3 = movie.title;
-    _$unforgetCache[5].e(_unwrappedJsxExp3);
+  _segment4();
+  _$unforgetCache[10].e(key);
+  if (_$unforgetCache[9].n) {
+    _segment5();
+    _$unforgetCache[9].e(value);
   }
-  let _unwrappedJsxEl = _$unforgetCache[3].v;
-  if (_$unforgetCache[3].n || _unwrappedJsxExp2 !== _$unforgetCache[4].v || _unwrappedJsxExp3 !== _$unforgetCache[5].v) {
-    _unwrappedJsxEl = <div key={_unwrappedJsxExp2}>
-          {_unwrappedJsxExp3} 
-          {/* {value} */}
-        </div>;
-    _$unforgetCache[3].e(_unwrappedJsxEl);
+  if (_$unforgetCache[8].n) {
+    _segment6();
+    _$unforgetCache[8].e(filteredMovies);
   }
-  let _unwrappedJsxExp = _$unforgetCache[1].v;
-  if (_$unforgetCache[1].n || filteredMovies !== _$unforgetCache[2].v || _unwrappedJsxEl !== _$unforgetCache[3].v) {
-    _unwrappedJsxExp = filteredMovies.map(movie => _unwrappedJsxEl);
-    _$unforgetCache[1].e(_unwrappedJsxExp);
+  _segment7();
+  if (_unwrapped !== _$unforgetCache[3].v || addRandomMovie !== _$unforgetCache[2].v || movies !== _$unforgetCache[4].v || key !== _$unforgetCache[10].v || value !== _$unforgetCache[9].v || filteredMovies !== _$unforgetCache[8].v || movies.length !== _$unforgetCache[4].v.length || _unwrappedJsxExp !== _$unforgetCache[6].v || _unwrappedJsxEl !== _$unforgetCache[1].v || _unwrappedJsxEl2 !== _$unforgetCache[5].v || _unwrappedJsxExp2 !== _$unforgetCache[7].v || _unwrappedJsxEl3 !== _$unforgetCache[0].v || filteredMovies.length !== _$unforgetCache[8].v.length) {
+    _segment8();
   }
-  let _unwrappedJsxEl2 = _$unforgetCache[0].v;
-  if (_$unforgetCache[0].n || _unwrappedJsxExp !== _$unforgetCache[1].v) {
-    _unwrappedJsxEl2 = <div>
-      {_unwrappedJsxExp}
-    </div>;
-    _$unforgetCache[0].e(_unwrappedJsxEl2);
+  if (_unwrapped !== _$unforgetCache[3].v || addRandomMovie !== _$unforgetCache[2].v || movies !== _$unforgetCache[4].v || key !== _$unforgetCache[10].v || value !== _$unforgetCache[9].v || filteredMovies !== _$unforgetCache[8].v || filteredMovies.length !== _$unforgetCache[8].v.length || _unwrappedJsxExp !== _$unforgetCache[6].v || _unwrappedJsxEl !== _$unforgetCache[1].v || _unwrappedJsxEl2 !== _$unforgetCache[5].v || _unwrappedJsxExp2 !== _$unforgetCache[7].v || _unwrappedJsxEl3 !== _$unforgetCache[0].v) {
+    _segment9();
+  }
+  if (_$unforgetCache[6].n) {
+    _segment10();
+    _$unforgetCache[6].e(_unwrappedJsxExp);
+  }
+  if (_$unforgetCache[7].n || filteredMovies !== _$unforgetCache[8].v || value !== _$unforgetCache[9].v) {
+    _segment11();
+    _$unforgetCache[7].e(_unwrappedJsxExp2);
+  }
+  if (_$unforgetCache[1].n || addRandomMovie !== _$unforgetCache[2].v) {
+    _segment12();
+    _$unforgetCache[1].e(_unwrappedJsxEl);
+  }
+  if (_$unforgetCache[5].n || _unwrappedJsxExp !== _$unforgetCache[6].v) {
+    _segment13();
+    _$unforgetCache[5].e(_unwrappedJsxEl2);
+  }
+  if (_$unforgetCache[0].n || _unwrappedJsxEl !== _$unforgetCache[1].v || _unwrappedJsxEl2 !== _$unforgetCache[5].v || _unwrappedJsxExp2 !== _$unforgetCache[7].v) {
+    _segment14();
+    _$unforgetCache[0].e(_unwrappedJsxEl3);
   }
-  _$unforgetCommit();
-  return _unwrappedJsxEl2;
-}"
+  return _unwrappedJsxEl3;
+}
+export default MyComponent;"
 `;

packages/compiler/src/fixtures/fixture_10.js

@@ -0,0 +1,111 @@
+import React, { useState, useEffect, useMemo, useCallback } from "react";
+
+const fetchUser = () => {
+  return fetch("https://api.github.com/users/mohebifar").then((response) => {
+    if (!response.ok) {
+      throw new Error("Network response was not ok");
+    }
+    return response.json();
+  });
+};
+
+const fetchUserFollowers = () => {
+  return fetch("https://api.github.com/users/mohebifar/followers").then(
+    (response) => {
+      if (!response.ok) {
+        throw new Error("Network response was not ok");
+      }
+      return response.json();
+    }
+  );
+};
+
+function UserList() {
+  const [user, setUser] = useState([]);
+  const [followers, setFollowers] = useState([]);
+  const [loading, setLoading] = useState(true);
+  const [error, setError] = useState(null);
+  const [evenOnes, setEvenOnes] = useState(false);
+
+  useEffect(() => {
+    Promise.all([
+      fetchUser(),
+      fetchUserFollowers(),
+      new Promise((resolve) => setTimeout(resolve, 2000)),
+    ])
+      .then(([userData, followersData]) => {
+        setFollowers(followersData);
+        setUser(userData);
+        setLoading(false);
+      })
+      .catch((error) => {
+        setError(error.toString());
+        setLoading(false);
+      });
+  }, []);
+
+  // useCallback to memoize a hypothetical handler function
+  const handleUserClick = useCallback((userId) => {
+    console.log("User clicked:", userId);
+    // Handler logic here...
+  }, []);
+
+  // without useCallback
+  const toggleEvenOnes = () => {
+    setEvenOnes((prev) => !prev);
+  };
+
+  const memoizedFollowers = useMemo(() => followers, [followers]);
+
+  const evenFollowers = memoizedFollowers.filter(
+    (_, index) => index % 2 === (evenOnes ? 0 : 1)
+  );
+
+  // Early return for loading state
+  if (loading) return <div>Loading...</div>;
+
+  // Early return for error state
+  if (error) return <div>Error: {error}</div>;
+
+  const userListElement = evenFollowers.map((follower) => (
+    <UserListItem key={follower.id} user={follower} />
+  ));
+
+  return (
+    <div>
+      <h1>Follwers of {user.name}</h1>
+      <h1>User List</h1>
+      <button onClick={toggleEvenOnes}>
+        {evenOnes ? "Show Odd" : "Show Even"}
+      </button>
+      <ul onClick={() => handleUserClick(user.id)}>{userListElement}</ul>
+    </div>
+  );
+}
+
+function UserListItem({ user: { login, avatar_url, html_url } }) {
+  return (
+    <li
+      style={{
+        display: "flex",
+        alignItems: "center",
+        justifyContent: "space-between",
+        gap: 10,
+        marginBottom: 10,
+      }}
+    >
+      <img
+        style={{
+          width: 100,
+          height: 100,
+          borderRadius: 20,
+        }}
+        src={avatar_url}
+        alt={login}
+      />
+      <a href={html_url}>{login}</a>
+    </li>
+  );
+}
+
+export default UserList;

packages/compiler/src/fixtures/fixture_5.js

@@ -1,38 +1,88 @@
+const words = [
+  "Shawshank",
+  "Redemption",
+  "Godfather",
+  "Dark",
+  "Knight",
+  "Forest",
+  "Gump",
+  "Private",
+  "Ryan",
+];
+
+const initalMovies = [
+  {
+    title: "The Shawshank Redemption",
+    year: 1994,
+  },
+  {
+    title: "The Godfather",
+    year: 1972,
+  },
+  {
+    title: "The Dark Knight",
+    year: 2008,
+  },
+];
+
 const useMovies = () => {
-  return [
-    { title: "The Shawshank Redemption", year: 1994 },
-    { title: "The Godfather", year: 1972 },
-    { title: "The Dark Knight", year: 2008 },
-  ];
+  const [movies, setMovies] = useState(initalMovies);
+
+  const addRandomMovie = () => {
+    setMovies((p) => [
+      ...p,
+      {
+        title:
+          "The " +
+          words[Math.floor(Math.random() * words.length)] +
+          " " +
+          words[Math.floor(Math.random() * words.length)],
+        year: 1990 + Math.floor(Math.random() * 20),
+      },
+    ]);
+  };
+
+  return [movies, addRandomMovie];
 };
 
+const useKey = () => "key1";
+
 const object = { key1: "value1", key2: "value2" };
 
 export function MyComponent() {
-  const movies = useMovies();
+  const [movies, addRandomMovie] = useMovies();
 
-  // const key = useKey();
+  const key = useKey();
 
-  // const value = object[key];
+  const value = object[key];
 
   const filteredMovies = [];
 
-  // const i = 5;
+  const i = 5;
 
   for (let i = 0; i < movies.length; i++) {
     if (movies[i].year > 2000) {
       filteredMovies.push(movies[i]);
     }
   }
 
+  if (filteredMovies.length > 0) {
+    // Just some side effect
+    console.log("Movies after 2000: ", filteredMovies);
+  }
+
   return (
     <div>
+      <button onClick={addRandomMovie}>Add random movie</button>
+      <div>total number of movies: {movies.length}</div>
       {filteredMovies.map((movie) => (
         <div key={movie.title}>
-          {movie.title} 
-          {/* {value} */}
+          {movie.title}
+          {value}
         </div>
       ))}
     </div>
   );
 }
+
+export default MyComponent;

packages/compiler/src/fixtures/fixture_6.js

@@ -0,0 +1,3 @@
+export function SimpleJSX() {
+  return <div></div>;
+}

packages/compiler/src/fixtures/fixture_7.js

@@ -0,0 +1,10 @@
+export function SimpleJSX() {
+  const value = useValue();
+
+  return (
+    <div>
+      <button>Hello</button>
+      <span>{value}</span>
+    </div>
+  );
+}

packages/compiler/src/fixtures/fixture_8.js

@@ -0,0 +1,9 @@
+export function SimpleJSX() {
+  const value = useValue();
+
+  if (value === "loading") {
+    return <div>Loading...</div>;
+  }
+
+  return <span />;
+}

packages/compiler/src/fixtures/fixture_9.js

@@ -0,0 +1,31 @@
+function useValue() {
+  const [value, setValue] = useState(Math.random());
+  useEffect(() => {
+    const interval = setInterval(() => {
+      setValue(Math.random());
+    }, 1000);
+
+    return () => clearInterval(interval);
+  }, []);
+  return value;
+}
+
+export function SimpleJSX() {
+  const value = useValue();
+
+  const valueWith2Decimal = value.toFixed(2);
+
+  if (value > 0.8) {
+    return <div>Loading because value is {valueWith2Decimal}...</div>;
+  }
+
+  const derivedValue = "state updated to: " + valueWith2Decimal;
+
+  return (
+    <div>
+      <button>Hello</button>
+      <div>{value}</div>
+      <div>{derivedValue}</div>
+    </div>
+  );
+}

packages/compiler/src/fixtures/fixture_unsupported_1.js

@@ -1,9 +1,12 @@
 const useMovies = () => {
-  return [
-    { title: "The Shawshank Redemption", year: 1994 },
-    { title: "The Godfather", year: 1972 },
-    { title: "The Dark Knight", year: 2008 },
-  ];
+  return {
+    data: [
+      { title: "The Shawshank Redemption", year: 1994 },
+      { title: "The Godfather", year: 1972 },
+      { title: "The Dark Knight", year: 2008 },
+    ],
+    loading: false,
+  };
 };
 
 const object = { key1: "value1", key2: "value2" };
@@ -15,7 +18,11 @@ export function MyComponent() {
 
   i = [{ title: "The Shawshank Redemption", year: 1994 }];
 
-  const movies = useMovies();
+  const { movies, loading } = useMovies();
+
+  if (loading) {
+    return <div>Loading...</div>;
+  }
 
   for (let i = 0; i < movies.length; i++) {
     if (movies[i].year > 2000) {
@@ -53,17 +60,9 @@ export function MyComponent() {
     i = [{ title: "The Shawshank Redemption", year: 1994 }];
   };
 
-  let movies;
+  let movies, loading;
   const updater_4 = () => {
-    movies = useMovies();
-  };
-
-  const updater_5 = () => {
-    for (let i = 0; i < movies.length; i++) {
-      if (movies[i].year > 2000) {
-        filteredMovies.push(movies[i]);
-      }
-    }
+    ({movies, loading}) = useMovies();
   };
 
   updater_1();
@@ -74,6 +73,19 @@ export function MyComponent() {
 
   updater_3();
   updater_4();
+
+  if (loading) {
+    return <div>Loading...</div>;
+  }
+
+  const updater_5 = () => {
+    for (let i = 0; i < movies.length; i++) {
+      if (movies[i].year > 2000) {
+        filteredMovies.push(movies[i]);
+      }
+    }
+  };
+
 
   if (movies_check || filteredMovies_check) {
     updater_5();

packages/compiler/src/utils/constants.ts

@@ -2,13 +2,16 @@ export const RUNTIME_MODULE_NAME = "@react-unforget/runtime";
 export const RUNTIME_MODULE_CREATE_CACHE_HOOK_NAME = "useCreateCache$unforget";
 export const DEFAULT_CACHE_VARIABLE_NAME = "$unforgetCache";
 export const DEFAULT_CACHE_COMMIT_VARIABLE_NAME = "$unforgetCommit";
+export const DEFAULT_CACHE_NULL_VARIABLE_NAME = "$unforgetNull";
 
 export const DEFAULT_UNUSED_VARIABLE_NAME = "_unused";
 export const DEFAULT_UNWRAPPED_VARIABLE_NAME = "_unwrapped";
 export const DEFAULT_UNWRAPPED_JSX_EXPRESSION_VARIABLE_NAME =
   "_unwrappedJsxExp";
+export const DEFAULT_UNWRAPPED_PROPS_VARIABLE_NAME = "_props";
 export const DEFAULT_UNWRAPPED_JSX_ELEMENT_VARIABLE_NAME = "_unwrappedJsxEl";
 export const DEFAULT_SIDE_EFFECT_VARIABLE_NAME = "_sideEffect";
+export const DEFAULT_SEGMENT_CALLABLE_VARIABLE_NAME = "_segment";
 
 export const RUNTIME_MODULE_CACHE_ENQUEUE_METHOD_NAME = "e";
 export const RUNTIME_MODULE_CACHE_VALUE_PROP_NAME = "v";

packages/compiler/src/utils/get-returns-of-function.ts

@@ -1,12 +1,12 @@
 import type * as babel from "@babel/core";
-import { getFunctionParent } from "./get-function-parent";
+import { isInTheSameFunctionScope } from "./is-in-the-same-function-scope";
 
 export function getReturnsOfFunction(fn: babel.NodePath<babel.types.Function>) {
   const returns: babel.NodePath<babel.types.ReturnStatement>[] = [];
 
   fn.traverse({
     ReturnStatement(path) {
-      if (getFunctionParent(path)?.node === fn.node) {
+      if (isInTheSameFunctionScope(path, fn)) {
         returns.push(path);
       }
     },

packages/compiler/src/utils/is-hook-call.ts

@@ -1,5 +1,6 @@
 import type * as babel from "@babel/core";
 import { getRightmostIdName } from "./get-rightmost-id-name";
+import { isInTheSameFunctionScope } from "./is-in-the-same-function-scope";
 
 export function isHookCall(path: babel.NodePath<babel.types.CallExpression>) {
   path.assertCallExpression();
@@ -18,3 +19,23 @@ export function isHookCall(path: babel.NodePath<babel.types.CallExpression>) {
 
   return /^use[A-Z]/.test(rightmostId);
 }
+
+export function hasHookCall(
+  path: babel.NodePath<babel.types.Node>,
+  componentPath: babel.NodePath<babel.types.Function>
+) {
+  let hasHookCall = false;
+  path.traverse({
+    CallExpression: (innerPath) => {
+      if (
+        isHookCall(innerPath) &&
+        isInTheSameFunctionScope(innerPath, componentPath)
+      ) {
+        hasHookCall = true;
+        return;
+      }
+    },
+  });
+
+  return hasHookCall;
+}

packages/compiler/src/utils/is-in-the-same-function-scope.ts

@@ -0,0 +1,19 @@
+import * as babel from "@babel/core";
+import * as t from "@babel/types";
+import { Scope } from "@babel/traverse";
+
+// We need this to properly detect if return statements belong to the same function
+export function isInTheSameFunctionScope(
+  path: babel.NodePath<babel.types.Node>,
+  fn: babel.NodePath<babel.types.Function>
+) {
+  let currentScope: Scope | null = path.scope;
+  do {
+    if (t.isFunction(currentScope.block)) {
+      return currentScope.block === fn.node;
+    }
+    currentScope = currentScope?.parent ?? null;
+} while (currentScope);
+
+  return false;
+}

packages/compiler/src/utils/tests/is-in-the-same-function-scope.test.ts

@@ -0,0 +1,76 @@
+import * as babel from "@babel/core";
+import { isInTheSameFunctionScope } from "../is-in-the-same-function-scope";
+import { parse } from "../testing";
+
+describe("isInTheSameFunctionScope", () => {
+  it("basic example", () => {
+    const path = parse(`
+  function testFunction() {
+    return true;
+  }
+  `);
+
+    const testFn = path.get("body.0") as babel.NodePath<babel.types.Function>;
+    const returnStatement = testFn.get(
+      "body.body.0"
+    ) as babel.NodePath<babel.types.ReturnStatement>;
+
+    expect(isInTheSameFunctionScope(returnStatement, testFn)).toStrictEqual(
+      true
+    );
+  });
+
+  it("with node in an if statement block", () => {
+    const path = parse(`
+function testFunction() {
+  if (true) {
+    return true;
+  }
+}
+  `);
+
+    const testFn = path.get("body.0") as babel.NodePath<babel.types.Function>;
+    const returnStatement = testFn.get(
+      "body.body.0.consequent.body.0"
+    ) as babel.NodePath<babel.types.ReturnStatement>;
+
+    expect(isInTheSameFunctionScope(returnStatement, testFn)).toStrictEqual(
+      true
+    );
+  });
+
+  it("detects when a path is not in the same function scope immediately", () => {
+    const path = parse(`
+function testFunction() {
+  const callback = () => {
+    return 'foo';
+  }
+  if (true) {
+    return true;
+  }
+}
+  `);
+
+    const testFn = path.get("body.0") as babel.NodePath<babel.types.Function>;
+    const callbackFn = testFn.get(
+      "body.body.0.declarations.0.init"
+    ) as babel.NodePath<babel.types.ArrowFunctionExpression>;
+    const returnStatementInCallback = callbackFn.get(
+      "body.body.0"
+    ) as babel.NodePath<babel.types.ReturnStatement>;
+
+    const returnStatement = testFn.get(
+      "body.body.1.consequent.body.0"
+    ) as babel.NodePath<babel.types.ReturnStatement>;
+
+    expect(isInTheSameFunctionScope(returnStatement, testFn)).toStrictEqual(
+      true
+    );
+    expect(
+      isInTheSameFunctionScope(returnStatementInCallback, testFn)
+    ).toStrictEqual(false);
+    expect(
+      isInTheSameFunctionScope(returnStatementInCallback, callbackFn)
+    ).toStrictEqual(true);
+  });
+});

packages/compiler/src/utils/unwrap-jsx-elements.ts

@@ -1,6 +1,7 @@
 import * as babel from "@babel/core";
 import * as t from "@babel/types";
 import { DEFAULT_UNWRAPPED_JSX_ELEMENT_VARIABLE_NAME } from "./constants";
+import { isInTheSameFunctionScope } from "./is-in-the-same-function-scope";
 import { unwrapGenericExpression } from "./unwrap-generic-expression";
 
 type JSXChild = t.JSXElement["children"][number];
@@ -23,6 +24,9 @@ export function unwrapJsxElements(fn: babel.NodePath<t.Function>) {
     }
 
     if (path.isJSXElement() || path.isJSXFragment()) {
+      if (!isInTheSameFunctionScope(path, fn)) {
+        return;
+      }
       const childrenPath = path.get("children") as babel.NodePath<
         t.JSXElement["children"][number]
       >[];

packages/compiler/src/utils/unwrap-jsx-expressions.ts

@@ -1,6 +1,7 @@
 import * as babel from "@babel/core";
 import * as t from "@babel/types";
 import { DEFAULT_UNWRAPPED_JSX_EXPRESSION_VARIABLE_NAME } from "./constants";
+import { isInTheSameFunctionScope } from "./is-in-the-same-function-scope";
 import { unwrapGenericExpression } from "./unwrap-generic-expression";
 
 export function unwrapJsxExpressions(fn: babel.NodePath<t.Function>) {
@@ -16,6 +17,9 @@ export function unwrapJsxExpressions(fn: babel.NodePath<t.Function>) {
       ) {
         return;
       }
+      if (!isInTheSameFunctionScope(path, fn)) {
+        return;
+      }
 
       unwrapGenericExpression(
         fn,

packages/runtime/src/index.ts

@@ -22,7 +22,7 @@ interface Commit {
 
 export const useCreateCache$unforget = <S extends number>(
   size: S
-): [FixedArray<CacheEntry, S>, Commit] => {
+): [FixedArray<CacheEntry, S>, Commit, UnassignedType] => {
   const valuesRef: MutableRefObject<FixedArray<CacheEntry, S> | null> =
     useRef(null);
 
@@ -68,5 +68,5 @@ export const useCreateCache$unforget = <S extends number>(
     throw new Error("Unreachable");
   }
 
-  return [valuesRef.current, commitRef.current];
+  return [valuesRef.current, commitRef.current, UNASSIGNED];
 };