[compiler] Add experimental trace-tape runtime and opt-in emission prototype

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

@@ -170,6 +170,12 @@ export type PluginOptions = Partial<{
    */
   enableReanimatedCheck: boolean;
 
+  /**
+   * Experimental research surface. When enabled, functions annotated with
+   * 'use trace tape' may emit a tiny trace-tape companion artifact.
+   */
+  enableEmitTraceTape: boolean;
+
   /**
    * The minimum major version of React that the compiler should emit code for. If the target is 19
    * or higher, the compiler emits direct imports of React runtime APIs needed by the compiler. On
@@ -317,6 +323,7 @@ export const defaultOptions: ParsedPluginOptions = {
     return filename.indexOf('node_modules') === -1;
   },
   enableReanimatedCheck: true,
+  enableEmitTraceTape: false,
   customOptOutDirectives: null,
   target: '19',
 };

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

@@ -47,6 +47,7 @@ export type CompilerPass = {
 export const OPT_IN_DIRECTIVES = new Set(['use forget', 'use memo']);
 export const OPT_OUT_DIRECTIVES = new Set(['use no forget', 'use no memo']);
 const DYNAMIC_GATING_DIRECTIVE = new RegExp('^use memo if\\(([^\\)]*)\\)$');
+const TRACE_TAPE_DIRECTIVE = 'use trace tape';
 
 export function tryFindDirectiveEnablingMemoization(
   directives: Array<t.Directive>,
@@ -280,6 +281,237 @@ export function createNewFunctionNode(
   return transformedFn;
 }
 
+function hasTraceTapeDirective(directives: Array<t.Directive>): boolean {
+  return directives.some(
+    directive => directive.value.value === TRACE_TAPE_DIRECTIVE,
+  );
+}
+
+function getTraceTapeParamName(fn: BabelFn): string | null {
+  if (fn.node.params.length !== 1) {
+    return null;
+  }
+  const [param] = fn.node.params;
+  return t.isIdentifier(param) ? param.name : null;
+}
+
+function getTraceTapePathSegments(
+  expression: t.Expression,
+  paramName: string,
+): Array<string> | null {
+  if (!t.isMemberExpression(expression) || expression.computed) {
+    return null;
+  }
+  if (!t.isIdentifier(expression.property)) {
+    return null;
+  }
+  if (t.isIdentifier(expression.object)) {
+    return expression.object.name === paramName ? [expression.property.name] : null;
+  }
+  if (!t.isMemberExpression(expression.object)) {
+    return null;
+  }
+  const prefix = getTraceTapePathSegments(expression.object, paramName);
+  return prefix == null ? null : [...prefix, expression.property.name];
+}
+
+function createTraceTapeInputExpression(
+  inputName: string,
+  pathSegments: Array<string>,
+): t.Expression {
+  let expression: t.Expression = t.identifier(inputName);
+  for (const segment of pathSegments) {
+    expression = t.memberExpression(expression, t.identifier(segment));
+  }
+  return expression;
+}
+
+function createTraceTapeSelector(
+  traceSelectorName: string,
+  inputName: string,
+  pathSegments: Array<string>,
+): t.Expression {
+  return t.callExpression(t.identifier(traceSelectorName), [
+    t.stringLiteral(pathSegments.join('.')),
+    t.arrowFunctionExpression(
+      [t.identifier(inputName)],
+      createTraceTapeInputExpression(inputName, pathSegments),
+    ),
+  ]);
+}
+
+function createTraceTapeComputeFn(
+  inputName: string,
+  pathSegments: Array<string>,
+): t.ArrowFunctionExpression {
+  return t.arrowFunctionExpression(
+    [t.identifier(inputName)],
+    createTraceTapeInputExpression(inputName, pathSegments),
+  );
+}
+
+function maybeCreateTraceTapeArtifact(
+  fn: BabelFn,
+  programContext: ProgramContext,
+): t.Statement | null {
+  if (!programContext.opts.enableEmitTraceTape || !fn.isFunctionDeclaration()) {
+    return null;
+  }
+  if (fn.parentPath.isExportDefaultDeclaration()) {
+    return null;
+  }
+  const functionName = fn.node.id?.name;
+  if (functionName == null || fn.node.body.type !== 'BlockStatement') {
+    return null;
+  }
+  if (!hasTraceTapeDirective(fn.node.body.directives)) {
+    return null;
+  }
+  const paramName = getTraceTapeParamName(fn);
+  if (paramName == null) {
+    return null;
+  }
+  if (fn.node.body.body.length !== 1) {
+    return null;
+  }
+  const [statement] = fn.node.body.body;
+  if (!t.isReturnStatement(statement) || statement.argument == null) {
+    return null;
+  }
+  if (!t.isJSXElement(statement.argument)) {
+    return null;
+  }
+
+  const traceSessionImport = programContext.addImportSpecifier(
+    {
+      source: programContext.reactRuntimeModule,
+      importSpecifierName: 'experimental_createRenderTraceSession',
+    },
+    '_traceTapeSession',
+  );
+  const traceSelectorImport = programContext.addImportSpecifier(
+    {
+      source: programContext.reactRuntimeModule,
+      importSpecifierName: 'experimental_createTraceSelector',
+    },
+    '_traceTapeSelector',
+  );
+
+  const root = statement.argument;
+  const traceName = programContext.newUid('trace');
+  const inputName = programContext.newUid('input');
+  const operations: Array<t.Statement> = [];
+
+  for (const attribute of root.openingElement.attributes) {
+    if (t.isJSXSpreadAttribute(attribute)) {
+      return null;
+    }
+    if (!t.isJSXIdentifier(attribute.name)) {
+      return null;
+    }
+    if (
+      attribute.value == null ||
+      t.isStringLiteral(attribute.value) ||
+      !t.isJSXExpressionContainer(attribute.value) ||
+      t.isJSXEmptyExpression(attribute.value.expression)
+    ) {
+      continue;
+    }
+    const pathSegments = getTraceTapePathSegments(
+      attribute.value.expression,
+      paramName,
+    );
+    if (pathSegments == null) {
+      return null;
+    }
+    operations.push(
+      t.expressionStatement(
+        t.callExpression(
+          t.memberExpression(t.identifier(traceName), t.identifier('attr')),
+          [
+            t.stringLiteral('root'),
+            t.stringLiteral(attribute.name.name),
+            t.arrayExpression([
+              createTraceTapeSelector(
+                traceSelectorImport.name,
+                inputName,
+                pathSegments,
+              ),
+            ]),
+            createTraceTapeComputeFn(inputName, pathSegments),
+          ],
+        ),
+      ),
+    );
+  }
+
+  for (const [index, child] of root.children.entries()) {
+    if (t.isJSXText(child)) {
+      if (child.value.trim().length !== 0) {
+        return null;
+      }
+      continue;
+    }
+    if (
+      !t.isJSXExpressionContainer(child) ||
+      t.isJSXEmptyExpression(child.expression)
+    ) {
+      return null;
+    }
+    const pathSegments = getTraceTapePathSegments(child.expression, paramName);
+    if (pathSegments == null) {
+      return null;
+    }
+    operations.push(
+      t.expressionStatement(
+        t.callExpression(
+          t.memberExpression(t.identifier(traceName), t.identifier('text')),
+          [
+            t.stringLiteral(`root.children.${index}`),
+            t.arrayExpression([
+              createTraceTapeSelector(
+                traceSelectorImport.name,
+                inputName,
+                pathSegments,
+              ),
+            ]),
+            createTraceTapeComputeFn(inputName, pathSegments),
+          ],
+        ),
+      ),
+    );
+  }
+
+  if (operations.length === 0) {
+    return null;
+  }
+
+  return t.expressionStatement(
+    t.assignmentExpression(
+      '=',
+      t.memberExpression(
+        t.identifier(functionName),
+        t.identifier('__traceTape'),
+      ),
+      t.functionExpression(
+        null,
+        [],
+        t.blockStatement([
+          t.returnStatement(
+            t.callExpression(t.identifier(traceSessionImport.name), [
+              t.functionExpression(
+                null,
+                [t.identifier(traceName), t.identifier(inputName)],
+                t.blockStatement(operations),
+              ),
+            ]),
+          ),
+        ]),
+      ),
+    ),
+  );
+}
+
 function insertNewOutlinedFunctionNode(
   program: NodePath<t.Program>,
   originalFn: BabelFn,
@@ -769,6 +1001,13 @@ function applyCompiledFunctions(
         referencedBeforeDeclared.has(result),
       );
     } else {
+      const traceTapeArtifact =
+        kind === 'original'
+          ? maybeCreateTraceTapeArtifact(originalFn, programContext)
+          : null;
+      if (traceTapeArtifact != null) {
+        originalFn.insertAfter(traceTapeArtifact);
+      }
       originalFn.replaceWith(transformedFn);
     }
   }

compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/trace-tape-annotation.expect.md

@@ -0,0 +1,67 @@
+
+## Input
+
+```javascript
+// @compilationMode:"annotation" @enableEmitTraceTape
+
+function Foo(props) {
+  'use memo';
+  'use trace tape';
+  return <div title={props.title}>{props.count}</div>;
+}
+
+export const FIXTURE_ENTRYPOINT = {
+  fn: Foo,
+  params: [{title: 'hello', count: 3}],
+};
+```
+
+## Code
+
+```javascript
+import {
+  c as _c,
+  experimental_createRenderTraceSession as _traceTapeSession,
+  experimental_createTraceSelector as _traceTapeSelector,
+} from "react/compiler-runtime"; // @compilationMode:"annotation" @enableEmitTraceTape
+
+function Foo(props) {
+  "use memo";
+  "use trace tape";
+  const $ = _c(3);
+  let t0;
+  if ($[0] !== props.count || $[1] !== props.title) {
+    t0 = <div title={props.title}>{props.count}</div>;
+    $[0] = props.count;
+    $[1] = props.title;
+    $[2] = t0;
+  } else {
+    t0 = $[2];
+  }
+  return t0;
+}
+Foo.__traceTape = function () {
+  return _traceTapeSession(function (trace, input) {
+    trace.attr(
+      "root",
+      "title",
+      [_traceTapeSelector("title", (input) => input.title)],
+      (input) => input.title,
+    );
+    trace.text(
+      "root.children.0",
+      [_traceTapeSelector("count", (input) => input.count)],
+      (input) => input.count,
+    );
+  });
+};
+
+export const FIXTURE_ENTRYPOINT = {
+  fn: Foo,
+  params: [{ title: "hello", count: 3 }],
+};
+
+```
+
+### Eval output
+(kind: ok) <div title="hello">3</div>

compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/trace-tape-annotation.js

@@ -0,0 +1,12 @@
+// @compilationMode:"annotation" @enableEmitTraceTape
+
+function Foo(props) {
+  'use memo';
+  'use trace tape';
+  return <div title={props.title}>{props.count}</div>;
+}
+
+export const FIXTURE_ENTRYPOINT = {
+  fn: Foo,
+  params: [{title: 'hello', count: 3}],
+};

compiler/packages/babel-plugin-react-compiler/src/__tests__/parseConfigPragma-test.ts

@@ -14,16 +14,18 @@ describe('parseConfigPragmaForTests()', () => {
 
     // Validate defaults first to make sure that the parser is getting the value from the pragma,
     // and not just missing it and getting the default value
+    expect(defaultOptions.enableEmitTraceTape).toBe(false);
     expect(defaultConfig.enableForest).toBe(false);
     expect(defaultConfig.validateNoSetStateInEffects).toBe(false);
     expect(defaultConfig.validateNoSetStateInRender).toBe(true);
 
     const config = parseConfigPragmaForTests(
-      '@enableForest @validateNoSetStateInEffects:true @validateNoSetStateInRender:false',
+      '@enableForest @enableEmitTraceTape @validateNoSetStateInEffects:true @validateNoSetStateInRender:false',
       {compilationMode: defaultOptions.compilationMode},
     );
     expect(config).toEqual({
       ...defaultOptions,
+      enableEmitTraceTape: true,
       panicThreshold: 'all_errors',
       environment: {
         ...defaultOptions.environment,