[compiler] Correctly insert (Arrow)FunctionExpressions

Previously we would insert new (Arrow)FunctionExpressions as a sibling
of the original function. However this would break in the outlining case
as it would cause the original function expression's parent to become a
SequenceExpression, breaking a bunch of assumptions in the babel plugin.

To get around this, we synthesize a new VariableDeclaration to contain
the newly inserted function expression and therefore insert it as a true
sibling to the original function.

Yeah, it's kinda gross

ghstack-source-id: aefc2cd
Pull Request resolved: #30446

Author: poteto

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

@@ -183,13 +183,99 @@ export function createNewFunctionNode(
       transformedFn = fn;
       break;
     }
+    default: {
+      assertExhaustive(
+        originalFn.node,
+        `Creating unhandled function: ${originalFn.node}`,
+      );
+    }
   }
-
   // Avoid visiting the new transformed version
   ALREADY_COMPILED.add(transformedFn);
   return transformedFn;
 }
 
+function insertNewFunctionNode(
+  originalFn: BabelFn,
+  compiledFn: CodegenFunction,
+): NodePath<t.Function> {
+  switch (originalFn.type) {
+    case 'FunctionDeclaration': {
+      return originalFn.insertAfter(
+        createNewFunctionNode(originalFn, compiledFn),
+      )[0]!;
+    }
+    /**
+     * This is pretty gross but we can't just append an (Arrow)FunctionExpression as a sibling of
+     * the original function. If the original function was itself an (Arrow)FunctionExpression,
+     * this would cause its parent to become a SequenceExpression instead which breaks a bunch of
+     * assumptions elsewhere in the plugin.
+     *
+     * To get around this, we synthesize a new VariableDeclaration holding the compiled function
+     * expression and insert it as a true sibling (ie within the Program's block statements).
+     */
+    case 'ArrowFunctionExpression':
+    case 'FunctionExpression': {
+      const funcExpr = createNewFunctionNode(originalFn, compiledFn);
+      CompilerError.invariant(
+        t.isArrowFunctionExpression(funcExpr) ||
+          t.isFunctionExpression(funcExpr),
+        {
+          reason: 'Expected an (arrow) function expression to be created',
+          description: `Got: ${funcExpr.type}`,
+          loc: null,
+        },
+      );
+      CompilerError.invariant(compiledFn.id != null, {
+        reason: 'Expected compiled function to have an identifier',
+        loc: null,
+      });
+      CompilerError.invariant(originalFn.parentPath.isVariableDeclarator(), {
+        reason: 'Expected a variable declarator parent',
+        loc: null,
+      });
+      const varDecl = originalFn.parentPath.parentPath;
+      varDecl.insertAfter(
+        t.variableDeclaration('const', [
+          t.variableDeclarator(compiledFn.id, funcExpr),
+        ]),
+      )[0]!;
+      CompilerError.invariant(typeof varDecl.key === 'number', {
+        reason:
+          'Just Babel things: expected the VariableDeclaration containing the compiled function expression to have a number key',
+        loc: null,
+      });
+      const insertedCompiledFnVarDecl = varDecl.getSibling(varDecl.key + 1);
+      CompilerError.invariant(
+        insertedCompiledFnVarDecl.isVariableDeclaration(),
+        {
+          reason: 'Expected inserted sibling to be a VariableDeclaration',
+          loc: null,
+        },
+      );
+      const insertedFuncExpr = insertedCompiledFnVarDecl
+        .get('declarations')[0]!
+        .get('init')!;
+      CompilerError.invariant(
+        insertedFuncExpr.isArrowFunctionExpression() ||
+          insertedFuncExpr.isFunctionExpression(),
+        {
+          reason: 'Expected inserted (arrow) function expression',
+          description: `Got: ${insertedFuncExpr}`,
+          loc: null,
+        },
+      );
+      return insertedFuncExpr;
+    }
+    default: {
+      assertExhaustive(
+        originalFn,
+        `Inserting unhandled function: ${originalFn}`,
+      );
+    }
+  }
+}
+
 /*
  * This is a hack to work around what seems to be a Babel bug. Babel doesn't
  * consistently respect the `skip()` function to avoid revisiting a node within
@@ -407,9 +493,7 @@ export function compileProgram(
         reason: 'Unexpected nested outlined functions',
         loc: outlined.fn.loc,
       });
-      const fn = current.fn.insertAfter(
-        createNewFunctionNode(current.fn, outlined.fn),
-      )[0]!;
+      const fn = insertNewFunctionNode(current.fn, outlined.fn);
       fn.skip();
       ALREADY_COMPILED.add(fn.node);
       if (outlined.type !== null) {

compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/error.bug-outlining-in-func-expr.expect.md

@@ -0,0 +1,49 @@
+
+## Input
+
+```javascript
+// @debug
+const Component2 = props => {
+  return (
+    <ul>
+      {props.items.map(item => (
+        <li key={item.id}>{item.name}</li>
+      ))}
+    </ul>
+  );
+};
+
+```
+
+## Code
+
+```javascript
+import { c as _c } from "react/compiler-runtime"; // @debug
+const Component2 = (props) => {
+  const $ = _c(4);
+  let t0;
+  if ($[0] !== props.items) {
+    t0 = props.items.map(_temp);
+    $[0] = props.items;
+    $[1] = t0;
+  } else {
+    t0 = $[1];
+  }
+  let t1;
+  if ($[2] !== t0) {
+    t1 = <ul>{t0}</ul>;
+    $[2] = t0;
+    $[3] = t1;
+  } else {
+    t1 = $[3];
+  }
+  return t1;
+};
+const _temp = (item) => {
+  return <li key={item.id}>{item.name}</li>;
+};
+
+```
+      
+### Eval output
+(kind: exception) Fixture not implemented

compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/outlining-in-func-expr.expect.md

@@ -1,3 +1,4 @@
+// @debug
 const Component2 = props => {
   return (
     <ul>