diff --git a/compiler/packages/babel-plugin-react-compiler/src/Entrypoint/Program.ts b/compiler/packages/babel-plugin-react-compiler/src/Entrypoint/Program.ts index 2880e9283c77..aae46f38801c 100644 --- a/compiler/packages/babel-plugin-react-compiler/src/Entrypoint/Program.ts +++ b/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: { diff --git a/compiler/packages/babel-plugin-react-compiler/src/Utils/OutlinedByReactCompiler.ts b/compiler/packages/babel-plugin-react-compiler/src/Utils/OutlinedByReactCompiler.ts new file mode 100644 index 000000000000..cf2767464eab --- /dev/null +++ b/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; +} diff --git a/compiler/packages/babel-plugin-react-compiler/src/__tests__/OutlinedByReactCompiler-test.ts b/compiler/packages/babel-plugin-react-compiler/src/__tests__/OutlinedByReactCompiler-test.ts new file mode 100644 index 000000000000..4e7a70f30d57 --- /dev/null +++ b/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 ( + + + + ); + })} + + ); + } + + 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); + }); +});