Fix cache header in loops

lxsmnsyc · Mar 12, 2024 · cd79ff1 · cd79ff1
1 parent 9f7ea33
commit cd79ff1
Show file tree

Hide file tree

Showing 3 changed files with 73 additions and 48 deletions.
diff --git a/packages/forgetti/src/core/optimizer-scope.ts b/packages/forgetti/src/core/optimizer-scope.ts
@@ -177,15 +177,22 @@ export default class OptimizerScope {
     return undefined;
   }
 
-  loop: t.Identifier | undefined;
+  loop: {
+    memo: t.Identifier | undefined;
+    ref: t.Identifier | undefined;
+  } = {
+    memo: undefined,
+    ref: undefined,
+  };
 
   loopIndex: t.Identifier | undefined;
 
-  createLoopHeader(): t.Identifier {
-    if (!this.loop) {
-      this.loop = this.path.scope.generateUidIdentifier('loop');
+  createLoopHeader(type: 'memo' | 'ref'): t.Identifier {
+    let current = this.loop[type];
+    if (!current) {
+      current = this.loop[type] = this.path.scope.generateUidIdentifier('loop');
     }
-    return this.loop;
+    return current;
   }
 
   createLoopIndex(): t.Identifier {
@@ -239,45 +246,55 @@ export default class OptimizerScope {
     ]);
   }
 
-  getLoopMemoHeaderDeclaration(): t.VariableDeclaration {
-    const header = this.createHeader('memo');
-    const index = this.createLoopIndex();
-    const localIndex = this.path.scope.generateUidIdentifier('loopId');
-    return t.variableDeclaration('let', [
-      t.variableDeclarator(localIndex, t.updateExpression('++', index)),
-      t.variableDeclarator(
-        this.createLoopHeader(),
-        t.callExpression(
-          getImportIdentifier(this.ctx, this.path, RUNTIME_BRANCH),
-          [header, localIndex, t.numericLiteral(this.indecesMemo)],
-        ),
-      ),
-    ]);
+  getLoopMemoHeaderDeclaration(): t.VariableDeclaration[] | undefined {
+    if (this.loop.memo) {
+      const header = this.createHeader('memo');
+      const index = this.createLoopIndex();
+      const localIndex = this.path.scope.generateUidIdentifier('loopId');
+      return [
+        t.variableDeclaration('let', [
+          t.variableDeclarator(localIndex, t.updateExpression('++', index)),
+          t.variableDeclarator(
+            this.createLoopHeader('memo'),
+            t.callExpression(
+              getImportIdentifier(this.ctx, this.path, RUNTIME_BRANCH),
+              [header, localIndex, t.numericLiteral(this.indecesMemo)],
+            ),
+          ),
+        ]),
+      ];
+    }
+    return undefined;
   }
 
-  getLoopRefHeaderDeclaration(): t.VariableDeclaration {
-    const header = this.createHeader('ref');
-    const index = this.createLoopIndex();
-    const localIndex = this.path.scope.generateUidIdentifier('loopId');
-    return t.variableDeclaration('let', [
-      t.variableDeclarator(localIndex, t.updateExpression('++', index)),
-      t.variableDeclarator(
-        this.createLoopHeader(),
-        t.callExpression(
-          getImportIdentifier(this.ctx, this.path, RUNTIME_BRANCH),
-          [header, localIndex, t.numericLiteral(this.indecesRef)],
-        ),
-      ),
-    ]);
+  getLoopRefHeaderDeclaration(): t.VariableDeclaration[] | undefined {
+    if (this.loop.ref) {
+      const header = this.createHeader('ref');
+      const index = this.createLoopIndex();
+      const localIndex = this.path.scope.generateUidIdentifier('loopId');
+      return [
+        t.variableDeclaration('let', [
+          t.variableDeclarator(localIndex, t.updateExpression('++', index)),
+          t.variableDeclarator(
+            this.createLoopHeader('ref'),
+            t.callExpression(
+              getImportIdentifier(this.ctx, this.path, RUNTIME_BRANCH),
+              [header, localIndex, t.numericLiteral(this.indecesRef)],
+            ),
+          ),
+        ]),
+      ];
+    }
+    return undefined;
   }
 
   getStatements(): t.Statement[] {
     const result = [...this.statements];
     const memoHeader = this.isInLoop
-      ? [this.getLoopMemoHeaderDeclaration()]
+      ? this.getLoopMemoHeaderDeclaration()
       : this.getMemoDeclarations();
     const refHeader = this.isInLoop
-      ? [this.getLoopRefHeaderDeclaration()]
+      ? this.getLoopRefHeaderDeclaration()
       : this.getRefDeclarations();
     return mergeVariableDeclaration([
       ...(memoHeader || []),

diff --git a/packages/forgetti/src/core/optimizer.ts b/packages/forgetti/src/core/optimizer.ts
@@ -91,7 +91,7 @@ export default class Optimizer {
     }
     // Creates the cache header
     const header = this.scope.isInLoop
-      ? this.scope.createLoopHeader()
+      ? this.scope.createLoopHeader(type)
       : this.scope.createHeader(type);
     // Get the memo index
     const index = this.scope.createIndex(type);

diff --git a/packages/forgetti/test/__snapshots__/hooks.test.ts.snap b/packages/forgetti/test/__snapshots__/hooks.test.ts.snap
@@ -265,11 +265,18 @@ function Example(props) {
 exports[`hooks > should optimize useRef 1`] = `
 "import { useRef as _useRef } from \\"react\\";
 import { $$ref as _$$ref } from \\"forgetti/runtime\\";
+import { useMemo as _useMemo } from \\"react\\";
+import { $$cache as _$$cache } from \\"forgetti/runtime\\";
+import { $$equals as _$$equals } from \\"forgetti/runtime\\";
 import { useRef } from 'react';
 function Example(props) {
-  let _ref = _$$ref(_useRef, 1);
+  let _cache = _$$cache(_useMemo, 2),
+    _ref = _$$ref(_useRef, 1),
+    _equals = _$$equals(_cache, 0, props),
+    _value = _equals ? _cache[0] : _cache[0] = props,
+    _value2 = _equals ? _cache[1] : _cache[1] = _value.value;
   return 0 in _ref ? _ref[0] : _ref[0] = {
-    current: props.value
+    current: _value2
   };
 }"
 `;
@@ -282,17 +289,18 @@ import { $$cache as _$$cache } from \\"forgetti/runtime\\";
 import { $$equals as _$$equals } from \\"forgetti/runtime\\";
 import { useRef, useMemo } from 'react';
 function Example(props) {
-  let _cache = _$$cache(_useMemo, 5),
-    _ref = _$$ref(_useRef, 1);
+  let _cache = _$$cache(_useMemo, 6),
+    _ref = _$$ref(_useRef, 1),
+    _equals = _$$equals(_cache, 0, props),
+    _value = _equals ? _cache[0] : _cache[0] = props,
+    _value2 = _equals ? _cache[1] : _cache[1] = _value.foo;
   const aRef = 0 in _ref ? _ref[0] : _ref[0] = {
-    current: props.foo
+    current: _value2
   };
-  let _equals = _$$equals(_cache, 0, props),
-    _value2 = _equals ? _cache[0] : _cache[0] = props,
-    _value3 = _equals ? _cache[1] : _cache[1] = _value2.bar,
-    _equals2 = _$$equals(_cache, 2, _value3),
-    _value4 = _equals2 ? _cache[2] : _cache[2] = _value3,
-    _value5 = _equals2 ? _cache[3] : _cache[3] = [_value4];
-  return _equals2 ? _cache[4] : _cache[4] = (() => props.bar())();
+  let _value4 = _equals ? _cache[2] : _cache[2] = _value.bar,
+    _equals2 = _$$equals(_cache, 3, _value4),
+    _value5 = _equals2 ? _cache[3] : _cache[3] = _value4,
+    _value6 = _equals2 ? _cache[4] : _cache[4] = [_value5];
+  return _equals2 ? _cache[5] : _cache[5] = (() => props.bar())();
 }"
 `;