[compiler] detect and throw on untransformed required features

Traverse program after running compiler transform to find untransformed references to compiler features (e.g. `inferEffectDeps`, `fire`).

Hard error to fail the babel pipeline when the compiler fails to transform these features to give predictable runtime semantics. Untransformed calls to functions like `fire` will throw at runtime anyways, so let's fail the build to catch these earlier.

Note that with this fails the build *regardless of panicThreshold*

Author: mofeiZ

compiler/packages/babel-plugin-react-compiler/src/Babel/BabelPlugin.ts

@@ -11,6 +11,7 @@ import {
   injectReanimatedFlag,
   pipelineUsesReanimatedPlugin,
 } from '../Entrypoint/Reanimated';
+import validateNoUntransformedReferences from '../Entrypoint/ValidateNoUntransformedReferences';
 
 const ENABLE_REACT_COMPILER_TIMINGS =
   process.env['ENABLE_REACT_COMPILER_TIMINGS'] === '1';
@@ -61,12 +62,17 @@ export default function BabelPluginReactCompiler(
               },
             };
           }
-          compileProgram(prog, {
+          const result = compileProgram(prog, {
             opts,
             filename: pass.filename ?? null,
             comments: pass.file.ast.comments ?? [],
             code: pass.file.code,
           });
+          validateNoUntransformedReferences(
+            prog,
+            opts.environment,
+            result?.retryErrors ?? [],
+          );
           if (ENABLE_REACT_COMPILER_TIMINGS === true) {
             performance.mark(`${filename}:end`, {
               detail: 'BabelPlugin:Program:end',

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

@@ -271,6 +271,9 @@ function isFilePartOfSources(
   return false;
 }
 
+type CompileProgramResult = {
+  retryErrors: Array<{fn: BabelFn; error: CompilerError}>;
+};
 /**
  * `compileProgram` is directly invoked by the react-compiler babel plugin, so
  * exceptions thrown by this function will fail the babel build.
@@ -285,16 +288,16 @@ function isFilePartOfSources(
 export function compileProgram(
   program: NodePath<t.Program>,
   pass: CompilerPass,
-): void {
+): CompileProgramResult | null {
   if (shouldSkipCompilation(program, pass)) {
-    return;
+    return null;
   }
 
   const environment = pass.opts.environment;
   const restrictedImportsErr = validateRestrictedImports(program, environment);
   if (restrictedImportsErr) {
     handleError(restrictedImportsErr, pass, null);
-    return;
+    return null;
   }
   const useMemoCacheIdentifier = program.scope.generateUidIdentifier('c');
 
@@ -365,7 +368,7 @@ export function compileProgram(
       filename: pass.filename ?? null,
     },
   );
-
+  const retryErrors: Array<{fn: BabelFn; error: CompilerError}> = [];
   const processFn = (
     fn: BabelFn,
     fnType: ReactFunctionType,
@@ -429,7 +432,9 @@ export function compileProgram(
         handleError(compileResult.error, pass, fn.node.loc ?? null);
       }
       // If non-memoization features are enabled, retry regardless of error kind
-      if (!environment.enableFire) {
+      if (
+        !(environment.enableFire || environment.inferEffectDependencies != null)
+      ) {
         return null;
       }
       try {
@@ -448,6 +453,9 @@ export function compileProgram(
         };
       } catch (err) {
         // TODO: we might want to log error here, but this will also result in duplicate logging
+        if (err instanceof CompilerError) {
+          retryErrors.push({fn, error: err});
+        }
         return null;
       }
     }
@@ -538,7 +546,7 @@ export function compileProgram(
     program.node.directives,
   );
   if (moduleScopeOptOutDirectives.length > 0) {
-    return;
+    return null;
   }
   let gating: null | {
     gatingFn: ExternalFunction;
@@ -596,7 +604,7 @@ export function compileProgram(
     }
   } catch (err) {
     handleError(err, pass, null);
-    return;
+    return null;
   }
 
   /*
@@ -638,6 +646,7 @@ export function compileProgram(
     }
     addImportsToProgram(program, externalFunctions);
   }
+  return {retryErrors};
 }
 
 function shouldSkipCompilation(

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

@@ -0,0 +1,220 @@
+import {NodePath} from '@babel/core';
+import * as t from '@babel/types';
+
+import {CompilerError, EnvironmentConfig} from '..';
+import {getOrInsertWith} from '../Utils/utils';
+import {Environment} from '../HIR';
+
+export default function validateNoUntransformedReferences(
+  path: NodePath<t.Program>,
+  env: EnvironmentConfig,
+  transformErrors: Array<{fn: NodePath<t.Node>; error: CompilerError}>,
+): void {
+  function assertValidEffectImportReference(
+    numArgs: number,
+    paths: Array<NodePath<t.Node>>,
+  ): void {
+    for (const path of paths) {
+      const parent = path.parentPath;
+      if (parent != null && parent.isCallExpression()) {
+        const args = parent.get('arguments');
+        /**
+         * Only error on untransformed references of the form `useMyEffect(...)`
+         * or `moduleNamespace.useMyEffect(...)`, with matching argument counts.
+         * TODO: do we also want a mode to also hard error on non-call references?
+         */
+        if (args.length === numArgs) {
+          const maybeErrorDiagnostic = matchCompilerDiagnostic(
+            path,
+            transformErrors,
+          );
+          /**
+           * Note that we cannot easily check the type of the first argument here,
+           * as it may have already been transformed by the compiler (and not
+           * memoized).
+           */
+          CompilerError.throwInvalidReact({
+            reason:
+              '[Fire] Untransformed reference to compiler-required feature. ' +
+              'Either remove this call or ensure it is successfully transformed by the compiler',
+            description: maybeErrorDiagnostic
+              ? `(Bailout reason: ${maybeErrorDiagnostic})`
+              : null,
+            loc: parent.node.loc ?? null,
+          });
+        }
+      }
+    }
+  }
+
+  function assertValidFireImportReference(
+    paths: Array<NodePath<t.Node>>,
+  ): void {
+    if (paths.length > 0) {
+      const maybeErrorDiagnostic = matchCompilerDiagnostic(
+        paths[0],
+        transformErrors,
+      );
+      CompilerError.throwInvalidReact({
+        reason:
+          '[Fire] Untransformed reference to compiler-required feature. ' +
+          'Either remove this `fire` call or ensure it is successfully transformed by the compiler',
+        description: maybeErrorDiagnostic
+          ? `(Bailout reason: ${maybeErrorDiagnostic})`
+          : null,
+        loc: paths[0].node.loc ?? null,
+      });
+    }
+  }
+
+  const moduleLoadChecks = new Map<
+    string,
+    Map<string, CheckInvalidReferenceFn>
+  >();
+  if (env.enableFire) {
+    /**
+     * Error on any untransformed references to `fire` (e.g. including non-call
+     * expressions)
+     */
+    for (const module of Environment.knownReactModules) {
+      const react = getOrInsertWith(moduleLoadChecks, module, () => new Map());
+      react.set('fire', assertValidFireImportReference);
+    }
+  }
+  if (env.inferEffectDependencies) {
+    for (const {
+      function: {source, importSpecifierName},
+      numRequiredArgs,
+    } of env.inferEffectDependencies) {
+      const module = getOrInsertWith(moduleLoadChecks, source, () => new Map());
+      module.set(
+        importSpecifierName,
+        assertValidEffectImportReference.bind(null, numRequiredArgs),
+      );
+    }
+  }
+  if (moduleLoadChecks.size > 0) {
+    transformProgram(path, moduleLoadChecks);
+  }
+}
+
+type TraversalState = {
+  shouldInvalidateScopes: boolean;
+  program: NodePath<t.Program>;
+};
+type CheckInvalidReferenceFn = (paths: Array<NodePath<t.Node>>) => void;
+function validateImportSpecifier(
+  specifier: NodePath<t.ImportSpecifier>,
+  importSpecifierChecks: Map<string, CheckInvalidReferenceFn>,
+  state: TraversalState,
+): void {
+  const imported = specifier.get('imported');
+  const specifierName: string =
+    imported.node.type === 'Identifier'
+      ? imported.node.name
+      : imported.node.value;
+  const checkFn = importSpecifierChecks.get(specifierName);
+  if (checkFn == null) {
+    return;
+  }
+  if (state.shouldInvalidateScopes) {
+    state.shouldInvalidateScopes = false;
+    state.program.scope.crawl();
+  }
+
+  const local = specifier.get('local');
+  const binding = local.scope.getBinding(local.node.name);
+  CompilerError.invariant(binding != null, {
+    reason: 'Expected binding to be found for import specifier',
+    loc: local.node.loc ?? null,
+  });
+  checkFn(binding.referencePaths);
+}
+
+function validateNamespacedImport(
+  specifier: NodePath<t.ImportNamespaceSpecifier | t.ImportDefaultSpecifier>,
+  importSpecifierChecks: Map<string, CheckInvalidReferenceFn>,
+  state: TraversalState,
+): void {
+  if (state.shouldInvalidateScopes) {
+    state.shouldInvalidateScopes = false;
+    state.program.scope.crawl();
+  }
+  const local = specifier.get('local');
+  const binding = local.scope.getBinding(local.node.name);
+
+  CompilerError.invariant(binding != null, {
+    reason: 'Expected binding to be found for import specifier',
+    loc: local.node.loc ?? null,
+  });
+  const filteredReferences = new Map<
+    CheckInvalidReferenceFn,
+    Array<NodePath<t.Node>>
+  >();
+  for (const reference of binding.referencePaths) {
+    const parent = reference.parentPath;
+    if (
+      parent != null &&
+      parent.isMemberExpression() &&
+      parent.get('object') === reference
+    ) {
+      if (parent.node.computed || parent.node.property.type !== 'Identifier') {
+        continue;
+      }
+      const checkFn = importSpecifierChecks.get(parent.node.property.name);
+      if (checkFn != null) {
+        getOrInsertWith(filteredReferences, checkFn, () => []).push(parent);
+      }
+    }
+  }
+
+  for (const [checkFn, references] of filteredReferences) {
+    checkFn(references);
+  }
+}
+function transformProgram(
+  path: NodePath<t.Program>,
+  moduleLoadChecks: Map<string, Map<string, CheckInvalidReferenceFn>>,
+): void {
+  const traversalState = {shouldInvalidateScopes: true, program: path};
+  path.traverse({
+    ImportDeclaration(path: NodePath<t.ImportDeclaration>) {
+      const importSpecifierChecks = moduleLoadChecks.get(
+        path.node.source.value,
+      );
+      if (importSpecifierChecks == null) {
+        return;
+      }
+      const specifiers = path.get('specifiers');
+      for (const specifier of specifiers) {
+        if (specifier.isImportSpecifier()) {
+          validateImportSpecifier(
+            specifier,
+            importSpecifierChecks,
+            traversalState,
+          );
+        } else {
+          validateNamespacedImport(
+            specifier as NodePath<
+              t.ImportNamespaceSpecifier | t.ImportDefaultSpecifier
+            >,
+            importSpecifierChecks,
+            traversalState,
+          );
+        }
+      }
+    },
+  });
+}
+
+function matchCompilerDiagnostic(
+  badReference: NodePath<t.Node>,
+  transformErrors: Array<{fn: NodePath<t.Node>; error: CompilerError}>,
+): string | null {
+  for (const {fn, error} of transformErrors) {
+    if (fn.isAncestor(badReference)) {
+      return error.toString();
+    }
+  }
+  return null;
+}

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

@@ -1121,6 +1121,7 @@ export class Environment {
       moduleName.toLowerCase() === 'react-dom'
     );
   }
