[compiler] detect and throw on untransformed required features

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,19 @@ 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,
+            pass.filename ?? null,
+            opts.logger,
+            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,275 @@
+import {NodePath} from '@babel/core';
+import * as t from '@babel/types';
+
+import {
+  CompilerError,
+  CompilerErrorDetailOptions,
+  EnvironmentConfig,
+  ErrorSeverity,
+  Logger,
+} from '..';
+import {getOrInsertWith} from '../Utils/utils';
+import {Environment} from '../HIR';
+import {DEFAULT_EXPORT} from '../HIR/Environment';
+
+function throwInvalidReact(
+  options: Omit<CompilerErrorDetailOptions, 'severity'>,
+  {logger, filename}: TraversalState,
+): never {
+  const detail: CompilerErrorDetailOptions = {
+    ...options,
+    severity: ErrorSeverity.InvalidReact,
+  };
+  logger?.logEvent(filename, {
+    kind: 'CompileError',
+    fnLoc: null,
+    detail,
+  });
+  CompilerError.throw(detail);
+}
+function assertValidEffectImportReference(
+  numArgs: number,
+  paths: Array<NodePath<t.Node>>,
+  context: TraversalState,
+): 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,
+          context.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).
+         */
+        throwInvalidReact(
+          {
+            reason:
+              '[InferEffectDependencies] React Compiler is unable to infer dependencies of this effect. ' +
+              'This will break your build! ' +
+              'To resolve, either pass your own dependency array or fix reported compiler bailout diagnostics.',
+            description: maybeErrorDiagnostic
+              ? `(Bailout reason: ${maybeErrorDiagnostic})`
+              : null,
+            loc: parent.node.loc ?? null,
+          },
+          context,
+        );
+      }
+    }
+  }
+}
+
+function assertValidFireImportReference(
+  paths: Array<NodePath<t.Node>>,
+  context: TraversalState,
+): void {
+  if (paths.length > 0) {
+    const maybeErrorDiagnostic = matchCompilerDiagnostic(
+      paths[0],
+      context.transformErrors,
+    );
+    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,
+      },
+      context,
+    );
+  }
+}
+export default function validateNoUntransformedReferences(
+  path: NodePath<t.Program>,
+  filename: string | null,
+  logger: Logger | null,
+  env: EnvironmentConfig,
+  transformErrors: Array<{fn: NodePath<t.Node>; error: CompilerError}>,
+): void {
+  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, filename, logger, transformErrors);
+  }
+}
+
+type TraversalState = {
+  shouldInvalidateScopes: boolean;
+  program: NodePath<t.Program>;
+  logger: Logger | null;
+  filename: string | null;
+  transformErrors: Array<{fn: NodePath<t.Node>; error: CompilerError}>;
+};
+type CheckInvalidReferenceFn = (
+  paths: Array<NodePath<t.Node>>,
+  context: TraversalState,
+) => 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, state);
+}
+
+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);
+  const defaultCheckFn = importSpecifierChecks.get(DEFAULT_EXPORT);
+
+  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) {
+    if (defaultCheckFn != null) {
+      getOrInsertWith(filteredReferences, defaultCheckFn, () => []).push(
+        reference,
+      );
+    }
+    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, state);
+  }
+}
+function transformProgram(
+  path: NodePath<t.Program>,
+
+  moduleLoadChecks: Map<string, Map<string, CheckInvalidReferenceFn>>,
+  filename: string | null,
+  logger: Logger | null,
+  transformErrors: Array<{fn: NodePath<t.Node>; error: CompilerError}>,
+): void {
+  const traversalState: TraversalState = {
+    shouldInvalidateScopes: true,
+    program: path,
+    filename,
+    logger,
+    transformErrors,
+  };
+  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,