[compile] Error on fire outside of effects and ensure correct compilation, correct import

compiler/packages/babel-plugin-react-compiler/src/Entrypoint/Program.ts

@@ -567,7 +567,10 @@ export function compileProgram(
 
     const hasFireRewrite = compiledFns.some(c => c.compiledFn.hasFireRewrite);
     if (environment.enableFire && hasFireRewrite) {
-      externalFunctions.push({source: 'react', importSpecifierName: 'useFire'});
+      externalFunctions.push({
+        source: getReactCompilerRuntimeModule(pass.opts),
+        importSpecifierName: 'useFire',
+      });
     }
   } catch (err) {
     handleError(err, pass, null);

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

@@ -897,6 +897,14 @@ export function printSourceLocation(loc: SourceLocation): string {
   }
 }
 
+export function printSourceLocationLine(loc: SourceLocation): string {
+  if (typeof loc === 'symbol') {
+    return 'generated';
+  } else {
+    return `${loc.start.line}:${loc.end.line}`;
+  }
+}
+
 export function printAliases(aliases: DisjointSet<Identifier>): string {
   const aliasSets = aliases.buildSets();
 

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

@@ -5,7 +5,12 @@
  * LICENSE file in the root directory of this source tree.
  */
 
-import {CompilerError, CompilerErrorDetailOptions, ErrorSeverity} from '..';
+import {
+  CompilerError,
+  CompilerErrorDetailOptions,
+  ErrorSeverity,
+  SourceLocation,
+} from '..';
 import {
   CallExpression,
   Effect,
@@ -28,14 +33,11 @@ import {
 import {createTemporaryPlace, markInstructionIds} from '../HIR/HIRBuilder';
 import {getOrInsertWith} from '../Utils/utils';
 import {BuiltInFireId, DefaultNonmutatingHook} from '../HIR/ObjectShape';
+import {eachInstructionOperand} from '../HIR/visitors';
+import {printSourceLocationLine} from '../HIR/PrintHIR';
 
 /*
  * TODO(jmbrown):
- *   In this stack:
- *     - Assert no lingering fire calls
- *     - Ensure a fired function is not called regularly elsewhere in the same effect
- *
- *   Future:
  *   - rewrite dep arrays
  *   - traverse object methods
  *   - method calls
@@ -47,6 +49,9 @@ const CANNOT_COMPILE_FIRE = 'Cannot compile `fire`';
 export function transformFire(fn: HIRFunction): void {
   const context = new Context(fn.env);
   replaceFireFunctions(fn, context);
+  if (!context.hasErrors()) {
+    ensureNoMoreFireUses(fn, context);
+  }
   context.throwIfErrorsFound();
 }
 
@@ -120,6 +125,11 @@ function replaceFireFunctions(fn: HIRFunction, context: Context): void {
               }
               rewriteInstrs.set(loadUseEffectInstrId, newInstrs);
             }
+            ensureNoRemainingCalleeCaptures(
+              lambda.loweredFunc.func,
+              context,
+              capturedCallees,
+            );
           }
         }
       } else if (
@@ -159,7 +169,10 @@ function replaceFireFunctions(fn: HIRFunction, context: Context): void {
             }
 
             const fireFunctionBinding =
-              context.getOrGenerateFireFunctionBinding(loadLocal.place);
+              context.getOrGenerateFireFunctionBinding(
+                loadLocal.place,
+                value.loc,
+              );
 
             loadLocal.place = {...fireFunctionBinding};
 
@@ -320,6 +333,69 @@ function visitFunctionExpressionAndPropagateFireDependencies(
   return calleesCapturedByFnExpression;
 }
 
+/*
+ * eachInstructionOperand is not sufficient for our cases because:
+ *  1. fire is a global, which will not appear
+ *  2. The HIR may be malformed, so can't rely on function deps and must
+ *     traverse the whole function.
+ */
+function* eachReachablePlace(fn: HIRFunction): Iterable<Place> {
+  for (const [, block] of fn.body.blocks) {
+    for (const instr of block.instructions) {
+      if (
+        instr.value.kind === 'FunctionExpression' ||
+        instr.value.kind === 'ObjectMethod'
+      ) {
+        yield* eachReachablePlace(instr.value.loweredFunc.func);
+      } else {
+        yield* eachInstructionOperand(instr);
+      }
+    }
+  }
+}
+
+function ensureNoRemainingCalleeCaptures(
+  fn: HIRFunction,
+  context: Context,
+  capturedCallees: FireCalleesToFireFunctionBinding,
+): void {
+  for (const place of eachReachablePlace(fn)) {
+    const calleeInfo = capturedCallees.get(place.identifier.id);
+    if (calleeInfo != null) {
+      const calleeName =
+        calleeInfo.capturedCalleeIdentifier.name?.kind === 'named'
+          ? calleeInfo.capturedCalleeIdentifier.name.value
+          : '<unknown>';
+      context.pushError({
+        loc: place.loc,
+        description: `All uses of ${calleeName} must be either used with a fire() call in \
+this effect or not used with a fire() call at all. ${calleeName} was used with fire() on line \
+${printSourceLocationLine(calleeInfo.fireLoc)} in this effect`,
+        severity: ErrorSeverity.InvalidReact,
+        reason: CANNOT_COMPILE_FIRE,
+        suggestions: null,
+      });
+    }
+  }
+}
+
+function ensureNoMoreFireUses(fn: HIRFunction, context: Context): void {
+  for (const place of eachReachablePlace(fn)) {
+    if (
+      place.identifier.type.kind === 'Function' &&
+      place.identifier.type.shapeId === BuiltInFireId
+    ) {
+      context.pushError({
+        loc: place.identifier.loc,
+        description: 'Cannot use `fire` outside of a useEffect function',
+        severity: ErrorSeverity.Invariant,
+        reason: CANNOT_COMPILE_FIRE,
+        suggestions: null,
+      });
+    }
+  }
+}
+
 function makeLoadUseFireInstruction(env: Environment): Instruction {
   const useFirePlace = createTemporaryPlace(env, GeneratedSource);
   useFirePlace.effect = Effect.Read;
@@ -422,6 +498,7 @@ type FireCalleesToFireFunctionBinding = Map<
   {
     fireFunctionBinding: Place;
     capturedCalleeIdentifier: Identifier;
+    fireLoc: SourceLocation;
   }
 >;
 
@@ -523,8 +600,10 @@ class Context {
   getLoadLocalInstr(id: IdentifierId): LoadLocal | undefined {
     return this.#loadLocals.get(id);
   }
-
-  getOrGenerateFireFunctionBinding(callee: Place): Place {
+  getOrGenerateFireFunctionBinding(
+    callee: Place,
+    fireLoc: SourceLocation,
+  ): Place {
     const fireFunctionBinding = getOrInsertWith(
       this.#fireCalleesToFireFunctions,
       callee.identifier.id,
@@ -534,6 +613,7 @@ class Context {
     this.#capturedCalleeIdentifierIds.set(callee.identifier.id, {
       fireFunctionBinding,
       capturedCalleeIdentifier: callee.identifier,
+      fireLoc,
     });
 
     return fireFunctionBinding;
@@ -575,8 +655,12 @@ class Context {
     return this.#loadGlobalInstructionIds.get(id);
   }
 
+  hasErrors(): boolean {
+    return this.#errors.hasErrors();
+  }
+
   throwIfErrorsFound(): void {
-    if (this.#errors.hasErrors()) throw this.#errors;
+    if (this.hasErrors()) throw this.#errors;
   }
 }
 

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

@@ -21,7 +21,7 @@ function Component(props) {
 ## Code
 
 ```javascript
-import { useFire } from "react";
+import { useFire } from "react/compiler-runtime";
 import { c as _c } from "react/compiler-runtime"; // @enableFire
 import { fire } from "react";
 

compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/transform-fire/deep-scope.expect.md

@@ -30,7 +30,7 @@ function Component(props) {
 ## Code
 
 ```javascript
-import { useFire } from "react";
+import { useFire } from "react/compiler-runtime";
 import { c as _c } from "react/compiler-runtime"; // @enableFire
 import { fire } from "react";
 

compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/transform-fire/error.invalid-mix-fire-and-no-fire.expect.md

@@ -0,0 +1,39 @@
+
+## Input
+
+```javascript
+// @enableFire
+import {fire} from 'react';
+
+function Component(props) {
+  const foo = props => {
+    console.log(props);
+  };
+  useEffect(() => {
+    function nested() {
+      fire(foo(props));
+      foo(props);
+    }
+
+    nested();
+  });
+
+  return null;
+}
+
+```
+
+
+## Error
+
+```
+   9 |     function nested() {
+  10 |       fire(foo(props));
+> 11 |       foo(props);
+     |       ^^^ InvalidReact: Cannot compile `fire`. All uses of foo must be either used with a fire() call in this effect or not used with a fire() call at all. foo was used with fire() on line 10:10 in this effect (11:11)
+  12 |     }
+  13 |
+  14 |     nested();
+```
+
+

compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/transform-fire/error.invalid-mix-fire-and-no-fire.js

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

compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/transform-fire/error.invalid-outside-effect.expect.md

@@ -0,0 +1,38 @@
+
+## Input
+
+```javascript
+// @enableFire
+import {fire, useCallback} from 'react';
+
+function Component({props, bar}) {
+  const foo = () => {
+    console.log(props);
+  };
+  fire(foo(props));
+
+  useCallback(() => {
+    fire(foo(props));
+  }, [foo, props]);
+
+  return null;
+}
+
+```
+
+
+## Error
+
+```
+   6 |     console.log(props);
+   7 |   };
+>  8 |   fire(foo(props));
+     |   ^^^^ Invariant: Cannot compile `fire`. Cannot use `fire` outside of a useEffect function (8:8)
+
+Invariant: Cannot compile `fire`. Cannot use `fire` outside of a useEffect function (11:11)
+   9 |
+  10 |   useCallback(() => {
+  11 |     fire(foo(props));
+```
+
+

compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/transform-fire/error.invalid-outside-effect.js

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

compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/transform-fire/multiple-scope.expect.md

@@ -29,7 +29,7 @@ function Component(props) {
 ## Code
 
 ```javascript
-import { useFire } from "react";
+import { useFire } from "react/compiler-runtime";
 import { c as _c } from "react/compiler-runtime"; // @enableFire
 import { fire } from "react";
 

compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/transform-fire/repeated-calls.expect.md

@@ -22,7 +22,7 @@ function Component(props) {
 ## Code
 
 ```javascript
-import { useFire } from "react";
+import { useFire } from "react/compiler-runtime";
 import { c as _c } from "react/compiler-runtime"; // @enableFire
 import { fire } from "react";
 

compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/transform-fire/shared-hook-calls.expect.md

@@ -26,7 +26,7 @@ function Component({bar, baz}) {
 ## Code
 
 ```javascript
-import { useFire } from "react";
+import { useFire } from "react/compiler-runtime";
 import { c as _c } from "react/compiler-runtime"; // @enableFire
 import { fire } from "react";