[compiler][rfc] Simple retry for fire

Start of a simple retry pipeline. Would love feedback on how we implement this to be extensible to other compiler non-memoization features (e.g. inlineJSX)

Author: mofeiZ

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

@@ -162,7 +162,8 @@ function runWithEnvironment(
   if (
     !env.config.enablePreserveExistingManualUseMemo &&
     !env.config.disableMemoizationForDebugging &&
-    !env.config.enableChangeDetectionForDebugging
+    !env.config.enableChangeDetectionForDebugging &&
+    !env.config.enableMinimalTransformsForRetry
   ) {
     dropManualMemoization(hir);
     log({kind: 'hir', name: 'DropManualMemoization', value: hir});
@@ -279,8 +280,10 @@ function runWithEnvironment(
     value: hir,
   });
 
-  inferReactiveScopeVariables(hir);
-  log({kind: 'hir', name: 'InferReactiveScopeVariables', value: hir});
+  if (!env.config.enableMinimalTransformsForRetry) {
+    inferReactiveScopeVariables(hir);
+    log({kind: 'hir', name: 'InferReactiveScopeVariables', value: hir});
+  }
 
   const fbtOperands = memoizeFbtAndMacroOperandsInSameScope(hir);
   log({

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

@@ -333,6 +333,35 @@ export function compileProgram(
     queue.push({kind: 'original', fn, fnType});
   };
 
+  const runMinimalCompilePipeline = (
+    fn: NodePath<
+      t.FunctionDeclaration | t.ArrowFunctionExpression | t.FunctionExpression
+    >,
+    fnType: ReactFunctionType,
+  ) => {
+    const retryEnvironment = {
+      ...environment,
+      validateHooksUsage: false,
+      validateRefAccessDuringRender: false,
+      validateNoSetStateInRender: false,
+      validateNoSetStateInPassiveEffects: false,
+      validateNoJSXInTryStatements: false,
+      validateMemoizedEffectDependencies: false,
+      validateNoCapitalizedCalls: null,
+      validateBlocklistedImports: null,
+      enableMinimalTransformsForRetry: true,
+    };
+    return compileFn(
+      fn,
+      retryEnvironment,
+      fnType,
+      useMemoCacheIdentifier.name,
+      pass.opts.logger,
+      pass.filename,
+      pass.code,
+    );
+  };
+
   // Main traversal to compile with Forget
   program.traverse(
     {
@@ -382,66 +411,81 @@ export function compileProgram(
       );
     }
 
-    let compiledFn: CodegenFunction;
-    try {
-      /**
-       * Note that Babel does not attach comment nodes to nodes; they are dangling off of the
-       * Program node itself. We need to figure out whether an eslint suppression range
-       * applies to this function first.
-       */
-      const suppressionsInFunction = filterSuppressionsThatAffectFunction(
-        suppressions,
-        fn,
-      );
-      if (suppressionsInFunction.length > 0) {
-        const lintError = suppressionsToCompilerError(suppressionsInFunction);
-        if (optOutDirectives.length > 0) {
-          logError(lintError, pass, fn.node.loc ?? null);
-        } else {
-          handleError(lintError, pass, fn.node.loc ?? null);
-        }
-        return null;
+    /**
+     * Note that Babel does not attach comment nodes to nodes; they are dangling off of the
+     * Program node itself. We need to figure out whether an eslint suppression range
+     * applies to this function first.
+     */
+    const suppressionsInFunction = filterSuppressionsThatAffectFunction(
+      suppressions,
+      fn,
+    );
+    let compileResult:
+      | {kind: 'compile'; compiledFn: CodegenFunction}
+      | {kind: 'error'; error: unknown};
+    if (suppressionsInFunction.length > 0) {
+      compileResult = {
+        kind: 'error',
+        error: suppressionsToCompilerError(suppressionsInFunction),
+      };
+    } else {
+      try {
+        compileResult = {
+          kind: 'compile',
+          compiledFn: compileFn(
+            fn,
+            environment,
+            fnType,
+            useMemoCacheIdentifier.name,
+            pass.opts.logger,
+            pass.filename,
+            pass.code,
+          ),
+        };
+      } catch (err) {
+        compileResult = {kind: 'error', error: err};
       }
-
-      compiledFn = compileFn(
-        fn,
-        environment,
-        fnType,
-        useMemoCacheIdentifier.name,
-        pass.opts.logger,
-        pass.filename,
-        pass.code,
-      );
-      pass.opts.logger?.logEvent(pass.filename, {
-        kind: 'CompileSuccess',
-        fnLoc: fn.node.loc ?? null,
-        fnName: compiledFn.id?.name ?? null,
-        memoSlots: compiledFn.memoSlotsUsed,
-        memoBlocks: compiledFn.memoBlocks,
-        memoValues: compiledFn.memoValues,
-        prunedMemoBlocks: compiledFn.prunedMemoBlocks,
-        prunedMemoValues: compiledFn.prunedMemoValues,
-      });
-    } catch (err) {
+    }
+    // If non-memoization features are enabled, retry regardless of error kind
+    if (compileResult.kind === 'error' && environment.enableFire) {
+      try {
+        compileResult = {
+          kind: 'compile',
+          compiledFn: runMinimalCompilePipeline(fn, fnType),
+        };
+      } catch (err) {
+        compileResult = {kind: 'error', error: err};
+      }
+    }
+    if (compileResult.kind === 'error') {
       /**
        * If an opt out directive is present, log only instead of throwing and don't mark as
        * containing a critical error.
        */
-      if (fn.node.body.type === 'BlockStatement') {
-        if (optOutDirectives.length > 0) {
-          logError(err, pass, fn.node.loc ?? null);
-          return null;
-        }
+      if (optOutDirectives.length > 0) {
+        logError(compileResult.error, pass, fn.node.loc ?? null);
+      } else {
+        handleError(compileResult.error, pass, fn.node.loc ?? null);
       }
-      handleError(err, pass, fn.node.loc ?? null);
       return null;
     }
 
+    pass.opts.logger?.logEvent(pass.filename, {
+      kind: 'CompileSuccess',
+      fnLoc: fn.node.loc ?? null,
+      fnName: compileResult.compiledFn.id?.name ?? null,
+      memoSlots: compileResult.compiledFn.memoSlotsUsed,
+      memoBlocks: compileResult.compiledFn.memoBlocks,
+      memoValues: compileResult.compiledFn.memoValues,
+      prunedMemoBlocks: compileResult.compiledFn.prunedMemoBlocks,
+      prunedMemoValues: compileResult.compiledFn.prunedMemoValues,
+    });
+
     /**
      * Always compile functions with opt in directives.
      */
     if (optInDirectives.length > 0) {
-      return compiledFn;
+      return compileResult.compiledFn;
     } else if (pass.opts.compilationMode === 'annotation') {
       /**
        * No opt-in directive in annotation mode, so don't insert the compiled function.
@@ -467,7 +511,7 @@ export function compileProgram(
     }
 
     if (!pass.opts.noEmit) {
-      return compiledFn;
+      return compileResult.compiledFn;
     }
     return null;
   };

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

@@ -552,6 +552,8 @@ const EnvironmentConfigSchema = z.object({
    */
   disableMemoizationForDebugging: z.boolean().default(false),
 
+  enableMinimalTransformsForRetry: z.boolean().default(false),
+
   /**
    * When true, rather using memoized values, the compiler will always re-compute
    * values, and then use a heuristic to compare the memoized value to the newly

compiler/packages/babel-plugin-react-compiler/src/Inference/InferReferenceEffects.ts

@@ -241,7 +241,7 @@ export default function inferReferenceEffects(
 
   if (options.isFunctionExpression) {
     fn.effects = functionEffects;
-  } else {
+  } else if (!fn.env.config.enableMinimalTransformsForRetry) {
     raiseFunctionEffectErrors(functionEffects);
   }
 }

compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/error.insert-fire.expect.md

@@ -0,0 +1,43 @@
+
+## Input
+
+```javascript
+// @enableFire
+import {fire} from 'react';
+
+function Component({props, bar}) {
+  const foo = () => {
+    console.log(props);
+  };
+  useEffect(() => {
+    fire(foo(props), bar);
+    fire(...foo);
+    fire(bar);
+    fire(props.foo());
+  });
+
+  return null;
+}
+
+```
+
+
+## Error
+
+```
+   7 |   };
+   8 |   useEffect(() => {
+>  9 |     fire(foo(props), bar);
+     |     ^^^^^^^^^^^^^^^^^^^^^ InvalidReact: Cannot compile `fire`. fire() can only take in a single call expression as an argument but received multiple arguments (9:9)
+
+InvalidReact: Cannot compile `fire`. fire() can only take in a single call expression as an argument but received a spread argument (10:10)
+
+InvalidReact: Cannot compile `fire`. `fire()` can only receive a function call such as `fire(fn(a,b)). Method calls and other expressions are not allowed (11:11)
+
+InvalidReact: Cannot compile `fire`. `fire()` can only receive a function call such as `fire(fn(a,b)). Method calls and other expressions are not allowed (12:12)
+  10 |     fire(...foo);
+  11 |     fire(bar);
+  12 |     fire(props.foo());
+```
+          
+      

compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/error.insert-fire.js

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

compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/fire/bailout-retry/bailout-capitalized-fn-call.expect.md

@@ -0,0 +1,50 @@
+
+## Input
+
+```javascript
+// @validateNoCapitalizedCalls @enableFire
+import {fire} from 'react';
+const CapitalizedCall = require('shared-runtime').sum;
+
+function Component({prop1, bar}) {
+  const foo = () => {
+    console.log(prop1);
+  };
+  useEffect(() => {
+    fire(foo(prop1));
+    fire(foo());
+    fire(bar());
+  });
+
+  return CapitalizedCall();
+}
+
+```
+
+## Code
+
+```javascript
+import { useFire } from "react/compiler-runtime"; // @validateNoCapitalizedCalls @enableFire
+import { fire } from "react";
+const CapitalizedCall = require("shared-runtime").sum;
+
+function Component(t0) {
+  const { prop1, bar } = t0;
+  const foo = () => {
+    console.log(prop1);
+  };
+  const t1 = useFire(foo);
+  const t2 = useFire(bar);
+
+  useEffect(() => {
+    t1(prop1);
+    t1();
+    t2();
+  });
+  return CapitalizedCall();
+}
+
+```
+      
+### Eval output
+(kind: exception) Fixture not implemented

compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/fire/bailout-retry/bailout-capitalized-fn-call.js

@@ -0,0 +1,16 @@
+// @validateNoCapitalizedCalls @enableFire
+import {fire} from 'react';
+const CapitalizedCall = require('shared-runtime').sum;
+
+function Component({prop1, bar}) {
+  const foo = () => {
+    console.log(prop1);
+  };
+  useEffect(() => {
+    fire(foo(prop1));
+    fire(foo());
+    fire(bar());
+  });
+
+  return CapitalizedCall();
+}