Repository: facebook/react
Pull request: [compiler][bugfix] Bail out when a memo block declares hoisted fns


Note that bailing out adds false positives for hoisted functions whose only references are within other functions. For example, this rewrite would be safe.
```js
// source program
  function foo() {
    return bar();
  }
  function bar() {
    return 42;
  }

// compiler output
let bar;
if (/* deps changed */) {
  function foo() {
    return bar();
  }
  bar = function bar() {
    return 42;
  }
}
```
These false positives are difficult to detect because any maybe-call of foo before the definition of bar would be invalid.

Instead of bailing out, we should rewrite hoisted function declarations to the following form.
```js
let bar$0;
if (/* deps changed */) {
  // All references within the declaring memo block
  // or before the function declaration should use
  // the original identifier `bar`
  function foo() {
    return bar();
  }
  function bar() {
    return 42;
  }
  bar$0 = bar;
}
// All references after the declaring memo block
// or after the function declaration should use
// the rewritten declaration `bar$0`
```


Changed files:
filename: compiler/packages/babel-plugin-react-compiler/src/ReactiveScopes/PruneHoistedContexts.ts
patch: @@ -5,10 +5,13 @@
  * LICENSE file in the root directory of this source tree.
  */
 
+import {CompilerError} from '..';
 import {
   convertHoistedLValueKind,
   IdentifierId,
+  InstructionId,
   InstructionKind,
+  Place,
   ReactiveFunction,
   ReactiveInstruction,
   ReactiveScopeBlock,
@@ -24,31 +27,78 @@ import {
 /*
  * Prunes DeclareContexts lowered for HoistedConsts, and transforms any references back to its
  * original instruction kind.
+ *
+ * Also detects and bails out on context variables which are:
+ * - function declarations, which are hoisted by JS engines to the nearest block scope
+ * - referenced before they are defined (i.e. having a `DeclareContext HoistedConst`)
+ * - declared
+ *
+ * This is because React Compiler converts a `function foo()` function declaration to
+ * 1. a `let foo;` declaration before reactive memo blocks
+ * 2. a `foo = function foo() {}` assignment within the block
+ *
+ * This means references before the assignment are invalid (see fixture
+ * error.todo-functiondecl-hoisting)
  */
 export function pruneHoistedContexts(fn: ReactiveFunction): void {
   visitReactiveFunction(fn, new Visitor(), {
     activeScopes: empty(),
+    uninitialized: new Map(),
   });
 }
 
 type VisitorState = {
   activeScopes: Stack<Set<IdentifierId>>;
+  uninitialized: Map<
+    IdentifierId,
+    | {
+        kind: 'unknown-kind';
+      }
+    | {
+        kind: 'func';
+        definition: Place | null;
+      }
+  >;
 };
 
 class Visitor extends ReactiveFunctionTransform<VisitorState> {
   override visitScope(scope: ReactiveScopeBlock, state: VisitorState): void {
     state.activeScopes = state.activeScopes.push(
       new Set(scope.scope.declarations.keys()),
     );
+    /**
+     * Add declared but not initialized / assigned variables. This may include
+     * function declarations that escape the memo block.
+     */
+    for (const decl of scope.scope.declarations.values()) {
+      state.uninitialized.set(decl.identifier.id, {kind: 'unknown-kind'});
+    }
     this.traverseScope(scope, state);
     state.activeScopes.pop();
+    for (const decl of scope.scope.declarations.values()) {
+      state.uninitialized.delete(decl.identifier.id);
+    }
+  }
+  override visitPlace(
+    _id: InstructionId,
+    place: Place,
+    state: VisitorState,
+  ): void {
+    const maybeHoistedFn = state.uninitialized.get(place.identifier.id);
+    if (
+      maybeHoistedFn?.kind === 'func' &&
+      maybeHoistedFn.definition !== place
+    ) {
+      CompilerError.throwTodo({
+        reason: '[PruneHoistedContexts] Rewrite hoisted function references',
+        loc: place.loc,
+      });
+    }
   }
   override transformInstruction(
     instruction: ReactiveInstruction,
     state: VisitorState,
   ): Transformed<ReactiveStatement> {
-    this.visitInstruction(instruction, state);
-
     /**
      * Remove hoisted declarations to preserve TDZ
      */
@@ -57,6 +107,18 @@ class Visitor extends ReactiveFunctionTransform<VisitorState> {
         instruction.value.lvalue.kind,
       );
       if (maybeNonHoisted != null) {
+        if (
+          maybeNonHoisted === InstructionKind.Function &&
+          state.uninitialized.has(instruction.value.lvalue.place.identifier.id)
+        ) {
+          state.uninitialized.set(
+            instruction.value.lvalue.place.identifier.id,
+            {
+              kind: 'func',
+              definition: null,
+            },
+          );
+        }
         return {kind: 'remove'};
       }
     }
@@ -65,18 +127,44 @@ class Visitor extends ReactiveFunctionTransform<VisitorState> {
       instruction.value.lvalue.kind !== InstructionKind.Reassign
     ) {
       /**
-       * Rewrite StoreContexts let/const/functions that will be pre-declared in
+       * Rewrite StoreContexts let/const that will be pre-declared in
        * codegen to reassignments.
        */
       const lvalueId = instruction.value.lvalue.place.identifier.id;
       const isDeclaredByScope = state.activeScopes.find(scope =>
         scope.has(lvalueId),
       );
       if (isDeclaredByScope) {
-        instruction.value.lvalue.kind = InstructionKind.Reassign;
+        if (
+          instruction.value.lvalue.kind === InstructionKind.Let ||
+          instruction.value.lvalue.kind === InstructionKind.Const
+        ) {
+          instruction.value.lvalue.kind = InstructionKind.Reassign;
+        } else if (instruction.value.lvalue.kind === InstructionKind.Function) {
+          const maybeHoistedFn = state.uninitialized.get(lvalueId);
+          if (maybeHoistedFn != null) {
+            CompilerError.invariant(maybeHoistedFn.kind === 'func', {
+              reason: '[PruneHoistedContexts] Unexpected hoisted function',
+              loc: instruction.loc,
+            });
+            maybeHoistedFn.definition = instruction.value.lvalue.place;
+            /**
+             * References to hoisted functions are now "safe" as variable assignments
+             * have finished.
+             */
+            state.uninitialized.delete(lvalueId);
+          }
+        } else {
+          CompilerError.throwTodo({
+            reason: '[PruneHoistedContexts] Unexpected kind',
+            description: `(${instruction.value.lvalue.kind})`,
+            loc: instruction.loc,
+          });
+        }
       }
     }
 
+    this.visitInstruction(instruction, state);
     return {kind: 'keep'};
   }
 }

filename: compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/bug-functiondecl-hoisting.expect.md
patch: @@ -1,79 +0,0 @@
-
-## Input
-
-```javascript
-import {Stringify} from 'shared-runtime';
-
-/**
- * Fixture currently fails with
- * Found differences in evaluator results
- *   Non-forget (expected):
- *   (kind: ok) <div>{"result":{"value":2},"fn":{"kind":"Function","result":{"value":2}},"shouldInvokeFns":true}</div>
- *   Forget:
- *   (kind: exception) bar is not a function
- */
-function Foo({value}) {
-  const result = bar();
-  function bar() {
-    return {value};
-  }
-  return <Stringify result={result} fn={bar} shouldInvokeFns={true} />;
-}
-
-export const FIXTURE_ENTRYPOINT = {
-  fn: Foo,
-  params: [{value: 2}],
-};
-
-```
-
-## Code
-
-```javascript
-import { c as _c } from "react/compiler-runtime";
-import { Stringify } from "shared-runtime";
-
-/**
- * Fixture currently fails with
- * Found differences in evaluator results
- *   Non-forget (expected):
- *   (kind: ok) <div>{"result":{"value":2},"fn":{"kind":"Function","result":{"value":2}},"shouldInvokeFns":true}</div>
- *   Forget:
- *   (kind: exception) bar is not a function
- */
-function Foo(t0) {
-  const $ = _c(6);
-  const { value } = t0;
-  let bar;
-  let result;
-  if ($[0] !== value) {
-    result = bar();
-    bar = function bar() {
-      return { value };
-    };
-    $[0] = value;
-    $[1] = bar;
-    $[2] = result;
-  } else {
-    bar = $[1];
-    result = $[2];
-  }
-  let t1;
-  if ($[3] !== bar || $[4] !== result) {
-    t1 = <Stringify result={result} fn={bar} shouldInvokeFns={true} />;
-    $[3] = bar;
-    $[4] = result;
-    $[5] = t1;
-  } else {
-    t1 = $[5];
-  }
-  return t1;
-}
-
-export const FIXTURE_ENTRYPOINT = {
-  fn: Foo,
-  params: [{ value: 2 }],
-};
-
-```
-      
\ No newline at end of file

filename: compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/error.todo-functiondecl-hoisting.expect.md
patch: @@ -0,0 +1,43 @@
+
+## Input
+
+```javascript
+import {Stringify} from 'shared-runtime';
+
+/**
+ * Fixture currently fails with
+ * Found differences in evaluator results
+ *   Non-forget (expected):
+ *   (kind: ok) <div>{"result":{"value":2},"fn":{"kind":"Function","result":{"value":2}},"shouldInvokeFns":true}</div>
+ *   Forget:
+ *   (kind: exception) bar is not a function
+ */
+function Foo({value}) {
+  const result = bar();
+  function bar() {
+    return {value};
+  }
+  return <Stringify result={result} fn={bar} shouldInvokeFns={true} />;
+}
+
+export const FIXTURE_ENTRYPOINT = {
+  fn: Foo,
+  params: [{value: 2}],
+};
+
+```
+
+
+## Error
+
+```
+  10 |  */
+  11 | function Foo({value}) {
+> 12 |   const result = bar();
+     |                  ^^^ Todo: [PruneHoistedContexts] Rewrite hoisted function references (12:12)
+  13 |   function bar() {
+  14 |     return {value};
+  15 |   }
+```
+          
+      
\ No newline at end of file

filename: compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/error.todo-valid-functiondecl-hoisting.expect.md
patch: @@ -0,0 +1,46 @@
+
+## Input
+
+```javascript
+import {Stringify} from 'shared-runtime';
+/**
+ * Also see error.todo-functiondecl-hoisting.tsx which shows *invalid*
+ * compilation cases.
+ *
+ * This bailout specifically is a false positive for since this function's only
+ * reference-before-definition are within other functions which are not invoked.
+ */
+function Foo() {
+  'use memo';
+
+  function foo() {
+    return bar();
+  }
+  function bar() {
+    return 42;
+  }
+
+  return <Stringify fn1={foo} fn2={bar} shouldInvokeFns={true} />;
+}
+
+export const FIXTURE_ENTRYPOINT = {
+  fn: Foo,
+  params: [],
+};
+
+```
+
+
+## Error
+
+```
+  13 |     return bar();
+  14 |   }
+> 15 |   function bar() {
+     |            ^^^ Todo: [PruneHoistedContexts] Rewrite hoisted function references (15:15)
+  16 |     return 42;
+  17 |   }
+  18 |
+```
+          
+      
\ No newline at end of file

filename: compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/error.todo-valid-functiondecl-hoisting.tsx
patch: @@ -0,0 +1,25 @@
+import {Stringify} from 'shared-runtime';
+/**
+ * Also see error.todo-functiondecl-hoisting.tsx which shows *invalid*
+ * compilation cases.
+ *
+ * This bailout specifically is a false positive for since this function's only
+ * reference-before-definition are within other functions which are not invoked.
+ */
+function Foo() {
+  'use memo';
+
+  function foo() {
+    return bar();
+  }
+  function bar() {
+    return 42;
+  }
+
+  return <Stringify fn1={foo} fn2={bar} shouldInvokeFns={true} />;
+}
+
+export const FIXTURE_ENTRYPOINT = {
+  fn: Foo,
+  params: [],
+};