From 76d603a72aa14e93e73a68e0f96024a1841902f3 Mon Sep 17 00:00:00 2001
From: Mike Vitousek <mmvitousek@gmail.com>
Date: Mon, 9 Dec 2024 10:48:43 -0800
Subject: [PATCH] [compiler] Support for non-declatation for in/of iterators

ghstack-source-id: a28801e022561029e2f46c3dcb858bd4a81dea6a
Pull Request resolved: https://github.com/facebook/react/pull/31710
---
 .../src/HIR/BuildHIR.ts                       | 58 +++++++++++--------
 .../compiler/error.todo-kitchensink.expect.md |  6 --
 2 files changed, 35 insertions(+), 29 deletions(-)

diff --git a/compiler/packages/babel-plugin-react-compiler/src/HIR/BuildHIR.ts b/compiler/packages/babel-plugin-react-compiler/src/HIR/BuildHIR.ts
index ecc22365dd03c..4d9ce6becc17a 100644
--- a/compiler/packages/babel-plugin-react-compiler/src/HIR/BuildHIR.ts
+++ b/compiler/packages/babel-plugin-react-compiler/src/HIR/BuildHIR.ts
@@ -1078,6 +1078,12 @@ function lowerStatement(
       const left = stmt.get('left');
       const leftLoc = left.node.loc ?? GeneratedSource;
       let test: Place;
+      const advanceIterator = lowerValueToTemporary(builder, {
+        kind: 'IteratorNext',
+        loc: leftLoc,
+        iterator: {...iterator},
+        collection: {...value},
+      });
       if (left.isVariableDeclaration()) {
         const declarations = left.get('declarations');
         CompilerError.invariant(declarations.length === 1, {
@@ -1087,12 +1093,6 @@ function lowerStatement(
           suggestions: null,
         });
         const id = declarations[0].get('id');
-        const advanceIterator = lowerValueToTemporary(builder, {
-          kind: 'IteratorNext',
-          loc: leftLoc,
-          iterator: {...iterator},
-          collection: {...value},
-        });
         const assign = lowerAssignment(
           builder,
           leftLoc,
@@ -1103,13 +1103,19 @@ function lowerStatement(
         );
         test = lowerValueToTemporary(builder, assign);
       } else {
-        builder.errors.push({
-          reason: `(BuildHIR::lowerStatement) Handle ${left.type} inits in ForOfStatement`,
-          severity: ErrorSeverity.Todo,
-          loc: left.node.loc ?? null,
-          suggestions: null,
+        CompilerError.invariant(left.isLVal(), {
+          loc: leftLoc,
+          reason: 'Expected ForOf init to be a variable declaration or lval',
         });
-        return;
+        const assign = lowerAssignment(
+          builder,
+          leftLoc,
+          InstructionKind.Reassign,
+          left,
+          advanceIterator,
+          'Assignment',
+        );
+        test = lowerValueToTemporary(builder, assign);
       }
       builder.terminateWithContinuation(
         {
@@ -1166,6 +1172,11 @@ function lowerStatement(
       const left = stmt.get('left');
       const leftLoc = left.node.loc ?? GeneratedSource;
       let test: Place;
+      const nextPropertyTemp = lowerValueToTemporary(builder, {
+        kind: 'NextPropertyOf',
+        loc: leftLoc,
+        value,
+      });
       if (left.isVariableDeclaration()) {
         const declarations = left.get('declarations');
         CompilerError.invariant(declarations.length === 1, {
@@ -1175,11 +1186,6 @@ function lowerStatement(
           suggestions: null,
         });
         const id = declarations[0].get('id');
-        const nextPropertyTemp = lowerValueToTemporary(builder, {
-          kind: 'NextPropertyOf',
-          loc: leftLoc,
-          value,
-        });
         const assign = lowerAssignment(
           builder,
           leftLoc,
@@ -1190,13 +1196,19 @@ function lowerStatement(
         );
         test = lowerValueToTemporary(builder, assign);
       } else {
-        builder.errors.push({
-          reason: `(BuildHIR::lowerStatement) Handle ${left.type} inits in ForInStatement`,
-          severity: ErrorSeverity.Todo,
-          loc: left.node.loc ?? null,
-          suggestions: null,
+        CompilerError.invariant(left.isLVal(), {
+          loc: leftLoc,
+          reason: 'Expected ForIn init to be a variable declaration or lval',
         });
-        return;
+        const assign = lowerAssignment(
+          builder,
+          leftLoc,
+          InstructionKind.Reassign,
+          left,
+          nextPropertyTemp,
+          'Assignment',
+        );
+        test = lowerValueToTemporary(builder, assign);
       }
       builder.terminateWithContinuation(
         {
diff --git a/compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/error.todo-kitchensink.expect.md b/compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/error.todo-kitchensink.expect.md
index d8360455bb303..776b16091327d 100644
--- a/compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/error.todo-kitchensink.expect.md
+++ b/compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/error.todo-kitchensink.expect.md
@@ -98,12 +98,6 @@ Todo: (BuildHIR::lowerExpression) Handle tagged template with interpolations (30
 
 Todo: (BuildHIR::lowerExpression) Handle tagged template where cooked value is different from raw value (34:34)
 
-Todo: (BuildHIR::lowerStatement) Handle Identifier inits in ForOfStatement (36:36)
-
-Todo: (BuildHIR::lowerStatement) Handle ArrayPattern inits in ForOfStatement (38:38)
-
-Todo: (BuildHIR::lowerStatement) Handle ObjectPattern inits in ForOfStatement (40:40)
-
 Todo: (BuildHIR::node.lowerReorderableExpression) Expression type `MemberExpression` cannot be safely reordered (57:57)
 
 Todo: (BuildHIR::node.lowerReorderableExpression) Expression type `BinaryExpression` cannot be safely reordered (53:53)