+  static knownReactModules: ReadonlyArray<string> = ['react', 'react-dom'];
 
   getFallthroughPropertyType(
     receiver: Type,

compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/infer-effect-dependencies/bailout-retry/error.callsite-in-non-react-fn.expect.md

@@ -0,0 +1,26 @@
+
+## Input
+
+```javascript
+// @inferEffectDependencies @compilationMode(infer) @panicThreshold(none)
+import {useEffect} from 'react';
+
+function nonReactFn(arg) {
+  useEffect(() => [1, 2, arg]);
+}
+
+```
+
+
+## Error
+
+```
+  3 |
+  4 | function nonReactFn(arg) {
+> 5 |   useEffect(() => [1, 2, arg]);
+    |   ^^^^^^^^^^^^^^^^^^^^^^^^^^^^ InvalidReact: [Fire] Untransformed reference to compiler-required feature. Either remove this call or ensure it is successfully transformed by the compiler (5:5)
+  6 | }
+  7 |
+```
+          
+      

compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/infer-effect-dependencies/bailout-retry/error.callsite-in-non-react-fn.js

@@ -0,0 +1,6 @@
+// @inferEffectDependencies @compilationMode(infer) @panicThreshold(none)
+import {useEffect} from 'react';
+
+function nonReactFn(arg) {
+  useEffect(() => [1, 2, arg]);
+}