Codegen comments to explain output

compiler/packages/babel-plugin-react-forget/src/HIR/Environment.ts

@@ -257,6 +257,14 @@ const EnvironmentConfigSchema = z.object({
    * ```
    */
   enableChangeVariableCodegen: z.boolean().default(false),
+
+  /**
+   * Enable emitting comments that explain Forget's output, and which
+   * values are being checked and which values produced by each memo block.
+   *
+   * Intended for use in demo purposes (incl playground)
+   */
+  enableMemoizationComments: z.boolean().default(false),
 });
 
 export type EnvironmentConfig = z.infer<typeof EnvironmentConfigSchema>;

compiler/packages/babel-plugin-react-forget/src/ReactiveScopes/CodegenReactiveFunction.ts

@@ -240,9 +240,12 @@ function codegenMemoBlockForReactiveScope(
   const cacheStoreStatements: Array<t.Statement> = [];
   const cacheLoadStatements: Array<t.Statement> = [];
   const changeExpressions: Array<t.Expression> = [];
+  const changeExpressionComments: Array<string> = [];
+  const outputComments: Array<string> = [];
   for (const dep of scope.dependencies) {
     const index = cx.nextCacheIndex;
     const depValue = codegenDependency(cx, dep);
+    changeExpressionComments.push(printDependencyComment(dep));
     const comparison = t.binaryExpression(
       "!==",
       t.memberExpression(t.identifier("$"), t.numericLiteral(index), true),
@@ -285,6 +288,7 @@ function codegenMemoBlockForReactiveScope(
     });
 
     const name = convertIdentifier(identifier);
+    outputComments.push(name.name);
     if (!cx.hasDeclared(identifier)) {
       statements.push(
         t.variableDeclaration("let", [t.variableDeclarator(name)])
@@ -316,6 +320,7 @@ function codegenMemoBlockForReactiveScope(
       firstOutputIndex = index;
     }
     const name = convertIdentifier(reassignment);
+    outputComments.push(name.name);
 
     cacheStoreStatements.push(
       t.expressionStatement(
@@ -369,7 +374,60 @@ function codegenMemoBlockForReactiveScope(
   const computationBlock = codegenBlock(cx, block);
   computationBlock.body.push(...cacheStoreStatements);
   const memoBlock = t.blockStatement(cacheLoadStatements);
-  statements.push(t.ifStatement(testCondition, computationBlock, memoBlock));
+  const memoStatement = t.ifStatement(
+    testCondition,
+    computationBlock,
+    memoBlock
+  );
+  if (cx.env.config.enableMemoizationComments) {
+    if (changeExpressionComments.length) {
+      t.addComment(
+        memoStatement,
+        "leading",
+        ` check if ${printDelimitedCommentList(
+          changeExpressionComments,
+          "or"
+        )} changed`,
+        true
+      );
+      t.addComment(
+        memoStatement,
+        "leading",
+        ` "useMemo" for ${printDelimitedCommentList(outputComments, "and")}:`,
+        true
+      );
+    } else {
+      t.addComment(
+        memoStatement,
+        "leading",
+        " cache value with no dependencies",
+        true
+      );
+      t.addComment(
+        memoStatement,
+        "leading",
+        ` "useMemo" for ${printDelimitedCommentList(outputComments, "and")}:`,
+        true
+      );
+    }
+    if (computationBlock.body.length > 0) {
+      t.addComment(
+        computationBlock.body[0]!,
+        "leading",
+        ` Inputs changed, recompute`,
+        true
+      );
+    }
+    if (memoBlock.body.length > 0) {
+      t.addComment(
+        memoBlock.body[0]!,
+        "leading",
+        ` Inputs did not change, use cached value`,
+        true
+      );
+    }
+  }
+  statements.push(memoStatement);
 }
 
 function codegenSignalBlockForReactiveScope(
@@ -789,6 +847,41 @@ function codegenForInit(
   }
 }
 
+function printDependencyComment(dependency: ReactiveScopeDependency): string {
+  const identifier = convertIdentifier(dependency.identifier);
+  let name = identifier.name;
+  if (dependency.path !== null) {
+    for (const path of dependency.path) {
+      name += `.${path}`;
+    }
+  }
+  return name;
+}
+
+function printDelimitedCommentList(
+  items: Array<string>,
+  finalCompletion: string
+): string {
+  if (items.length === 2) {
+    return items.join(` ${finalCompletion} `);
+  } else if (items.length <= 1) {
+    return items.join("");
+  }
+
+  let output = [];
+  for (let i = 0; i < items.length; i++) {
+    const item = items[i]!;
+    if (i < items.length - 2) {
+      output.push(`${item}, `);
+    } else if (i === items.length - 2) {
+      output.push(`${item}, ${finalCompletion} `);
+    } else {
+      output.push(item);
+    }
+  }
+  return output.join("");
+}
+
 function codegenDependency(
   cx: Context,
   dependency: ReactiveScopeDependency

compiler/packages/babel-plugin-react-forget/src/__tests__/fixtures/compiler/memoization-comments.expect.md

@@ -0,0 +1,76 @@
+
+## Input
+
+```javascript
+// @enableMemoizationComments
+import { addOne, getNumber, identity } from "shared-runtime";
+
+function Component(props) {
+  const x = identity(props.a);
+  const y = addOne(x);
+  const z = identity(props.b);
+  return [x, y, z];
+}
+
+export const FIXTURE_ENTRYPOINT = {
+  fn: Component,
+  params: [{ a: 1, b: 10 }],
+};
+
+```
+
+## Code
+
+```javascript
+import { unstable_useMemoCache as useMemoCache } from "react"; // @enableMemoizationComments
+import { addOne, getNumber, identity } from "shared-runtime";
+
+function Component(props) {
+  const $ = useMemoCache(9);
+  let t0;
+  let x; // "useMemo" for t0 and x:
+  // check if props.a changed
+  if ($[0] !== props.a) {
+    // Inputs changed, recompute
+    x = identity(props.a);
+    t0 = addOne(x);
+    $[0] = props.a;
+    $[1] = t0;
+    $[2] = x;
+  } else {
+    // Inputs did not change, use cached value
+    t0 = $[1];
+    x = $[2];
+  }
+  const y = t0;
+  let t1; // "useMemo" for t1:
+  // check if props.b changed
+  if ($[3] !== props.b) {
+    // Inputs changed, recompute
+    t1 = identity(props.b);
+    $[3] = props.b;
+    $[4] = t1;
+  } else {
+    // Inputs did not change, use cached value
+    t1 = $[4];
+  }
+  const z = t1;
+  let t2; // "useMemo" for t2:
+  // check if x, y, or z changed
+  if ($[5] !== x || $[6] !== y || $[7] !== z) {
+    // Inputs changed, recompute
+    t2 = [x, y, z];
+    $[5] = x;
+    $[6] = y;
+    $[7] = z;
+    $[8] = t2;
+  } else {
+    // Inputs did not change, use cached value
+    t2 = $[8];
+  }
+  return t2;
+}
+export const FIXTURE_ENTRYPOINT = { fn: Component, params: [{ a: 1, b: 10 }] };
+
+```
+      

compiler/packages/babel-plugin-react-forget/src/__tests__/fixtures/compiler/memoization-comments.js

@@ -0,0 +1,14 @@
+// @enableMemoizationComments
+import { addOne, getNumber, identity } from "shared-runtime";
+
+function Component(props) {
+  const x = identity(props.a);
+  const y = addOne(x);
+  const z = identity(props.b);
+  return [x, y, z];
+}
+
+export const FIXTURE_ENTRYPOINT = {
+  fn: Component,
+  params: [{ a: 1, b: 10 }],
+};