Mark compiler-outlined functions in AST

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

@@ -16,6 +16,10 @@ import {ExternalFunction, ReactFunctionType} from '../HIR/Environment';
 import {CodegenFunction} from '../ReactiveScopes';
 import {isComponentDeclaration} from '../Utils/ComponentDeclaration';
 import {isHookDeclaration} from '../Utils/HookDeclaration';
+import {
+  copyOutlinedByReactCompilerMarker,
+  markOutlinedByReactCompiler,
+} from '../Utils/OutlinedByReactCompiler';
 import {assertExhaustive} from '../Utils/utils';
 import {insertGatedFunctionDeclaration} from './Gating';
 import {
@@ -240,7 +244,7 @@ export function createNewFunctionNode(
         params: compiledFn.params,
         body: compiledFn.body,
       };
-      transformedFn = fn;
+      transformedFn = copyOutlinedByReactCompilerMarker(originalFn.node, fn);
       break;
     }
     case 'ArrowFunctionExpression': {
@@ -287,9 +291,20 @@ function insertNewOutlinedFunctionNode(
 ): BabelFn {
   switch (originalFn.type) {
     case 'FunctionDeclaration': {
-      return originalFn.insertAfter(
+      const insertedFuncDecl = originalFn.insertAfter(
         createNewFunctionNode(originalFn, compiledFn),
       )[0]!;
+      CompilerError.invariant(insertedFuncDecl.isFunctionDeclaration(), {
+        reason: 'Expected inserted outlined function declaration',
+        description: `Got: ${insertedFuncDecl}`,
+        loc: insertedFuncDecl.node?.loc ?? GeneratedSource,
+      });
+      /*
+       * Subsequent Babel passes need to distinguish compiler-outlined helpers
+       * from user-authored top-level declarations.
+       */
+      markOutlinedByReactCompiler(insertedFuncDecl.node);
+      return insertedFuncDecl;
     }
     /**
      * We can't just append the outlined function as a sibling of the original function if it is an
@@ -317,6 +332,11 @@ function insertNewOutlinedFunctionNode(
         description: `Got: ${insertedFuncDecl}`,
         loc: insertedFuncDecl.node?.loc ?? GeneratedSource,
       });
+      /*
+       * Subsequent Babel passes need to distinguish compiler-outlined helpers
+       * from user-authored top-level declarations.
+       */
+      markOutlinedByReactCompiler(insertedFuncDecl.node);
       return insertedFuncDecl;
     }
     default: {

compiler/packages/babel-plugin-react-compiler/src/Utils/OutlinedByReactCompiler.ts

@@ -0,0 +1,37 @@
+/**
+ * Copyright (c) Meta Platforms, Inc. and affiliates.
+ *
+ * This source code is licensed under the MIT license found in the
+ * LICENSE file in the root directory of this source tree.
+ */
+
+import * as t from '@babel/types';
+
+export type OutlinedByReactCompilerFunctionDeclaration =
+  t.FunctionDeclaration & {
+    isOutlinedByReactCompiler: boolean;
+  };
+
+export function isOutlinedByReactCompiler(
+  node: t.FunctionDeclaration,
+): node is OutlinedByReactCompilerFunctionDeclaration {
+  return Object.prototype.hasOwnProperty.call(node, 'isOutlinedByReactCompiler');
+}
+
+export function markOutlinedByReactCompiler(
+  node: t.FunctionDeclaration,
+): OutlinedByReactCompilerFunctionDeclaration {
+  const outlinedNode = node as OutlinedByReactCompilerFunctionDeclaration;
+  outlinedNode.isOutlinedByReactCompiler = true;
+  return outlinedNode;
+}
+
+export function copyOutlinedByReactCompilerMarker(
+  source: t.FunctionDeclaration,
+  target: t.FunctionDeclaration,
+): t.FunctionDeclaration {
+  if (isOutlinedByReactCompiler(source)) {
+    return markOutlinedByReactCompiler(target);
+  }
+  return target;
+}

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

@@ -0,0 +1,143 @@
+/**
+ * Copyright (c) Meta Platforms, Inc. and affiliates.
+ *
+ * This source code is licensed under the MIT license found in the
+ * LICENSE file in the root directory of this source tree.
+ */
+
+import * as t from '@babel/types';
+import {runBabelPluginReactCompiler} from '../Babel/RunReactCompilerBabelPlugin';
+import {PluginOptions} from '../Entrypoint';
+import {isOutlinedByReactCompiler} from '../Utils/OutlinedByReactCompiler';
+
+function compile(source: string, options: PluginOptions = {}): t.File {
+  const result = runBabelPluginReactCompiler(
+    source,
+    'test.js',
+    'flow',
+    {
+      ...options,
+      compilationMode: 'all',
+      environment: {
+        enableCustomTypeDefinitionForReanimated: true,
+        ...options.environment,
+      },
+    },
+    true,
+  );
+
+  expect(result.ast).not.toBeNull();
+  expect(result.ast).not.toBeUndefined();
+
+  return result.ast as t.File;
+}
+
+function getFunctionDeclaration(
+  program: t.Program,
+  name: string,
+): t.FunctionDeclaration {
+  const fn = program.body.find(
+    statement =>
+      t.isFunctionDeclaration(statement) && statement.id?.name === name,
+  );
+
+  expect(fn).toBeDefined();
+  expect(t.isFunctionDeclaration(fn)).toBe(true);
+
+  return fn as t.FunctionDeclaration;
+}
+
+function isRecompiledByReactCompiler(fn: t.FunctionDeclaration): boolean {
+  return fn.body.body.some(
+    statement =>
+      t.isVariableDeclaration(statement) &&
+      statement.declarations.some(
+        declarator =>
+          t.isIdentifier(declarator.id, {name: '$'}) &&
+          t.isCallExpression(declarator.init) &&
+          t.isIdentifier(declarator.init.callee, {name: '_c'}),
+      ),
+  );
+}
+
+describe('outlined function markers', () => {
+  it('marks outlined helpers in the reanimated outlining repro', () => {
+    const ast = compile(`
+      import {useDerivedValue} from 'react-native-reanimated';
+
+      const TestComponent = ({number}) => {
+        const keyToIndex = useDerivedValue(() => [1, 2, 3].map(() => null));
+        return null;
+      };
+    `);
+
+    expect(
+      isOutlinedByReactCompiler(getFunctionDeclaration(ast.program, '_temp')),
+    ).toBe(true);
+    expect(
+      isOutlinedByReactCompiler(getFunctionDeclaration(ast.program, '_temp2')),
+    ).toBe(true);
+  });
+
+  it('does not mark user-authored declarations', () => {
+    const ast = compile(`
+      function useFoo() {
+        return [1, 2, 3].map(() => 1);
+      }
+    `);
+
+    expect(
+      isOutlinedByReactCompiler(getFunctionDeclaration(ast.program, '_temp')),
+    ).toBe(true);
+    expect(
+      isOutlinedByReactCompiler(getFunctionDeclaration(ast.program, 'useFoo')),
+    ).toBe(false);
+  });
+
+  it('preserves the marker for JSX-outlined functions after recompilation', () => {
+    const ast = compile(
+      `
+      function Component({arr}) {
+        const x = useX();
+        return (
+          <>
+            {arr.map((i, id) => {
+              return (
+                <Bar key={id} x={x}>
+                  <Baz i={i}></Baz>
+                </Bar>
+              );
+            })}
+          </>
+        );
+      }
+
+      function Bar({x, children}) {
+        return (
+          <>
+            {x}
+            {children}
+          </>
+        );
+      }
+
+      function Baz({i}) {
+        return i;
+      }
+
+      function useX() {
+        return 'x';
+      }
+    `,
+      {
+        environment: {
+          enableJsxOutlining: true,
+        },
+      },
+    );
+
+    const outlinedComponent = getFunctionDeclaration(ast.program, '_temp');
+    expect(isRecompiledByReactCompiler(outlinedComponent)).toBe(true);
+    expect(isOutlinedByReactCompiler(outlinedComponent)).toBe(true);
+  });
+});