[compiler][rfc] Hacky retry pipeline for fire (facebook#32164)

Hacky retry pipeline for when transforming `fire(...)` calls encounters
validation, todo, or memoization invariant bailouts. Would love feedback
on how we implement this to be extensible to other compiler
non-memoization features (e.g. inlineJSX)

Some observations:
- Compiler "front-end" passes (e.g. lower, type, effect, and mutability
inferences) should be shared for all compiler features -- memo and
otherwise
- Many passes (anything dealing with reactive scope ranges, scope blocks
/ dependencies, and optimizations such as ReactiveIR facebook#31974) can be left
out of the retry pipeline. This PR hackily skips memoization features by
removing reactive scope creation, but we probably should restructure the
pipeline to skip these entirely on a retry
- We should maintain a canonical set of "validation flags"

Note the newly added fixtures are prefixed with `bailout-...` when the
retry fire pipeline is used. These fixture outputs contain correctly
inserted `useFire` calls and no memoization.

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

@@ -16,6 +16,7 @@ import {
   EnvironmentConfig,
   ExternalFunction,
   ReactFunctionType,
+  MINIMAL_RETRY_CONFIG,
   tryParseExternalFunction,
 } from '../HIR/Environment';
 import {CodegenFunction} from '../ReactiveScopes';
@@ -382,66 +383,92 @@ 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: compileFn(
+            fn,
+            {
+              ...environment,
+              ...MINIMAL_RETRY_CONFIG,
+            },
+            fnType,
+            useMemoCacheIdentifier.name,
+            pass.opts.logger,
+            pass.filename,
+            pass.code,
+          ),
+        };
+      } 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 +494,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
@@ -626,6 +628,17 @@ const EnvironmentConfigSchema = z.object({
 
 export type EnvironmentConfig = z.infer<typeof EnvironmentConfigSchema>;
 
+export const MINIMAL_RETRY_CONFIG: PartialEnvironmentConfig = {
+  validateHooksUsage: false,
+  validateRefAccessDuringRender: false,
+  validateNoSetStateInRender: false,
+  validateNoSetStateInPassiveEffects: false,
+  validateNoJSXInTryStatements: false,
+  validateMemoizedEffectDependencies: false,
+  validateNoCapitalizedCalls: null,
+  validateBlocklistedImports: null,
+  enableMinimalTransformsForRetry: true,
+};
 /**
  * For test fixtures and playground only.
  *

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/transform-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/transform-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();
+}

compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/transform-fire/bailout-retry/bailout-eslint-suppressions.expect.md

@@ -0,0 +1,55 @@
+
+## Input
+
+```javascript
+// @enableFire
+import {useRef} from 'react';
+
+function Component({props, bar}) {
+  const foo = () => {
+    console.log(props);
+  };
+  useEffect(() => {
+    fire(foo(props));
+    fire(foo());
+    fire(bar());
+  });
+
+  const ref = useRef(null);
+  // eslint-disable-next-line react-hooks/rules-of-hooks
+  ref.current = 'bad';
+  return <button ref={ref} />;
+}
+
+```
+
+## Code
+
+```javascript
+import { useFire } from "react/compiler-runtime"; // @enableFire
+import { useRef } from "react";
+
+function Component(t0) {
+  const { props, bar } = t0;
+  const foo = () => {
+    console.log(props);
+  };
+  const t1 = useFire(foo);
+  const t2 = useFire(bar);
+
+  useEffect(() => {
+    t1(props);
+    t1();
+    t2();
+  });
+
+  const ref = useRef(null);
+
+  ref.current = "bad";
+  return <button ref={ref} />;
+}
+
+```
+      
+### Eval output
+(kind: exception) Fixture not implemented

compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/transform-fire/bailout-retry/bailout-eslint-suppressions.js

@@ -0,0 +1,18 @@
+// @enableFire
+import {useRef} from 'react';
+
+function Component({props, bar}) {
+  const foo = () => {
+    console.log(props);
+  };
+  useEffect(() => {
+    fire(foo(props));
+    fire(foo());
+    fire(bar());
+  });
+
+  const ref = useRef(null);
+  // eslint-disable-next-line react-hooks/rules-of-hooks
+  ref.current = 'bad';
+  return <button ref={ref} />;
+}