[compiler] Infer deps configuration

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,41 @@ 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({
+          module: z.string(),
+          imported: z.string(),
+          numRequiredArgs: z.number(),
+        }),
+      ),
+    )
+    .default(null),
 
   /**
    * Enables inlining ReactElement object literals in place of JSX

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,31 @@ 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.module,
+      () => new Map<string, number>(),
+    );
+    moduleTargets.set(effectTarget.imported, effectTarget.numRequiredArgs);
+    autodepFnConfigs.set(effectTarget.module, moduleTargets);
+  }
+  const autodepFnLoads = new Map<IdentifierId, number>();
+
   const scopeInfos = new Map<
     ScopeId,
     {pruned: boolean; deps: ReactiveScopeDependencies; hasSingleInstr: boolean}
@@ -74,15 +84,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 +150,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 +160,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,63 @@
+
+## Input
+
+```javascript
+// @inferEffectDependencies
+import {useSpecialEffect} from 'react';
+import {print} 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 { useSpecialEffect } from "react";
+import { print } 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,10 @@
+// @inferEffectDependencies
+import {useSpecialEffect} from 'react';
+import {print} 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,9 +174,20 @@ function makePluginOptions(
       .filter(s => s.length > 0);
   }
 
-  let inferEffectDependencies = false;
+  let inferEffectDependencies = null;
   if (firstLine.includes('@inferEffectDependencies')) {
-    inferEffectDependencies = true;
+    inferEffectDependencies = [
+      {
+        module: 'react',
+        imported: 'useEffect',
+        numRequiredArgs: 1,
+      },
+      {
+        module: 'react',
+        imported: 'useSpecialEffect',
+        numRequiredArgs: 2,
+      },
+    ];
   }
 
   let logs: Array<{filename: string | null; event: LoggerEvent}> = [];