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: 815cb8adc9de0134d8da134e06d6b632d9a86370
Pull Request resolved: #29810

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

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

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

@@ -329,6 +329,8 @@ export function compileProgram(
         fnName: compiledFn.id?.name ?? null,
         memoSlots: compiledFn.memoSlotsUsed,
         memoBlocks: compiledFn.memoBlocks,
+        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,18 @@ export type CodegenFunction = {
    * how many inputs/outputs each block has
    */
   memoBlocks: 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 +285,7 @@ function codegenReactiveFunction(
     return Err(cx.errors);
   }
 
-  const countMemoBlockVisitor = new CountMemoBlockVisitor();
+  const countMemoBlockVisitor = new CountMemoBlockVisitor(fn.env);
   visitReactiveFunction(fn, countMemoBlockVisitor, undefined);
 
   return Ok({
@@ -285,16 +298,56 @@ function codegenReactiveFunction(
     async: fn.async,
     memoSlotsUsed: cx.nextCacheIndex,
     memoBlocks: countMemoBlockVisitor.count,
+    prunedMemoBlocks: countMemoBlockVisitor.prunedMemoBlocks,
+    prunedMemoValues: countMemoBlockVisitor.prunedMemoValues,
   });
 }
 
 class CountMemoBlockVisitor extends ReactiveFunctionVisitor<void> {
+  env: Environment;
   count: number = 0;
+  prunedMemoBlocks: number = 0;
+  prunedMemoValues: number = 0;
+
+  constructor(env: Environment) {
+    super();
+    this.env = env;
+  }
 
   override visitScope(scope: ReactiveScopeBlock, state: void): void {
     this.count += 1;
     this.traverseScope(scope, 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);
+  }
 }
 
 function convertParameter(

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

@@ -0,0 +1,79 @@
+
+## Input
+
+```javascript
+// @logger
+import { useState } from "react";
+import { identity, makeObject_Primitives, useHook } from "shared-runtime";
+
+function Component() {
+  const x = makeObject_Primitives();
+  const x2 = makeObject_Primitives();
+  useState(null);
+  identity(x);
+  identity(x2);
+
+  const y = useHook();
+
+  const z = [];
+
+  for (let i = 0; i < 10; i++) {
+    const obj = makeObject_Primitives();
+    z.push(obj);
+  }
+
+  return [x, x2, y, z];
+}
+
+```
+
+## 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;
+}
+
+```
+
+## Logs
+
+```
+{"kind":"CompileSuccess","fnLoc":{"start":{"line":5,"column":0,"index":121},"end":{"line":22,"column":1,"index":431},"filename":"log-pruned-memoization.ts"},"fnName":"Component","memoSlots":5,"memoBlocks":2,"prunedMemoBlocks":2,"prunedMemoValues":3}
+```
+      
+### Eval output
+(kind: exception) Fixture not implemented

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

@@ -0,0 +1,22 @@
+// @logger
+import { useState } from "react";
+import { identity, makeObject_Primitives, useHook } from "shared-runtime";
+
+function Component() {
+  const x = makeObject_Primitives();
+  const x2 = makeObject_Primitives();
+  useState(null);
+  identity(x);
+  identity(x2);
+
+  const y = useHook();
+
+  const z = [];
+
+  for (let i = 0; i < 10; i++) {
+    const obj = makeObject_Primitives();
+    z.push(obj);
+  }
+
+  return [x, x2, y, z];
+}

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

compiler/packages/snap/src/runner-worker.ts

@@ -189,6 +189,7 @@ export async function transformFixture(
     input,
     snapOutput,
     sproutOutput,
+    compileResult?.logs ?? null,
     error
   );