[compiler] Support for context variable loop iterators

Summary:
Fixing a compiler todo

ghstack-source-id: c4d9226b1745d003dc9945df1ac5c5a01712f909
Pull Request resolved: facebook#31709

Author: mvitousek

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

@@ -1354,20 +1354,6 @@ function codegenForInit(
   init: ReactiveValue,
 ): t.Expression | t.VariableDeclaration | null {
   if (init.kind === 'SequenceExpression') {
-    for (const instr of init.instructions) {
-      if (instr.value.kind === 'DeclareContext') {
-        CompilerError.throwTodo({
-          reason: `Support for loops where the index variable is a context variable`,
-          loc: instr.loc,
-          description:
-            instr.value.lvalue.place.identifier.name != null
-              ? `\`${instr.value.lvalue.place.identifier.name.value}\` is a context variable`
-              : null,
-          suggestions: null,
-        });
-      }
-    }
-
     const body = codegenBlock(
       cx,
       init.instructions.map(instruction => ({
@@ -1378,20 +1364,33 @@ function codegenForInit(
     const declarators: Array<t.VariableDeclarator> = [];
     let kind: 'let' | 'const' = 'const';
     body.forEach(instr => {
-      CompilerError.invariant(
-        instr.type === 'VariableDeclaration' &&
-          (instr.kind === 'let' || instr.kind === 'const'),
-        {
-          reason: 'Expected a variable declaration',
-          loc: init.loc,
-          description: `Got ${instr.type}`,
-          suggestions: null,
-        },
-      );
-      if (instr.kind === 'let') {
-        kind = 'let';
+      let top: undefined | t.VariableDeclarator = undefined;
+      if (
+        instr.type === 'ExpressionStatement' &&
+        instr.expression.type === 'AssignmentExpression' &&
+        instr.expression.operator === '=' &&
+        instr.expression.left.type === 'Identifier' &&
+        (top = declarators.at(-1))?.id.type === 'Identifier' &&
+        top?.id.name === instr.expression.left.name &&
+        top?.init == null
+      ) {
+        top.init = instr.expression.right;
+      } else {
+        CompilerError.invariant(
+          instr.type === 'VariableDeclaration' &&
+            (instr.kind === 'let' || instr.kind === 'const'),
+          {
+            reason: 'Expected a variable declaration',
+            loc: init.loc,
+            description: `Got ${instr.type}`,
+            suggestions: null,
+          },
+        );
+        if (instr.kind === 'let') {
+          kind = 'let';
+        }
+        declarators.push(...instr.declarations);
       }
-      declarators.push(...instr.declarations);
     });
     CompilerError.invariant(declarators.length > 0, {
       reason: 'Expected a variable declaration',

compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/error.todo-for-loop-with-context-variable-iterator.expect.md

@@ -0,0 +1,78 @@
+
+## Input
+
+```javascript
+function Component() {
+  const data = useData();
+  const items = [];
+  // NOTE: `i` is a context variable because it's reassigned and also referenced
+  // within a closure, the `onClick` handler of each item
+  for (let i = MIN; i <= MAX; i += INCREMENT) {
+    items.push(<div key={i} onClick={() => data.set(i)} />);
+  }
+  return <>{items}</>;
+}
+
+const MIN = 0;
+const MAX = 3;
+const INCREMENT = 1;
+
+function useData() {
+  return new Map();
+}
+
+export const FIXTURE_ENTRYPOINT = {
+  params: [],
+  fn: Component,
+};
+
+```
+
+## Code
+
+```javascript
+import { c as _c } from "react/compiler-runtime";
+function Component() {
+  const $ = _c(2);
+  const data = useData();
+  let t0;
+  if ($[0] !== data) {
+    const items = [];
+    for (let i = MIN; i <= MAX; i = i + INCREMENT, i) {
+      items.push(<div key={i} onClick={() => data.set(i)} />);
+    }
+
+    t0 = <>{items}</>;
+    $[0] = data;
+    $[1] = t0;
+  } else {
+    t0 = $[1];
+  }
+  return t0;
+}
+
+const MIN = 0;
+const MAX = 3;
+const INCREMENT = 1;
+
+function useData() {
+  const $ = _c(1);
+  let t0;
+  if ($[0] === Symbol.for("react.memo_cache_sentinel")) {
+    t0 = new Map();
+    $[0] = t0;
+  } else {
+    t0 = $[0];
+  }
+  return t0;
+}
+
+export const FIXTURE_ENTRYPOINT = {
+  params: [],
+  fn: Component,
+};
+
+```
+      
+### Eval output
+(kind: ok) <div></div><div></div><div></div><div></div>

compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/for-loop-with-context-variable-iterator.expect.md

@@ -4,7 +4,20 @@ function Component() {
   // NOTE: `i` is a context variable because it's reassigned and also referenced
   // within a closure, the `onClick` handler of each item
   for (let i = MIN; i <= MAX; i += INCREMENT) {
-    items.push(<Stringify key={i} onClick={() => data.set(i)} />);
+    items.push(<div key={i} onClick={() => data.set(i)} />);
   }
-  return items;
+  return <>{items}</>;
 }
+
+const MIN = 0;
+const MAX = 3;
+const INCREMENT = 1;
+
+function useData() {
+  return new Map();
+}
+
+export const FIXTURE_ENTRYPOINT = {
+  params: [],
+  fn: Component,
+};