[compiler] Infer deps configuration (facebook#31616)

Adds a way to configure how we insert deps for experimental purposes.

```
[
  {
    module: 'react',
    imported: 'useEffect',
    numRequiredArgs: 1,
  },
  {
    module: 'MyExperimentalEffectHooks',
    imported: 'useExperimentalEffect',
    numRequiredArgs: 2,
  },
]
```

would insert dependencies for calls of `useEffect` imported from `react`
if they have 1 argument and calls of useExperimentalEffect` from
`MyExperimentalEffectHooks` if they have 2 arguments. The pushed dep
array is appended to the arg list.

Author: jbrown215

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

@@ -356,7 +356,7 @@ function* runWithEnvironment(
   });
 
   if (env.config.inferEffectDependencies) {
-    inferEffectDependencies(env, hir);
+    inferEffectDependencies(hir);
   }
 
   if (env.config.inlineJsxTransform) {

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

@@ -242,9 +242,40 @@ const EnvironmentConfigSchema = z.object({
   enableOptionalDependencies: z.boolean().default(true),
 
   /**
-   * Enables inference and auto-insertion of effect dependencies. Still experimental.
+   * Enables inference and auto-insertion of effect dependencies. Takes in an array of
+   * configurable module and import pairs to allow for user-land experimentation. For example,
+   * [
+   *   {
+   *     module: 'react',
+   *     imported: 'useEffect',
+   *     numRequiredArgs: 1,
+   *   },{
+   *     module: 'MyExperimentalEffectHooks',
+   *     imported: 'useExperimentalEffect',
+   *     numRequiredArgs: 2,
+   *   },
+   * ]
+   * would insert dependencies for calls of `useEffect` imported from `react` and calls of
+   * useExperimentalEffect` from `MyExperimentalEffectHooks`.
+   *
+   * `numRequiredArgs` tells the compiler the amount of arguments required to append a dependency
+   *  array to the end of the call. With the configuration above, we'd insert dependencies for
+   *  `useEffect` if it is only given a single argument and it would be appended to the argument list.
+   *
+   * numRequiredArgs must always be greater than 0, otherwise there is no function to analyze for dependencies
+   *
+   * Still experimental.
    */
-  inferEffectDependencies: z.boolean().default(false),
+  inferEffectDependencies: z
+    .nullable(
+      z.array(
+        z.object({
+          function: ExternalFunctionSchema,
+          numRequiredArgs: z.number(),
+        }),
+      ),
+    )
+    .default(null),
 
   /**
    * Enables inlining ReactElement object literals in place of JSX
@@ -614,6 +645,22 @@ const testComplexConfigDefaults: PartialEnvironmentConfig = {
     source: 'react-compiler-runtime',
     importSpecifierName: 'useContext_withSelector',
   },
+  inferEffectDependencies: [
+    {
+      function: {
+        source: 'react',
+        importSpecifierName: 'useEffect',
+      },
+      numRequiredArgs: 1,
+    },
+    {
+      function: {
+        source: 'shared-runtime',
+        importSpecifierName: 'useSpecialEffect',
+      },
+      numRequiredArgs: 2,
+    },
+  ],
 };
 
 /**

compiler/packages/babel-plugin-react-compiler/src/Inference/InferEffectDependencies.ts

@@ -8,7 +8,6 @@ import {
   HIRFunction,
   IdentifierId,
   Instruction,
-  isUseEffectHookType,
   makeInstructionId,
   TInstruction,
   InstructionId,
@@ -23,20 +22,33 @@ import {
   markInstructionIds,
 } from '../HIR/HIRBuilder';
 import {eachInstructionOperand, eachTerminalOperand} from '../HIR/visitors';
+import {getOrInsertWith} from '../Utils/utils';
 
 /**
  * Infers reactive dependencies captured by useEffect lambdas and adds them as
  * a second argument to the useEffect call if no dependency array is provided.
  */
-export function inferEffectDependencies(
-  env: Environment,
-  fn: HIRFunction,
-): void {
+export function inferEffectDependencies(fn: HIRFunction): void {
   let hasRewrite = false;
   const fnExpressions = new Map<
     IdentifierId,
     TInstruction<FunctionExpression>
   >();
+
+  const autodepFnConfigs = new Map<string, Map<string, number>>();
+  for (const effectTarget of fn.env.config.inferEffectDependencies!) {
+    const moduleTargets = getOrInsertWith(
+      autodepFnConfigs,
+      effectTarget.function.source,
+      () => new Map<string, number>(),
+    );
+    moduleTargets.set(
+      effectTarget.function.importSpecifierName,
+      effectTarget.numRequiredArgs,
+    );
+  }
+  const autodepFnLoads = new Map<IdentifierId, number>();
+
   const scopeInfos = new Map<
     ScopeId,
     {pruned: boolean; deps: ReactiveScopeDependencies; hasSingleInstr: boolean}
@@ -74,15 +86,23 @@ export function inferEffectDependencies(
           lvalue.identifier.id,
           instr as TInstruction<FunctionExpression>,
         );
+      } else if (
+        value.kind === 'LoadGlobal' &&
+        value.binding.kind === 'ImportSpecifier'
+      ) {
+        const moduleTargets = autodepFnConfigs.get(value.binding.module);
+        if (moduleTargets != null) {
+          const numRequiredArgs = moduleTargets.get(value.binding.imported);
+          if (numRequiredArgs != null) {
+            autodepFnLoads.set(lvalue.identifier.id, numRequiredArgs);
+          }
+        }
       } else if (
         /*
-         * This check is not final. Right now we only look for useEffects without a dependency array.
-         * This is likely not how we will ship this feature, but it is good enough for us to make progress
-         * on the implementation and test it.
+         * TODO: Handle method calls
          */
         value.kind === 'CallExpression' &&
-        isUseEffectHookType(value.callee.identifier) &&
-        value.args.length === 1 &&
+        autodepFnLoads.get(value.callee.identifier.id) === value.args.length &&
         value.args[0].kind === 'Identifier'
       ) {
         const fnExpr = fnExpressions.get(value.args[0].identifier.id);
@@ -132,7 +152,7 @@ export function inferEffectDependencies(
             loc: GeneratedSource,
           };
 
-          const depsPlace = createTemporaryPlace(env, GeneratedSource);
+          const depsPlace = createTemporaryPlace(fn.env, GeneratedSource);
           depsPlace.effect = Effect.Read;
 
           newInstructions.push({
@@ -142,8 +162,8 @@ export function inferEffectDependencies(
             value: deps,
           });
 
-          // Step 2: insert the deps array as an argument of the useEffect
-          value.args[1] = {...depsPlace, effect: Effect.Freeze};
+          // Step 2: push the inferred deps array as an argument of the useEffect
+          value.args.push({...depsPlace, effect: Effect.Freeze});
           rewriteInstrs.set(instr.id, newInstructions);
         }
       }

compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/infer-deps-custom-config.expect.md

@@ -0,0 +1,61 @@
+
+## Input
+
+```javascript
+// @inferEffectDependencies
+import {print, useSpecialEffect} from 'shared-runtime';
+
+function CustomConfig({propVal}) {
+  // Insertion
+  useSpecialEffect(() => print(propVal), [propVal]);
+  // No insertion
+  useSpecialEffect(() => print(propVal), [propVal], [propVal]);
+}
+
+```
+
+## Code
+
+```javascript
+import { c as _c } from "react/compiler-runtime"; // @inferEffectDependencies
+import { print, useSpecialEffect } from "shared-runtime";
+
+function CustomConfig(t0) {
+  const $ = _c(7);
+  const { propVal } = t0;
+  let t1;
+  let t2;
+  if ($[0] !== propVal) {
+    t1 = () => print(propVal);
+    t2 = [propVal];
+    $[0] = propVal;
+    $[1] = t1;
+    $[2] = t2;
+  } else {
+    t1 = $[1];
+    t2 = $[2];
+  }
+  useSpecialEffect(t1, t2, [propVal]);
+  let t3;
+  let t4;
+  let t5;
+  if ($[3] !== propVal) {
+    t3 = () => print(propVal);
+    t4 = [propVal];
+    t5 = [propVal];
+    $[3] = propVal;
+    $[4] = t3;
+    $[5] = t4;
+    $[6] = t5;
+  } else {
+    t3 = $[4];
+    t4 = $[5];
+    t5 = $[6];
+  }
+  useSpecialEffect(t3, t4, t5);
+}
+
+```
+      
+### Eval output
+(kind: exception) Fixture not implemented

compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/infer-deps-custom-config.js

@@ -0,0 +1,9 @@
+// @inferEffectDependencies
+import {print, useSpecialEffect} from 'shared-runtime';
+
+function CustomConfig({propVal}) {
+  // Insertion
+  useSpecialEffect(() => print(propVal), [propVal]);
+  // No insertion
+  useSpecialEffect(() => print(propVal), [propVal], [propVal]);
+}

compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/infer-effect-dependencies.expect.md

@@ -3,6 +3,8 @@
 
 ```javascript
 // @inferEffectDependencies
+import {useEffect, useRef} from 'react';
+
 const moduleNonReactive = 0;
 
 function Component({foo, bar}) {
@@ -45,6 +47,8 @@ function Component({foo, bar}) {
 
 ```javascript
 import { c as _c } from "react/compiler-runtime"; // @inferEffectDependencies
+import { useEffect, useRef } from "react";
+
 const moduleNonReactive = 0;
 
 function Component(t0) {

compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/infer-effect-dependencies.js

@@ -1,4 +1,6 @@
 // @inferEffectDependencies
+import {useEffect, useRef} from 'react';
+
 const moduleNonReactive = 0;
 
 function Component({foo, bar}) {

compiler/packages/snap/src/compiler.ts

@@ -174,11 +174,6 @@ function makePluginOptions(
       .filter(s => s.length > 0);
   }
 
-  let inferEffectDependencies = false;
-  if (firstLine.includes('@inferEffectDependencies')) {
-    inferEffectDependencies = true;
-  }
-
   let logs: Array<{filename: string | null; event: LoggerEvent}> = [];
   let logger: Logger | null = null;
   if (firstLine.includes('@logger')) {
@@ -202,7 +197,6 @@ function makePluginOptions(
       hookPattern,
       validatePreserveExistingMemoizationGuarantees,
       validateBlocklistedImports,
-      inferEffectDependencies,
     },
     compilationMode,
     logger,

compiler/packages/snap/src/sprout/shared-runtime.ts

@@ -363,11 +363,20 @@ export function useFragment(..._args: Array<any>): object {
   };
 }
 
+export function useSpecialEffect(
+  fn: () => any,
+  _secondArg: any,
+  deps: Array<any>,
+) {
+  React.useEffect(fn, deps);
+}
+
 export function typedArrayPush<T>(array: Array<T>, item: T): void {
   array.push(item);
 }
 
 export function typedLog(...values: Array<any>): void {
   console.log(...values);
 }
+
 export default typedLog;