[compiler] Implement support for non-declaration for initializers

ghstack-source-id: 894914a
Pull Request resolved: facebook/react#31712

Author: mvitousek

compiler/packages/babel-plugin-react-compiler/src/HIR/BuildHIR.ts

@@ -519,27 +519,23 @@ function lowerStatement(
 
       const initBlock = builder.enter('loop', _blockId => {
         const init = stmt.get('init');
-        if (!init.isVariableDeclaration()) {
-          builder.errors.push({
-            reason:
-              '(BuildHIR::lowerStatement) Handle non-variable initialization in ForStatement',
-            severity: ErrorSeverity.Todo,
-            loc: stmt.node.loc ?? null,
-            suggestions: null,
+        if (init.isVariableDeclaration()) {
+          lowerStatement(builder, init);
+        } else if (init.isExpression()) {
+          lowerExpressionToTemporary(builder, init);
+        } else {
+          lowerValueToTemporary(builder, {
+            kind: 'Primitive',
+            loc: stmt.node.loc ?? GeneratedSource,
+            value: undefined,
           });
-          return {
-            kind: 'unsupported',
-            id: makeInstructionId(0),
-            loc: init.node?.loc ?? GeneratedSource,
-          };
         }
-        lowerStatement(builder, init);
         return {
           kind: 'goto',
           block: testBlock.id,
           variant: GotoVariant.Break,
           id: makeInstructionId(0),
-          loc: init.node.loc ?? GeneratedSource,
+          loc: init.node?.loc ?? GeneratedSource,
         };
       });
 

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

@@ -18,6 +18,7 @@ import {
 } from '../HIR';
 import {
   HIRFunction,
+  InstructionKind,
   ReactiveBreakTerminal,
   ReactiveContinueTerminal,
   ReactiveFunction,
@@ -364,26 +365,25 @@ class Driver {
         const init = this.visitValueBlock(terminal.init, terminal.loc);
         const initBlock = this.cx.ir.blocks.get(init.block)!;
         let initValue = init.value;
-        if (initValue.kind === 'SequenceExpression') {
+        const initBlockHasDeclarations = initBlock.instructions.some(
+          instr =>
+            instr.value.kind === 'DeclareContext' ||
+            instr.value.kind === 'DeclareLocal' ||
+            ((instr.value.kind === 'StoreContext' ||
+              instr.value.kind === 'StoreLocal') &&
+              instr.value.lvalue.kind !== InstructionKind.Reassign),
+        );
+        if (
+          initValue.kind === 'SequenceExpression' &&
+          initBlockHasDeclarations
+        ) {
           const last = initBlock.instructions.at(-1)!;
           initValue.instructions.push(last);
           initValue.value = {
             kind: 'Primitive',
             value: undefined,
             loc: terminal.loc,
           };
-        } else {
-          initValue = {
-            kind: 'SequenceExpression',
-            instructions: [initBlock.instructions.at(-1)!],
-            id: terminal.id,
-            loc: terminal.loc,
-            value: {
-              kind: 'Primitive',
-              value: undefined,
-              loc: terminal.loc,
-            },
-          };
         }
 
         const testValue = this.visitValueBlock(

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

@@ -1354,6 +1354,8 @@ function codegenForInit(
   init: ReactiveValue,
 ): t.Expression | t.VariableDeclaration | null {
   if (init.kind === 'SequenceExpression') {
+    // We may end up emitti
+    const temp = cx.temp;
     const body = codegenBlock(
       cx,
       init.instructions.map(instruction => ({
@@ -1363,7 +1365,7 @@ function codegenForInit(
     ).body;
     const declarators: Array<t.VariableDeclarator> = [];
     let kind: 'let' | 'const' = 'const';
-    body.forEach(instr => {
+    for (const instr of body) {
       let top: undefined | t.VariableDeclarator = undefined;
       if (
         instr.type === 'ExpressionStatement' &&
@@ -1375,30 +1377,29 @@ function codegenForInit(
         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,
-          },
-        );
+      } else if (
+        instr.type === 'VariableDeclaration' &&
+        (instr.kind === 'let' || instr.kind === 'const')
+      ) {
         if (instr.kind === 'let') {
           kind = 'let';
         }
         declarators.push(...instr.declarations);
+      } else {
+        /*
+         * We found instructions in the initializer that don't correspond to a variable declarator:
+         * emit as an expression instead
+         */
+        cx.temp = temp;
+        return codegenInstructionValueToExpression(cx, init);
       }
-    });
-    CompilerError.invariant(declarators.length > 0, {
-      reason: 'Expected a variable declaration',
-      loc: init.loc,
-      description: null,
-      suggestions: null,
-    });
-    return t.variableDeclaration(kind, declarators);
+    }
+    if (declarators.length > 0) {
+      return t.variableDeclaration(kind, declarators);
+    } else {
+      cx.temp = temp;
+      return codegenInstructionValueToExpression(cx, init);
+    }
   } else {
     return codegenInstructionValueToExpression(cx, init);
   }

compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/error.todo-kitchensink.expect.md

@@ -86,12 +86,6 @@ let moduleLocal = false;
 
 Todo: (BuildHIR::lowerStatement) Handle ClassDeclaration statements (5:10)
 
-Todo: (BuildHIR::lowerStatement) Handle non-variable initialization in ForStatement (20:22)
-
-Todo: (BuildHIR::lowerStatement) Handle non-variable initialization in ForStatement (23:25)
-
-Todo: (BuildHIR::lowerStatement) Handle non-variable initialization in ForStatement (26:28)
-
 Todo: (BuildHIR::lowerStatement) Handle empty test in ForStatement (26:28)
 
 Todo: (BuildHIR::lowerExpression) Handle tagged template with interpolations (30:32)

compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/for-inits.expect.md

@@ -0,0 +1,43 @@
+
+## Input
+
+```javascript
+function Foo() {
+  let i = 0;
+  for (42; i < 1; i += 1) {}
+  for (bar(); i < 1; i += 1) {}
+  for (; i < 1; i += 1) {}
+  for (i = 0; i < 1; i += 1) {}
+}
+
+function bar() {}
+
+export const FIXTURE_ENTRYPOINT = {
+  fn: Foo,
+  params: [],
+};
+
+```
+
+## Code
+
+```javascript
+function Foo() {
+  let i = 0;
+  for (42; i < 1; i = i + 1, i) {}
+  for (bar(); i < 1; i = i + 1, i) {}
+  for (undefined; i < 1; i = i + 1, i) {}
+  for (i = 0; i < 1; i = i + 1, i) {}
+}
+
+function bar() {}
+
+export const FIXTURE_ENTRYPOINT = {
+  fn: Foo,
+  params: [],
+};
+
+```
+      
+### Eval output
+(kind: ok) 

compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/for-inits.js

@@ -0,0 +1,14 @@
+function Foo() {
+  let i = 0;
+  for (42; i < 1; i += 1) {}
+  for (bar(); i < 1; i += 1) {}
+  for (; i < 1; i += 1) {}
+  for (i = 0; i < 1; i += 1) {}
+}
+
+function bar() {}
+
+export const FIXTURE_ENTRYPOINT = {
+  fn: Foo,
+  params: [],
+};