[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

@@ -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} />;
+}