compiler: Log metrics on pruned memo blocks/values

Adds additional information to the CompileSuccess LoggerEvent:
* `prunedMemoBlocks` is the number of reactive scopes that were pruned for some reason.
* `prunedMemoValues` is the number of unique _values_ produced by those scopes.

Both numbers exclude blocks that are just a hook call - ie although we create and prune a scope for eg `useState()`, that's just an artifact of the sequencing of our pipeline. So what this metric is counting is cases of _other_ values that go unmemoized. See the new fixture, which takes advantage of improvements in the snap runner to optionally emit the logger events in the .expect.md file if you include the "logger" pragma in a fixture.

ghstack-source-id: c2015bb
Pull Request resolved: #29810

Author: josephsavona

compiler/packages/babel-plugin-react-compiler/src/Entrypoint/Options.ts

@@ -169,6 +169,9 @@ export type LoggerEvent =
       fnName: string | null;
       memoSlots: number;
       memoBlocks: number;
+      memoValues: number;
+      prunedMemoBlocks: number;
+      prunedMemoValues: number;
     }
   | {
       kind: "PipelineError";

compiler/packages/babel-plugin-react-compiler/src/Entrypoint/Program.ts

@@ -329,6 +329,9 @@ export function compileProgram(
         fnName: compiledFn.id?.name ?? null,
         memoSlots: compiledFn.memoSlotsUsed,
         memoBlocks: compiledFn.memoBlocks,
+        memoValues: compiledFn.memoValues,
+        prunedMemoBlocks: compiledFn.prunedMemoBlocks,
+        prunedMemoValues: compiledFn.prunedMemoValues,
       });
     } catch (err) {
       hasCriticalError ||= isCriticalError(err);

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

@@ -22,6 +22,7 @@ import {
   ObjectPropertyKey,
   Pattern,
   Place,
+  PrunedReactiveScopeBlock,
   ReactiveBlock,
   ReactiveFunction,
   ReactiveInstruction,
@@ -67,6 +68,23 @@ export type CodegenFunction = {
    * how many inputs/outputs each block has
    */
   memoBlocks: number;
+
+  /**
+   * Number of memoized values across all reactive scopes
+   */
+  memoValues: number;
+
+  /**
+   * The number of reactive scopes that were created but had to be discarded
+   * because they contained hook calls.
+   */
+  prunedMemoBlocks: number;
+
+  /**
+   * The total number of values that should have been memoized but weren't
+   * because they were part of a pruned memo block.
+   */
+  prunedMemoValues: number;
 };
 
 export function codegenFunction(
@@ -272,7 +290,7 @@ function codegenReactiveFunction(
     return Err(cx.errors);
   }
 
-  const countMemoBlockVisitor = new CountMemoBlockVisitor();
+  const countMemoBlockVisitor = new CountMemoBlockVisitor(fn.env);
   visitReactiveFunction(fn, countMemoBlockVisitor, undefined);
 
   return Ok({
@@ -284,16 +302,59 @@ function codegenReactiveFunction(
     generator: fn.generator,
     async: fn.async,
     memoSlotsUsed: cx.nextCacheIndex,
-    memoBlocks: countMemoBlockVisitor.count,
+    memoBlocks: countMemoBlockVisitor.memoBlocks,
+    memoValues: countMemoBlockVisitor.memoValues,
+    prunedMemoBlocks: countMemoBlockVisitor.prunedMemoBlocks,
+    prunedMemoValues: countMemoBlockVisitor.prunedMemoValues,
   });
 }
 
 class CountMemoBlockVisitor extends ReactiveFunctionVisitor<void> {
-  count: number = 0;
+  env: Environment;
+  memoBlocks: number = 0;
+  memoValues: number = 0;
+  prunedMemoBlocks: number = 0;
+  prunedMemoValues: number = 0;
 
-  override visitScope(scope: ReactiveScopeBlock, state: void): void {
-    this.count += 1;
-    this.traverseScope(scope, state);
+  constructor(env: Environment) {
+    super();
+    this.env = env;
+  }
+
+  override visitScope(scopeBlock: ReactiveScopeBlock, state: void): void {
+    this.memoBlocks += 1;
+    this.memoValues += scopeBlock.scope.declarations.size;
+    this.traverseScope(scopeBlock, state);
+  }
+
+  override visitPrunedScope(
+    scopeBlock: PrunedReactiveScopeBlock,
+    state: void
+  ): void {
+    let isHookOnlyMemoBlock = false;
+    if (
+      scopeBlock.instructions.length === 1 &&
+      scopeBlock.instructions[0].kind === "instruction"
+    ) {
+      const instr = scopeBlock.instructions[0]!.instruction;
+      if (
+        instr.value.kind === "MethodCall" ||
+        instr.value.kind === "CallExpression"
+      ) {
+        const callee =
+          instr.value.kind === "MethodCall"
+            ? instr.value.property
+            : instr.value.callee;
+        if (getHookKind(this.env, callee.identifier) != null) {
+          isHookOnlyMemoBlock = true;
+        }
+      }
+    }
+    if (!isHookOnlyMemoBlock) {
+      this.prunedMemoBlocks += 1;
+      this.prunedMemoValues += scopeBlock.scope.declarations.size;
+    }
+    this.traversePrunedScope(scopeBlock, state);
   }
 }
 

compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/log-pruned-memoization.expect.md

@@ -0,0 +1,94 @@
+
+## Input
+
+```javascript
+// @logger
+import { useState } from "react";
+import { identity, makeObject_Primitives, useHook } from "shared-runtime";
+
+function Component() {
+  // The scopes for x and x2 are interleaved, so this is one scope with two values
+  const x = makeObject_Primitives();
+  const x2 = makeObject_Primitives();
+  useState(null);
+  identity(x);
+  identity(x2);
+
+  // We create a scope for all call expressions, but prune those with hook calls
+  // in this case it's _just_ a hook call, so we don't count this as pruned
+  const y = useHook();
+
+  const z = [];
+  for (let i = 0; i < 10; i++) {
+    // The scope for obj is pruned bc it's in a loop
+    const obj = makeObject_Primitives();
+    z.push(obj);
+  }
+
+  // Overall we expect two pruned scopes (for x+x2, and obj), with 3 pruned scope values.
+  return [x, x2, y, z];
+}
+
+export const FIXTURE_ENTRYPOINT = {
+  fn: Component,
+  params: [{}],
+};
+
+```
+
+## Code
+
+```javascript
+import { c as _c } from "react/compiler-runtime"; // @logger
+import { useState } from "react";
+import { identity, makeObject_Primitives, useHook } from "shared-runtime";
+
+function Component() {
+  const $ = _c(5);
+
+  const x = makeObject_Primitives();
+  const x2 = makeObject_Primitives();
+  useState(null);
+  identity(x);
+  identity(x2);
+
+  const y = useHook();
+  let z;
+  if ($[0] === Symbol.for("react.memo_cache_sentinel")) {
+    z = [];
+    for (let i = 0; i < 10; i++) {
+      const obj = makeObject_Primitives();
+      z.push(obj);
+    }
+    $[0] = z;
+  } else {
+    z = $[0];
+  }
+  let t0;
+  if ($[1] !== x || $[2] !== x2 || $[3] !== y) {
+    t0 = [x, x2, y, z];
+    $[1] = x;
+    $[2] = x2;
+    $[3] = y;
+    $[4] = t0;
+  } else {
+    t0 = $[4];
+  }
+  return t0;
+}
+
+export const FIXTURE_ENTRYPOINT = {
+  fn: Component,
+  params: [{}],
+};
+
+```
+
+## Logs
+
+```
+{"kind":"CompileSuccess","fnLoc":{"start":{"line":5,"column":0,"index":121},"end":{"line":26,"column":1,"index":813},"filename":"log-pruned-memoization.ts"},"fnName":"Component","memoSlots":5,"memoBlocks":2,"memoValues":2,"prunedMemoBlocks":2,"prunedMemoValues":3}
+```
+      
+### Eval output
+(kind: ok) [{"a":0,"b":"value1","c":true},{"a":0,"b":"value1","c":true},{"a":0,"b":"value1","c":true},[{"a":0,"b":"value1","c":true},{"a":0,"b":"value1","c":true},{"a":0,"b":"value1","c":true},{"a":0,"b":"value1","c":true},{"a":0,"b":"value1","c":true},{"a":0,"b":"value1","c":true},{"a":0,"b":"value1","c":true},{"a":0,"b":"value1","c":true},{"a":0,"b":"value1","c":true},{"a":0,"b":"value1","c":true}]]

compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/log-pruned-memoization.js

@@ -0,0 +1,31 @@
+// @logger
+import { useState } from "react";
+import { identity, makeObject_Primitives, useHook } from "shared-runtime";
+
+function Component() {
+  // The scopes for x and x2 are interleaved, so this is one scope with two values
+  const x = makeObject_Primitives();
+  const x2 = makeObject_Primitives();
+  useState(null);
+  identity(x);
+  identity(x2);
+
+  // We create a scope for all call expressions, but prune those with hook calls
+  // in this case it's _just_ a hook call, so we don't count this as pruned
+  const y = useHook();
+
+  const z = [];
+  for (let i = 0; i < 10; i++) {
+    // The scope for obj is pruned bc it's in a loop
+    const obj = makeObject_Primitives();
+    z.push(obj);
+  }
+
+  // Overall we expect two pruned scopes (for x+x2, and obj), with 3 pruned scope values.
+  return [x, x2, y, z];
+}
+
+export const FIXTURE_ENTRYPOINT = {
+  fn: Component,
+  params: [{}],
+};

compiler/packages/snap/src/compiler.ts

@@ -14,6 +14,8 @@ import * as t from "@babel/types";
 import assert from "assert";
 import type {
   CompilationMode,
+  Logger,
+  LoggerEvent,
   PanicThresholdOptions,
   PluginOptions,
 } from "babel-plugin-react-compiler/src/Entrypoint";
@@ -32,7 +34,7 @@ export function parseLanguage(source: string): "flow" | "typescript" {
 function makePluginOptions(
   firstLine: string,
   parseConfigPragmaFn: typeof ParseConfigPragma
-): PluginOptions {
+): [PluginOptions, Array<{ filename: string | null; event: LoggerEvent }>] {
   let gating = null;
   let enableEmitInstrumentForget = null;
   let enableEmitFreeze = null;
@@ -140,8 +142,18 @@ function makePluginOptions(
     );
   }
 
+  let logs: Array<{ filename: string | null; event: LoggerEvent }> = [];
+  let logger: Logger | null = null;
+  if (firstLine.includes("@logger")) {
+    logger = {
+      logEvent(filename: string | null, event: LoggerEvent): void {
+        logs.push({ filename, event });
+      },
+    };
+  }
+
   const config = parseConfigPragmaFn(firstLine);
-  return {
+  const options = {
     environment: {
       ...config,
       customHooks: new Map([
@@ -183,7 +195,7 @@ function makePluginOptions(
       enableChangeDetectionForDebugging,
     },
     compilationMode,
-    logger: null,
+    logger,
     gating,
     panicThreshold,
     noEmit: false,
@@ -193,6 +205,7 @@ function makePluginOptions(
     ignoreUseNoForget,
     enableReanimatedCheck: false,
   };
+  return [options, logs];
 }
 
 export function parseInput(
@@ -294,6 +307,7 @@ const FlowEvaluatorPresets = getEvaluatorPresets("flow");
 
 export type TransformResult = {
   forgetOutput: string;
+  logs: string | null;
   evaluatorCode: {
     original: string;
     forget: string;
@@ -330,12 +344,13 @@ export async function transformFixtureInput(
   /**
    * Get Forget compiled code
    */
+  const [options, logs] = makePluginOptions(firstLine, parseConfigPragmaFn);
   const forgetResult = transformFromAstSync(inputAst, input, {
     filename: virtualFilepath,
     highlightCode: false,
     retainLines: true,
     plugins: [
-      [plugin, makePluginOptions(firstLine, parseConfigPragmaFn)],
+      [plugin, options],
       "babel-plugin-fbt",
       "babel-plugin-fbt-runtime",
     ],
@@ -349,7 +364,7 @@ export async function transformFixtureInput(
     forgetResult?.code != null,
     "Expected BabelPluginReactForget to codegen successfully."
   );
-  const forgetOutput = forgetResult.code;
+  const forgetCode = forgetResult.code;
   let evaluatorCode = null;
 
   if (
@@ -363,7 +378,7 @@ export async function transformFixtureInput(
         forgetResult?.ast != null,
         "Expected BabelPluginReactForget ast."
       );
-      const result = transformFromAstSync(forgetResult.ast, forgetOutput, {
+      const result = transformFromAstSync(forgetResult.ast, forgetCode, {
         presets,
         filename: virtualFilepath,
         configFile: false,
@@ -415,10 +430,20 @@ export async function transformFixtureInput(
       original: originalEval,
     };
   }
+  const forgetOutput = await format(forgetCode, language);
+  let formattedLogs = null;
+  if (logs.length !== 0) {
+    formattedLogs = logs
+      .map(({ event }) => {
+        return JSON.stringify(event);
+      })
+      .join("\n");
+  }
   return {
     kind: "ok",
     value: {
-      forgetOutput: await format(forgetOutput, language),
+      forgetOutput,
+      logs: formattedLogs,
       evaluatorCode,
     },
   };

compiler/packages/snap/src/reporter.ts

@@ -22,6 +22,7 @@ export function writeOutputToString(
   input: string,
   compilerOutput: string | null,
   evaluatorOutput: string | null,
+  logs: string | null,
   errorMessage: string | null
 ) {
   // leading newline intentional
@@ -41,6 +42,14 @@ ${wrapWithTripleBackticks(compilerOutput, "javascript")}
     result += "\n";
   }
 
+  if (logs != null) {
+    result += `
+## Logs
+
+${wrapWithTripleBackticks(logs, null)}
+`;
+  }
+
   if (errorMessage != null) {
     result += `
 ## Error