[compiler] Rewrite effect dep arrays that use fire (#31811)

If an effect uses a dep array, also rewrite the dep array to use the
fire binding

--

compiler/packages/babel-plugin-react-compiler/src/Transform/TransformFire.ts

@@ -12,6 +12,7 @@ import {
   SourceLocation,
 } from '..';
 import {
+  ArrayExpression,
   CallExpression,
   Effect,
   Environment,
@@ -38,7 +39,6 @@ import {printSourceLocationLine} from '../HIR/PrintHIR';
 
 /*
  * TODO(jmbrown):
- *   - rewrite dep arrays
  *   - traverse object methods
  *   - method calls
  *   - React.useEffect calls
@@ -125,11 +125,58 @@ function replaceFireFunctions(fn: HIRFunction, context: Context): void {
               }
               rewriteInstrs.set(loadUseEffectInstrId, newInstrs);
             }
-            ensureNoRemainingCalleeCaptures(
-              lambda.loweredFunc.func,
-              context,
-              capturedCallees,
+          }
+          ensureNoRemainingCalleeCaptures(
+            lambda.loweredFunc.func,
+            context,
+            capturedCallees,
+          );
+
+          if (
+            value.args.length > 1 &&
+            value.args[1] != null &&
+            value.args[1].kind === 'Identifier'
+          ) {
+            const depArray = value.args[1];
+            const depArrayExpression = context.getArrayExpression(
+              depArray.identifier.id,
             );
+            if (depArrayExpression != null) {
+              for (const dependency of depArrayExpression.elements) {
+                if (dependency.kind === 'Identifier') {
+                  const loadOfDependency = context.getLoadLocalInstr(
+                    dependency.identifier.id,
+                  );
+                  if (loadOfDependency != null) {
+                    const replacedDepArrayItem = capturedCallees.get(
+                      loadOfDependency.place.identifier.id,
+                    );
+                    if (replacedDepArrayItem != null) {
+                      loadOfDependency.place =
+                        replacedDepArrayItem.fireFunctionBinding;
+                    }
+                  }
+                }
+              }
+            } else {
+              context.pushError({
+                loc: value.args[1].loc,
+                description:
+                  'You must use an array literal for an effect dependency array when that effect uses `fire()`',
+                severity: ErrorSeverity.Invariant,
+                reason: CANNOT_COMPILE_FIRE,
+                suggestions: null,
+              });
+            }
+          } else if (value.args.length > 1 && value.args[1].kind === 'Spread') {
+            context.pushError({
+              loc: value.args[1].place.loc,
+              description:
+                'You must use an array literal for an effect dependency array when that effect uses `fire()`',
+              severity: ErrorSeverity.Invariant,
+              reason: CANNOT_COMPILE_FIRE,
+              suggestions: null,
+            });
           }
         }
       } else if (
@@ -231,6 +278,8 @@ function replaceFireFunctions(fn: HIRFunction, context: Context): void {
         deleteInstrs.add(instr.id);
       } else if (value.kind === 'LoadGlobal') {
         context.addLoadGlobalInstrId(lvalue.identifier.id, instr.id);
+      } else if (value.kind === 'ArrayExpression') {
+        context.addArrayExpression(lvalue.identifier.id, value);
       }
     }
     block.instructions = rewriteInstructions(rewriteInstrs, block.instructions);
@@ -561,6 +610,12 @@ class Context {
     this.#env = env;
   }
 
+  /*
+   * We keep track of array expressions so we can rewrite dependency arrays passed to useEffect
+   * to use the fire functions
+   */
+  #arrayExpressions = new Map<IdentifierId, ArrayExpression>();
+
   pushError(error: CompilerErrorDetailOptions): void {
     this.#errors.push(error);
   }
@@ -655,6 +710,14 @@ class Context {
     return this.#loadGlobalInstructionIds.get(id);
   }
 
+  addArrayExpression(id: IdentifierId, array: ArrayExpression): void {
+    this.#arrayExpressions.set(id, array);
+  }
+
+  getArrayExpression(id: IdentifierId): ArrayExpression | undefined {
+    return this.#arrayExpressions.get(id);
+  }
+
   hasErrors(): boolean {
     return this.#errors.hasErrors();
   }

compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/transform-fire/error.invalid-rewrite-deps-no-array-literal.expect.md

@@ -0,0 +1,37 @@
+
+## Input
+
+```javascript
+// @enableFire
+import {fire} from 'react';
+
+function Component(props) {
+  const foo = props => {
+    console.log(props);
+  };
+
+  const deps = [foo, props];
+
+  useEffect(() => {
+    fire(foo(props));
+  }, deps);
+
+  return null;
+}
+
+```
+
+
+## Error
+
+```
+  11 |   useEffect(() => {
+  12 |     fire(foo(props));
+> 13 |   }, deps);
+     |      ^^^^ Invariant: Cannot compile `fire`. You must use an array literal for an effect dependency array when that effect uses `fire()` (13:13)
+  14 |
+  15 |   return null;
+  16 | }
+```
+          
+      

compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/transform-fire/error.invalid-rewrite-deps-no-array-literal.js

@@ -0,0 +1,16 @@
+// @enableFire
+import {fire} from 'react';
+
+function Component(props) {
+  const foo = props => {
+    console.log(props);
+  };
+
+  const deps = [foo, props];
+
+  useEffect(() => {
+    fire(foo(props));
+  }, deps);
+
+  return null;
+}

compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/transform-fire/error.invalid-rewrite-deps-spread.expect.md

@@ -0,0 +1,37 @@
+
+## Input
+
+```javascript
+// @enableFire
+import {fire} from 'react';
+
+function Component(props) {
+  const foo = props => {
+    console.log(props);
+  };
+
+  const deps = [foo, props];
+
+  useEffect(() => {
+    fire(foo(props));
+  }, ...deps);
+
+  return null;
+}
+
+```
+
+
+## Error
+
+```
+  11 |   useEffect(() => {
+  12 |     fire(foo(props));
+> 13 |   }, ...deps);
+     |         ^^^^ Invariant: Cannot compile `fire`. You must use an array literal for an effect dependency array when that effect uses `fire()` (13:13)
+  14 |
+  15 |   return null;
+  16 | }
+```
+          
+      

compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/transform-fire/error.invalid-rewrite-deps-spread.js

@@ -0,0 +1,19 @@
+// @enableFire
+import {fire} from 'react';
+
+function Component(props) {
+  const foo = props => {
+    console.log(props);
+  };
+
+  const deps = [foo, props];
+
+  useEffect(
+    () => {
+      fire(foo(props));
+    },
+    ...deps
+  );
+
+  return null;
+}

compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/transform-fire/fire-and-autodeps.expect.md

@@ -0,0 +1,60 @@
+
+## Input
+
+```javascript
+// @enableFire @inferEffectDependencies
+import {fire, useEffect} from 'react';
+
+function Component(props) {
+  const foo = arg => {
+    console.log(arg, props.bar);
+  };
+  useEffect(() => {
+    fire(foo(props));
+  });
+
+  return null;
+}
+
+```
+
+## Code
+
+```javascript
+import { useFire } from "react/compiler-runtime";
+import { c as _c } from "react/compiler-runtime"; // @enableFire @inferEffectDependencies
+import { fire, useEffect } from "react";
+
+function Component(props) {
+  const $ = _c(5);
+  let t0;
+  if ($[0] !== props.bar) {
+    t0 = (arg) => {
+      console.log(arg, props.bar);
+    };
+    $[0] = props.bar;
+    $[1] = t0;
+  } else {
+    t0 = $[1];
+  }
+  const foo = t0;
+  const t1 = useFire(foo);
+  let t2;
+  if ($[2] !== props || $[3] !== t1) {
+    t2 = () => {
+      t1(props);
+    };
+    $[2] = props;
+    $[3] = t1;
+    $[4] = t2;
+  } else {
+    t2 = $[4];
+  }
+  useEffect(t2, [t1, props]);
+  return null;
+}
+
+```
+      
+### Eval output
+(kind: exception) Fixture not implemented

compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/transform-fire/fire-and-autodeps.js

@@ -0,0 +1,13 @@
+// @enableFire @inferEffectDependencies
+import {fire, useEffect} from 'react';
+
+function Component(props) {
+  const foo = arg => {
+    console.log(arg, props.bar);
+  };
+  useEffect(() => {
+    fire(foo(props));
+  });
+
+  return null;
+}

compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/transform-fire/rewrite-deps.expect.md

@@ -0,0 +1,57 @@
+
+## Input
+
+```javascript
+// @enableFire
+import {fire} from 'react';
+
+function Component(props) {
+  const foo = props => {
+    console.log(props);
+  };
+  useEffect(() => {
+    fire(foo(props));
+  }, [foo, props]);
+
+  return null;
+}
+
+```
+
+## Code
+
+```javascript
+import { useFire } from "react/compiler-runtime";
+import { c as _c } from "react/compiler-runtime"; // @enableFire
+import { fire } from "react";
+
+function Component(props) {
+  const $ = _c(4);
+  const foo = _temp;
+  const t0 = useFire(foo);
+  let t1;
+  let t2;
+  if ($[0] !== props || $[1] !== t0) {
+    t1 = () => {
+      t0(props);
+    };
+    t2 = [t0, props];
+    $[0] = props;
+    $[1] = t0;
+    $[2] = t1;
+    $[3] = t2;
+  } else {
+    t1 = $[2];
+    t2 = $[3];
+  }
+  useEffect(t1, t2);
+  return null;
+}
+function _temp(props_0) {
+  console.log(props_0);
+}
+
+```
+      
+### Eval output
+(kind: exception) Fixture not implemented

compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/transform-fire/rewrite-deps.js

@@ -0,0 +1,13 @@
+// @enableFire
+import {fire} from 'react';
+
+function Component(props) {
+  const foo = props => {
+    console.log(props);
+  };
+  useEffect(() => {
+    fire(foo(props));
+  }, [foo, props]);
+
+  return null;
+}