[compiler] detect and hard error 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.

TODO[next PR]: throw retry pipeline errors (either directly or after aggregating across functions). This should show better errors than the current generic TODO message

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';
@@ -73,7 +74,13 @@ export default function BabelPluginReactCompiler(
             });
           }
         },
-        exit(_, pass): void {
+        exit(prog, pass): void {
+          prog.scope.crawl();
+          validateNoUntransformedReferences(
+            prog,
+            (pass.opts as any).environment,
+          );
+
           if (ENABLE_REACT_COMPILER_TIMINGS === true) {
             const filename = pass.filename ?? 'unknown';
             const measurement = performance.measure(filename, {

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

@@ -0,0 +1,187 @@
+import {NodePath} from '@babel/core';
+import * as t from '@babel/types';
+
+import {CompilerError, EnvironmentConfig} from '..';
+import {getOrInsertWith} from '../Utils/utils';
+
+export default function validateNoUntransformedReferences(
+  path: NodePath<t.Program>,
+  env: EnvironmentConfig,
+): 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)
+     */
+    const react = getOrInsertWith(moduleLoadChecks, 'react', () => new Map());
+    react.set('fire', assertNone);
+  }
+  if (env.inferEffectDependencies) {
+    /**
+     * Only error on untransformed references of the form `useMyEffect(...)` or
+     * `moduleNamespace.useMyEffect(...)`, with matching argument counts.
+     * TODO: add mode to also hard error on non-call references
+     */
+    for (const {
+      function: {source, importSpecifierName},
+      numRequiredArgs,
+    } of env.inferEffectDependencies) {
+      const module = getOrInsertWith(moduleLoadChecks, source, () => new Map());
+      module.set(
+        importSpecifierName,
+        assertNoAutoDepCalls.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 {
+  const local = (
+    specifier as NodePath<t.ImportNamespaceSpecifier | t.ImportDefaultSpecifier>
+  ).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) {
+    if (state.shouldInvalidateScopes) {
+      state.shouldInvalidateScopes = false;
+      state.program.scope.crawl();
+    }
+    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 assertNoAutoDepCalls(
+  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');
+      if (args.length === numArgs) {
+        /**
+         * 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.throwTodo({
+          reason:
+            'Untransformed reference to experimental compiler-only feature',
+          loc: parent.node.loc ?? null,
+        });
+      }
+    }
+  }
+}
+
+function assertNone(paths: Array<NodePath<t.Node>>): void {
+  if (paths.length > 0) {
+    CompilerError.throwTodo({
+      reason: 'Untransformed reference to experimental compiler-only feature',
+      loc: paths[0].node.loc ?? null,
+    });
+  }
+}

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]);
+    |   ^^^^^^^^^^^^^^^^^^^^^^^^^^^^ Todo: Untransformed reference to experimental compiler-only feature (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]);
+}

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

@@ -0,0 +1,41 @@
+
+## Input
+
+```javascript
+// @inferEffectDependencies @panicThreshold(none)
+import {useEffect} from 'react';
+
+/**
+ * Error on non-inlined effect functions:
+ * 1. From the effect hook callee's perspective, it only makes sense
+ *    to either
+ *    (a) never hard error (i.e. failing to infer deps is acceptable) or
+ *    (b) always hard error,
+ *    regardless of whether the callback function is an inline fn.
+ * 2. (Technical detail) it's harder to support detecting cases in which
+ *    function (pre-Forget transform) was inline but becomes memoized
+ */
+function Component({foo}) {
+  function f() {
+    console.log(foo);
+  }
+
+  // No inferred dep array, the argument is not a lambda
+  useEffect(f);
+}
+
+```
+
+
+## Error
+
+```
+  18 |
+  19 |   // No inferred dep array, the argument is not a lambda
+> 20 |   useEffect(f);
+     |   ^^^^^^^^^^^^ Todo: Untransformed reference to experimental compiler-only feature (20:20)
+  21 | }
+  22 |
+```
+          
+      

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

@@ -0,0 +1,21 @@
+// @inferEffectDependencies @panicThreshold(none)
+import {useEffect} from 'react';
+
+/**
+ * Error on non-inlined effect functions:
+ * 1. From the effect hook callee's perspective, it only makes sense
+ *    to either
+ *    (a) never hard error (i.e. failing to infer deps is acceptable) or
+ *    (b) always hard error,
+ *    regardless of whether the callback function is an inline fn.
+ * 2. (Technical detail) it's harder to support detecting cases in which
+ *    function (pre-Forget transform) was inline but becomes memoized
+ */
+function Component({foo}) {
+  function f() {
+    console.log(foo);
+  }
+
+  // No inferred dep array, the argument is not a lambda
+  useEffect(f);
+}

compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/infer-effect-dependencies/bailout-retry/error.todo-import-default-property-useEffect.expect.md

@@ -0,0 +1,27 @@
+
+## Input
+
+```javascript
+// @inferEffectDependencies @panicThreshold(none)
+import React from 'react';
+
+function Component() {
+  const obj = makeObject_Primitives();
+  React.useEffect(() => print(obj));
+}
+
+```
+
+
+## Error
+
+```
+  4 | function Component() {
+  5 |   const obj = makeObject_Primitives();
+> 6 |   React.useEffect(() => print(obj));
+    |   ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ Todo: Untransformed reference to experimental compiler-only feature (6:6)
+  7 | }
+  8 |
+```
+          
+      

compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/infer-effect-dependencies/todo.import-default-property-useEffect.js renamed to compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/infer-effect-dependencies/bailout-retry/error.todo-import-default-property-useEffect.js

@@ -1,7 +1,7 @@
-// @inferEffectDependencies
+// @inferEffectDependencies @panicThreshold(none)
 import React from 'react';
 
-function NonReactiveDepInEffect() {
+function Component() {
   const obj = makeObject_Primitives();
   React.useEffect(() => print(obj));
 }

compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/infer-effect-dependencies/bailout-retry/error.use-no-memo.expect.md

@@ -0,0 +1,27 @@
+
+## Input
+
+```javascript
+// @inferEffectDependencies @panicThreshold(none)
+import {useEffect} from 'react';
+
+function Component({propVal}) {
+  'use no memo';
+  useEffect(() => [propVal]);
+}
+
+```
+
+
+## Error
+
+```
+  4 | function Component({propVal}) {
+  5 |   'use no memo';
+> 6 |   useEffect(() => [propVal]);
+    |   ^^^^^^^^^^^^^^^^^^^^^^^^^^ Todo: Untransformed reference to experimental compiler-only feature (6:6)
+  7 | }
+  8 |
+```
+          
+      

compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/infer-effect-dependencies/bailout-retry/error.use-no-memo.js

@@ -0,0 +1,7 @@
+// @inferEffectDependencies @panicThreshold(none)
+import {useEffect} from 'react';
+
+function Component({propVal}) {
+  'use no memo';
+  useEffect(() => [propVal]);
+}