[compiler][bugfix] Fix hoisting of let declarations

compiler/packages/babel-plugin-react-compiler/src/ReactiveScopes/PruneHoistedContexts.ts

@@ -5,6 +5,7 @@
  * LICENSE file in the root directory of this source tree.
  */
 
+import {CompilerError} from '..';
 import {
   DeclarationId,
   InstructionKind,
@@ -27,14 +28,28 @@ export function pruneHoistedContexts(fn: ReactiveFunction): void {
   visitReactiveFunction(fn, new Visitor(), hoistedIdentifiers);
 }
 
-type HoistedIdentifiers = Map<DeclarationId, InstructionKind>;
+const REWRITTEN_HOISTED_CONST: unique symbol = Symbol(
+  'REWRITTEN_HOISTED_CONST',
+);
+const REWRITTEN_HOISTED_LET: unique symbol = Symbol('REWRITTEN_HOISTED_LET');
+
+type HoistedIdentifiers = Map<
+  DeclarationId,
+  | InstructionKind
+  | typeof REWRITTEN_HOISTED_CONST
+  | typeof REWRITTEN_HOISTED_LET
+>;
 
 class Visitor extends ReactiveFunctionTransform<HoistedIdentifiers> {
   override transformInstruction(
     instruction: ReactiveInstruction,
     state: HoistedIdentifiers,
   ): Transformed<ReactiveStatement> {
     this.visitInstruction(instruction, state);
+
+    /**
+     * Remove hoisted declarations to preserve TDZ
+     */
     if (
       instruction.value.kind === 'DeclareContext' &&
       instruction.value.lvalue.kind === 'HoistedConst'
@@ -68,31 +83,75 @@ class Visitor extends ReactiveFunctionTransform<HoistedIdentifiers> {
       return {kind: 'remove'};
     }
 
-    if (
-      instruction.value.kind === 'StoreContext' &&
-      state.has(instruction.value.lvalue.place.identifier.declarationId)
-    ) {
+    if (instruction.value.kind === 'StoreContext') {
       const kind = state.get(
         instruction.value.lvalue.place.identifier.declarationId,
-      )!;
-      return {
-        kind: 'replace',
-        value: {
-          kind: 'instruction',
-          instruction: {
-            ...instruction,
+      );
+      if (kind != null) {
+        CompilerError.invariant(kind !== REWRITTEN_HOISTED_CONST, {
+          reason: 'Expected exactly one store to a hoisted const variable',
+          loc: instruction.loc,
+        });
+        if (
+          kind === InstructionKind.Const ||
+          kind === InstructionKind.Function
+        ) {
+          state.set(
+            instruction.value.lvalue.place.identifier.declarationId,
+            REWRITTEN_HOISTED_CONST,
+          );
+          return {
+            kind: 'replace',
             value: {
-              ...instruction.value,
-              lvalue: {
-                ...instruction.value.lvalue,
-                kind,
+              kind: 'instruction',
+              instruction: {
+                ...instruction,
+                value: {
+                  ...instruction.value,
+                  lvalue: {
+                    ...instruction.value.lvalue,
+                    kind,
+                  },
+                  type: null,
+                  kind: 'StoreLocal',
+                },
               },
-              type: null,
-              kind: 'StoreLocal',
             },
-          },
-        },
-      };
+          };
+        } else if (kind !== REWRITTEN_HOISTED_LET) {
+          /**
+           * Context variables declared with let may have reassignments. Only
+           * insert a `DeclareContext` for the first encountered `StoreContext`
+           * instruction.
+           */
+          state.set(
+            instruction.value.lvalue.place.identifier.declarationId,
+            REWRITTEN_HOISTED_LET,
+          );
+          return {
+            kind: 'replace-many',
+            value: [
+              {
+                kind: 'instruction',
+                instruction: {
+                  id: instruction.id,
+                  lvalue: null,
+                  value: {
+                    kind: 'DeclareContext',
+                    lvalue: {
+                      kind: InstructionKind.Let,
+                      place: {...instruction.value.lvalue.place},
+                    },
+                    loc: instruction.value.loc,
+                  },
+                  loc: instruction.loc,
+                },
+              },
+              {kind: 'instruction', instruction},
+            ],
+          };
+        }
+      }
     }
 
     return {kind: 'keep'};

compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/hoisting-invalid-tdz-let.expect.md

@@ -0,0 +1,59 @@
+
+## Input
+
+```javascript
+function Foo() {
+  const getX = () => x;
+  console.log(getX());
+
+  let x = 4;
+  x += 5;
+
+  return <Stringify getX={getX} shouldInvokeFns={true} />;
+}
+
+export const FIXTURE_ENTRYPOINT = {
+  fn: Foo,
+  params: [],
+};
+
+```
+
+## Code
+
+```javascript
+import { c as _c } from "react/compiler-runtime";
+function Foo() {
+  const $ = _c(2);
+  let getX;
+  if ($[0] === Symbol.for("react.memo_cache_sentinel")) {
+    getX = () => x;
+    console.log(getX());
+
+    let x;
+    x = 4;
+    x = x + 5;
+    $[0] = getX;
+  } else {
+    getX = $[0];
+  }
+  x;
+  let t0;
+  if ($[1] === Symbol.for("react.memo_cache_sentinel")) {
+    t0 = <Stringify getX={getX} shouldInvokeFns={true} />;
+    $[1] = t0;
+  } else {
+    t0 = $[1];
+  }
+  return t0;
+}
+
+export const FIXTURE_ENTRYPOINT = {
+  fn: Foo,
+  params: [],
+};
+
+```
+
+### Eval output
+(kind: exception) Cannot access 'x' before initialization

compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/hoisting-invalid-tdz-let.js

@@ -0,0 +1,14 @@
+function Foo() {
+  const getX = () => x;
+  console.log(getX());
+
+  let x = 4;
+  x += 5;
+
+  return <Stringify getX={getX} shouldInvokeFns={true} />;
+}
+
+export const FIXTURE_ENTRYPOINT = {
+  fn: Foo,
+  params: [],
+};

compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/hoisting-nested-let-declaration-2.expect.md

@@ -36,7 +36,8 @@ function hoisting(cond) {
         items.push(bar());
       };
 
-      let bar = _temp;
+      let bar;
+      bar = _temp;
       foo();
     }
     $[0] = cond;

compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/hoisting-nested-let-declaration.expect.md

@@ -41,9 +41,11 @@ function hoisting() {
       return result;
     };
 
-    let foo = () => bar + baz;
+    let foo;
+    foo = () => bar + baz;
 
-    let bar = 3;
+    let bar;
+    bar = 3;
     const baz = 2;
     t0 = qux();
     $[0] = t0;

compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/hoisting-reassigned-let-declaration.expect.md

@@ -0,0 +1,67 @@
+
+## Input
+
+```javascript
+import {CONST_NUMBER0, CONST_NUMBER1, Stringify} from 'shared-runtime';
+
+function useHook({cond}) {
+  'use memo';
+  const getX = () => x;
+
+  let x = CONST_NUMBER0;
+  if (cond) {
+    x += CONST_NUMBER1;
+  }
+  return <Stringify getX={getX} shouldInvokeFns={true} />;
+}
+
+export const FIXTURE_ENTRYPOINT = {
+  fn: useHook,
+  params: [{cond: true}],
+  sequentialRenders: [{cond: true}, {cond: true}, {cond: false}],
+};
+
+```
+
+## Code
+
+```javascript
+import { c as _c } from "react/compiler-runtime";
+import { CONST_NUMBER0, CONST_NUMBER1, Stringify } from "shared-runtime";
+
+function useHook(t0) {
+  "use memo";
+  const $ = _c(2);
+  const { cond } = t0;
+  let t1;
+  if ($[0] !== cond) {
+    const getX = () => x;
+
+    let x;
+    x = CONST_NUMBER0;
+    if (cond) {
+      x = x + CONST_NUMBER1;
+      x;
+    }
+
+    t1 = <Stringify getX={getX} shouldInvokeFns={true} />;
+    $[0] = cond;
+    $[1] = t1;
+  } else {
+    t1 = $[1];
+  }
+  return t1;
+}
+
+export const FIXTURE_ENTRYPOINT = {
+  fn: useHook,
+  params: [{ cond: true }],
+  sequentialRenders: [{ cond: true }, { cond: true }, { cond: false }],
+};
+
+```
+
+### Eval output
+(kind: ok) <div>{"getX":{"kind":"Function","result":1},"shouldInvokeFns":true}</div>
+<div>{"getX":{"kind":"Function","result":1},"shouldInvokeFns":true}</div>
+<div>{"getX":{"kind":"Function","result":0},"shouldInvokeFns":true}</div>

compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/hoisting-reassigned-let-declaration.js

@@ -0,0 +1,18 @@
+import {CONST_NUMBER0, CONST_NUMBER1, Stringify} from 'shared-runtime';
+
+function useHook({cond}) {
+  'use memo';
+  const getX = () => x;
+
+  let x = CONST_NUMBER0;
+  if (cond) {
+    x += CONST_NUMBER1;
+  }
+  return <Stringify getX={getX} shouldInvokeFns={true} />;
+}
+
+export const FIXTURE_ENTRYPOINT = {
+  fn: useHook,
+  params: [{cond: true}],
+  sequentialRenders: [{cond: true}, {cond: true}, {cond: false}],
+};

compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/hoisting-reassigned-twice-let-declaration.expect.md

@@ -0,0 +1,69 @@
+
+## Input
+
+```javascript
+import {CONST_NUMBER0, CONST_NUMBER1, Stringify} from 'shared-runtime';
+
+function useHook({cond}) {
+  'use memo';
+  const getX = () => x;
+
+  let x = CONST_NUMBER0;
+  if (cond) {
+    x += CONST_NUMBER1;
+    x = Math.min(x, 100);
+  }
+  return <Stringify getX={getX} shouldInvokeFns={true} />;
+}
+
+export const FIXTURE_ENTRYPOINT = {
+  fn: useHook,
+  params: [{cond: true}],
+  sequentialRenders: [{cond: true}, {cond: true}, {cond: false}],
+};
+
+```
+
+## Code
+
+```javascript
+import { c as _c } from "react/compiler-runtime";
+import { CONST_NUMBER0, CONST_NUMBER1, Stringify } from "shared-runtime";
+
+function useHook(t0) {
+  "use memo";
+  const $ = _c(2);
+  const { cond } = t0;
+  let t1;
+  if ($[0] !== cond) {
+    const getX = () => x;
+
+    let x;
+    x = CONST_NUMBER0;
+    if (cond) {
+      x = x + CONST_NUMBER1;
+      x;
+      x = Math.min(x, 100);
+    }
+
+    t1 = <Stringify getX={getX} shouldInvokeFns={true} />;
+    $[0] = cond;
+    $[1] = t1;
+  } else {
+    t1 = $[1];
+  }
+  return t1;
+}
+
+export const FIXTURE_ENTRYPOINT = {
+  fn: useHook,
+  params: [{ cond: true }],
+  sequentialRenders: [{ cond: true }, { cond: true }, { cond: false }],
+};
+
+```
+
+### Eval output
+(kind: ok) <div>{"getX":{"kind":"Function","result":1},"shouldInvokeFns":true}</div>
+<div>{"getX":{"kind":"Function","result":1},"shouldInvokeFns":true}</div>
+<div>{"getX":{"kind":"Function","result":0},"shouldInvokeFns":true}</div>