Separate $$ref from $$cache

packages/forgetti/runtime/index.ts

@@ -12,14 +12,20 @@ export interface Ref<T> {
 }
 export type RefHook = <T>(callback: T) => Ref<T>;
 
-export function $$cache(hook: RefHook, size: number): unknown[] {
+export function $$ref(hook: RefHook, size: number): unknown[] {
   const ref = hook<unknown[] | undefined>(undefined);
   if (!ref.current) {
     ref.current = new Array(size);
   }
   return ref.current;
 }
 
+export type MemoHook = <T>(callback: () => T, dependencies: unknown[]) => T;
+
+export function $$cache(hook: MemoHook, size: number): unknown[] {
+  return hook(() => new Array<unknown>(size), []);
+}
+
 export function $$branch(parent: unknown[], index: number, size: number): unknown[] {
   parent[index] ||= new Array(size);
   return parent[index] as unknown[];

packages/forgetti/src/core/imports.ts

@@ -6,6 +6,12 @@ export const RUNTIME_EQUALS: ImportDefinition = {
   kind: 'named',
 };
 
+export const RUNTIME_REF: ImportDefinition = {
+  name: '$$ref',
+  source: 'forgetti/runtime',
+  kind: 'named',
+};
+
 export const RUNTIME_CACHE: ImportDefinition = {
   name: '$$cache',
   source: 'forgetti/runtime',

packages/forgetti/src/core/optimizer-scope.ts

@@ -2,7 +2,7 @@
 import * as t from '@babel/types';
 import type { OptimizedExpression, StateContext } from './types';
 import getImportIdentifier from './get-import-identifier';
-import { RUNTIME_BRANCH, RUNTIME_CACHE } from './imports';
+import { RUNTIME_BRANCH, RUNTIME_CACHE, RUNTIME_REF } from './imports';
 
 function mergeVariableDeclaration(statements: t.Statement[]): t.Statement[] {
   let stack: t.VariableDeclarator[] = [];
@@ -26,7 +26,11 @@ function mergeVariableDeclaration(statements: t.Statement[]): t.Statement[] {
 export default class OptimizerScope {
   memo: t.Identifier | undefined;
 
-  indeces = 0;
+  ref: t.Identifier | undefined;
+
+  indecesMemo = 0;
+
+  indecesRef = 0;
 
   ctx: StateContext;
 
@@ -50,64 +54,112 @@ export default class OptimizerScope {
     this.isInLoop = isInLoop;
   }
 
-  createHeader(): t.Identifier {
+  createHeader(type: 'memo' | 'ref' = 'memo'): t.Identifier {
+    if (type === 'ref') {
+      if (!this.ref) {
+        this.ref = this.path.scope.generateUidIdentifier('ref');
+      }
+      return this.ref;
+    }
+
     if (!this.memo) {
       this.memo = this.path.scope.generateUidIdentifier('cache');
     }
     return this.memo;
   }
 
-  createIndex(): t.NumericLiteral {
-    const current = this.indeces;
-    this.indeces += 1;
+  createIndex(type: 'memo' | 'ref'): t.NumericLiteral {
+    const current = type === 'memo' ? this.indecesMemo : this.indecesRef;
+    if (type === 'memo') {
+      this.indecesMemo += 1;
+    } else {
+      this.indecesRef += 1;
+    }
+
     return t.numericLiteral(current);
   }
 
-  getMemoDeclaration(): t.VariableDeclaration | undefined {
-    if (!this.memo) {
+  getMemoDeclarations(): t.VariableDeclaration[] | undefined {
+    if (!this.memo && !this.ref) {
       return undefined;
     }
     // This is for generating branched caching.
     // Parent means that we want to create the cache
     // from the parent (or root)
     if (this.parent) {
       const header = this.parent.createHeader();
-      const index = this.parent.createIndex();
-
-      return t.variableDeclaration('let', [
-        t.variableDeclarator(
-          this.createHeader(),
-          t.callExpression(
-            getImportIdentifier(
-              this.ctx,
-              this.path,
-              RUNTIME_BRANCH,
+      const index = this.parent.createIndex('memo');
+
+      return [
+        t.variableDeclaration('let', [
+          t.variableDeclarator(
+            this.createHeader(),
+            t.callExpression(
+              getImportIdentifier(
+                this.ctx,
+                this.path,
+                RUNTIME_BRANCH,
+              ),
+              [header, index, t.numericLiteral(this.indecesMemo)],
             ),
-            [header, index, t.numericLiteral(this.indeces)],
           ),
-        ),
-      ]);
+        ]),
+      ];
     }
-    return t.variableDeclaration('let', [
-      t.variableDeclarator(
-        this.memo,
-        t.callExpression(
-          getImportIdentifier(
-            this.ctx,
-            this.path,
-            RUNTIME_CACHE,
+
+    const outputDeclarations = [];
+
+    if (this.memo) {
+      outputDeclarations.push(
+
+        t.variableDeclaration('let', [
+          t.variableDeclarator(
+            this.memo,
+            t.callExpression(
+              getImportIdentifier(
+                this.ctx,
+                this.path,
+                RUNTIME_CACHE,
+              ),
+              [
+                getImportIdentifier(
+                  this.ctx,
+                  this.path,
+                  this.ctx.preset.runtime.useMemo,
+                ),
+                t.numericLiteral(this.indecesMemo),
+              ],
+            ),
           ),
-          [
-            getImportIdentifier(
-              this.ctx,
-              this.path,
-              this.ctx.preset.runtime.useRef,
+        ]),
+      );
+    }
+    if (this.ref) {
+      outputDeclarations.push(
+        t.variableDeclaration('let', [
+          t.variableDeclarator(
+            this.ref,
+            t.callExpression(
+              getImportIdentifier(
+                this.ctx,
+                this.path,
+                RUNTIME_REF,
+              ),
+              [
+                getImportIdentifier(
+                  this.ctx,
+                  this.path,
+                  this.ctx.preset.runtime.useRef,
+                ),
+                t.numericLiteral(this.indecesRef),
+              ],
             ),
-            t.numericLiteral(this.indeces),
-          ],
-        ),
-      ),
-    ]);
+          ),
+        ]),
+      );
+    }
+
+    return outputDeclarations;
   }
 
   loop: t.Identifier | undefined;
@@ -133,7 +185,7 @@ export default class OptimizerScope {
       return undefined;
     }
     const header = this.parent.createHeader();
-    const index = this.parent.createIndex();
+    const index = this.parent.createIndex('memo');
     const id = this.createLoopIndex();
 
     return t.variableDeclaration('let', [
@@ -168,7 +220,7 @@ export default class OptimizerScope {
             this.path,
             RUNTIME_BRANCH,
           ),
-          [header, localIndex, t.numericLiteral(this.indeces)],
+          [header, localIndex, t.numericLiteral(this.indecesMemo)],
         ),
       ),
     ]);
@@ -177,11 +229,11 @@ export default class OptimizerScope {
   getStatements(): t.Statement[] {
     const result = [...this.statements];
     const header = this.isInLoop
-      ? this.getLoopDeclaration()
-      : this.getMemoDeclaration();
+      ? [this.getLoopDeclaration()]
+      : this.getMemoDeclarations();
     if (header) {
       return mergeVariableDeclaration([
-        header,
+        ...header,
         ...result,
       ]);
     }

packages/forgetti/src/core/optimizer.ts

@@ -73,6 +73,7 @@ export default class Optimizer {
   createMemo(
     current: t.Expression,
     dependencies?: t.Expression | (t.Expression | undefined)[] | boolean,
+    type: 'memo' | 'ref' = 'memo',
   ): OptimizedExpression {
     // Check if the identifier is an already optimized
     // identifier so that we can skip it.
@@ -86,10 +87,10 @@ export default class Optimizer {
     const header = (
       this.scope.isInLoop
         ? this.scope.createLoopHeader()
-        : this.scope.createHeader()
+        : this.scope.createHeader(type)
     );
     // Get the memo index
-    const index = this.scope.createIndex();
+    const index = this.scope.createIndex(type);
     // Generate the access expression
     const pos = t.memberExpression(header, index, true);
     // Generate the `v` identifier
@@ -484,7 +485,7 @@ export default class Optimizer {
         init || t.identifier('undefined'),
       ),
     ]);
-    return this.createMemo(expr, true);
+    return this.createMemo(expr, true, 'ref');
   }
 
   optimizeCallExpression(

packages/forgetti/src/core/presets.ts

@@ -33,6 +33,7 @@ export interface Preset {
   };
   runtime: {
     useRef: ImportDefinition;
+    useMemo: ImportDefinition;
     memo?: ImportDefinition;
   };
   imports: {
@@ -67,6 +68,11 @@ export const PRESETS = {
         source: 'react',
         kind: 'named',
       },
+      useMemo: {
+        name: 'useMemo',
+        source: 'react',
+        kind: 'named',
+      },
       memo: {
         name: 'memo',
         source: 'react',
@@ -143,6 +149,11 @@ export const PRESETS = {
         source: 'preact/hooks',
         kind: 'named',
       },
+      useMemo: {
+        name: 'useMemo',
+        source: 'preact/hooks',
+        kind: 'named',
+      },
       memo: {
         name: 'memo',
         source: 'preact/compat',