[compiler] Rewrite useContext callee

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

@@ -205,8 +205,8 @@ function* runWithEnvironment(
     validateNoCapitalizedCalls(hir);
   }
 
-  if (env.config.enableLowerContextAccess) {
-    lowerContextAccess(hir);
+  if (env.config.lowerContextAccess) {
+    lowerContextAccess(hir, env.config.lowerContextAccess);
   }
 
   analyseFunctions(hir);

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

@@ -511,6 +511,11 @@ export function compileProgram(
       externalFunctions.push(gating);
     }
 
+    const lowerContextAccess = pass.opts.environment?.lowerContextAccess;
+    if (lowerContextAccess) {
+      externalFunctions.push(tryParseExternalFunction(lowerContextAccess));
+    }
+
     const enableEmitInstrumentForget =
       pass.opts.environment?.enableEmitInstrumentForget;
     if (enableEmitInstrumentForget != null) {

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

@@ -436,8 +436,11 @@ const EnvironmentConfigSchema = z.object({
   enableTreatRefLikeIdentifiersAsRefs: z.boolean().nullable().default(false),
 
   /*
-   * If enabled, this lowers any calls to `useContext` hook to use a selector
-   * function.
+   * If specified a value, the compiler lowers any calls to `useContext` to use
+   * this value as the callee.
+   *
+   * A selector function is compiled and passed as an argument along with the
+   * context to this function call.
    *
    * The compiler automatically figures out the keys by looking for the immediate
    * destructuring of the return value from the useContext call. In the future,
@@ -449,10 +452,10 @@ const EnvironmentConfigSchema = z.object({
    * const {foo, bar} = useContext(MyContext);
    *
    * // output
-   * const {foo, bar} = useContext(MyContext, (c) => [c.foo, c.bar]);
+   * const {foo, bar} = useCompiledContext(MyContext, (c) => [c.foo, c.bar]);
    * ```
    */
-  enableLowerContextAccess: z.boolean().nullable().default(false),
+  lowerContextAccess: ExternalFunctionSchema.nullish(),
 });
 
 export type EnvironmentConfig = z.infer<typeof EnvironmentConfigSchema>;

compiler/packages/babel-plugin-react-compiler/src/Optimization/LowerContextAccess.ts

@@ -11,10 +11,12 @@ import {
   CallExpression,
   Destructure,
   Environment,
+  ExternalFunction,
   GeneratedSource,
   HIRFunction,
   IdentifierId,
   Instruction,
+  LoadGlobal,
   LoadLocal,
   Place,
   PropertyLoad,
@@ -29,7 +31,10 @@ import {createTemporaryPlace} from '../HIR/HIRBuilder';
 import {enterSSA} from '../SSA';
 import {inferTypes} from '../TypeInference';
 
-export function lowerContextAccess(fn: HIRFunction): void {
+export function lowerContextAccess(
+  fn: HIRFunction,
+  loweredContextCallee: ExternalFunction,
+): void {
   const contextAccess: Map<IdentifierId, CallExpression> = new Map();
   const contextKeys: Map<IdentifierId, Array<string>> = new Map();
 
@@ -84,13 +89,23 @@ export function lowerContextAccess(fn: HIRFunction): void {
           isUseContextHookType(value.callee.identifier) &&
           contextKeys.has(lvalue.identifier.id)
         ) {
-          const keys = contextKeys.get(lvalue.identifier.id)!;
-          const selectorFnInstr = emitSelectorFn(fn.env, keys);
+          const loweredContextCalleeInstr = emitLoadLoweredContextCallee(
+            fn.env,
+            loweredContextCallee,
+          );
+
           if (nextInstructions === null) {
             nextInstructions = block.instructions.slice(0, i);
           }
+          nextInstructions.push(loweredContextCalleeInstr);
+
+          const keys = contextKeys.get(lvalue.identifier.id)!;
+          const selectorFnInstr = emitSelectorFn(fn.env, keys);
           nextInstructions.push(selectorFnInstr);
 
+          const lowerContextCallId = loweredContextCalleeInstr.lvalue;
+          value.callee = lowerContextCallId;
+
           const selectorFn = selectorFnInstr.lvalue;
           value.args.push(selectorFn);
         }
@@ -104,9 +119,32 @@ export function lowerContextAccess(fn: HIRFunction): void {
       }
     }
     markInstructionIds(fn.body);
+    inferTypes(fn);
   }
 }
 
+function emitLoadLoweredContextCallee(
+  env: Environment,
+  loweredContextCallee: ExternalFunction,
+): Instruction {
+  const loadGlobal: LoadGlobal = {
+    kind: 'LoadGlobal',
+    binding: {
+      kind: 'ImportNamespace',
+      module: loweredContextCallee.source,
+      name: loweredContextCallee.importSpecifierName,
+    },
+    loc: GeneratedSource,
+  };
+
+  return {
+    id: makeInstructionId(0),
+    loc: GeneratedSource,
+    lvalue: createTemporaryPlace(env, GeneratedSource),
+    value: loadGlobal,
+  };
+}
+
 function getContextKeys(value: Destructure): Array<string> | null {
   const keys = [];
   const pattern = value.lvalue.pattern;

compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/lower-context-acess-multiple.expect.md

@@ -2,7 +2,7 @@
 ## Input
 
 ```javascript
-// @enableLowerContextAccess
+// @lowerContextAccess
 function App() {
   const {foo} = useContext(MyContext);
   const {bar} = useContext(MyContext);
@@ -14,11 +14,12 @@ function App() {
 ## Code
 
 ```javascript
-import { c as _c } from "react/compiler-runtime"; // @enableLowerContextAccess
+import { useContext_withSelector } from "react-compiler-runtime";
+import { c as _c } from "react/compiler-runtime"; // @lowerContextAccess
 function App() {
   const $ = _c(3);
-  const { foo } = useContext(MyContext, _temp);
-  const { bar } = useContext(MyContext, _temp2);
+  const { foo } = useContext_withSelector(MyContext, _temp);
+  const { bar } = useContext_withSelector(MyContext, _temp2);
   let t0;
   if ($[0] !== foo || $[1] !== bar) {
     t0 = <Bar foo={foo} bar={bar} />;

compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/lower-context-acess-multiple.js

@@ -1,4 +1,4 @@
-// @enableLowerContextAccess
+// @lowerContextAccess
 function App() {
   const {foo} = useContext(MyContext);
   const {bar} = useContext(MyContext);

compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/lower-context-selector-simple.expect.md

@@ -2,7 +2,7 @@
 ## Input
 
 ```javascript
-// @enableLowerContextAccess
+// @lowerContextAccess
 function App() {
   const {foo, bar} = useContext(MyContext);
   return <Bar foo={foo} bar={bar} />;
@@ -13,10 +13,11 @@ function App() {
 ## Code
 
 ```javascript
-import { c as _c } from "react/compiler-runtime"; // @enableLowerContextAccess
+import { useContext_withSelector } from "react-compiler-runtime";
+import { c as _c } from "react/compiler-runtime"; // @lowerContextAccess
 function App() {
   const $ = _c(3);
-  const { foo, bar } = useContext(MyContext, _temp);
+  const { foo, bar } = useContext_withSelector(MyContext, _temp);
   let t0;
   if ($[0] !== foo || $[1] !== bar) {
     t0 = <Bar foo={foo} bar={bar} />;

compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/lower-context-selector-simple.js

@@ -1,4 +1,4 @@
-// @enableLowerContextAccess
+// @lowerContextAccess
 function App() {
   const {foo, bar} = useContext(MyContext);
   return <Bar foo={foo} bar={bar} />;

compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/todo.lower-context-access-array-destructuring.expect.md

@@ -2,7 +2,7 @@
 ## Input
 
 ```javascript
-// @enableLowerContextAccess
+// @lowerContextAccess
 function App() {
   const [foo, bar] = useContext(MyContext);
   return <Bar foo={foo} bar={bar} />;
@@ -13,7 +13,8 @@ function App() {
 ## Code
 
 ```javascript
-import { c as _c } from "react/compiler-runtime"; // @enableLowerContextAccess
+import { useContext_withSelector } from "react-compiler-runtime";
+import { c as _c } from "react/compiler-runtime"; // @lowerContextAccess
 function App() {
   const $ = _c(3);
   const [foo, bar] = useContext(MyContext);

compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/todo.lower-context-access-array-destructuring.js

@@ -1,4 +1,4 @@
-// @enableLowerContextAccess
+// @lowerContextAccess
 function App() {
   const [foo, bar] = useContext(MyContext);
   return <Bar foo={foo} bar={bar} />;

compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/todo.lower-context-access-destructure-multiple.expect.md

@@ -2,7 +2,7 @@
 ## Input
 
 ```javascript
-// @enableLowerContextAccess
+// @lowerContextAccess
 function App() {
   const context = useContext(MyContext);
   const {foo} = context;
@@ -15,7 +15,8 @@ function App() {
 ## Code
 
 ```javascript
-import { c as _c } from "react/compiler-runtime"; // @enableLowerContextAccess
+import { useContext_withSelector } from "react-compiler-runtime";
+import { c as _c } from "react/compiler-runtime"; // @lowerContextAccess
 function App() {
   const $ = _c(3);
   const context = useContext(MyContext);

compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/todo.lower-context-access-destructure-multiple.js

@@ -1,4 +1,4 @@
-// @enableLowerContextAccess
+// @lowerContextAccess
 function App() {
   const context = useContext(MyContext);
   const {foo} = context;

compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/todo.lower-context-access-mixed-array-obj.expect.md

@@ -2,7 +2,7 @@
 ## Input
 
 ```javascript
-// @enableLowerContextAccess
+// @lowerContextAccess
 function App() {
   const context = useContext(MyContext);
   const [foo] = context;
@@ -15,7 +15,8 @@ function App() {
 ## Code
 
 ```javascript
-import { c as _c } from "react/compiler-runtime"; // @enableLowerContextAccess
+import { useContext_withSelector } from "react-compiler-runtime";
+import { c as _c } from "react/compiler-runtime"; // @lowerContextAccess
 function App() {
   const $ = _c(3);
   const context = useContext(MyContext);

compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/todo.lower-context-access-mixed-array-obj.js

@@ -1,4 +1,4 @@
-// @enableLowerContextAccess
+// @lowerContextAccess
 function App() {
   const context = useContext(MyContext);
   const [foo] = context;

compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/todo.lower-context-access-nested-destructuring.expect.md

@@ -2,7 +2,7 @@
 ## Input
 
 ```javascript
-// @enableLowerContextAccess
+// @lowerContextAccess
 function App() {
   const {
     joe: {foo},
@@ -16,7 +16,8 @@ function App() {
 ## Code
 
 ```javascript
-import { c as _c } from "react/compiler-runtime"; // @enableLowerContextAccess
+import { useContext_withSelector } from "react-compiler-runtime";
+import { c as _c } from "react/compiler-runtime"; // @lowerContextAccess
 function App() {
   const $ = _c(3);
   const { joe: t0, bar } = useContext(MyContext);

compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/todo.lower-context-access-nested-destructuring.js

@@ -1,4 +1,4 @@
-// @enableLowerContextAccess
+// @lowerContextAccess
 function App() {
   const {
     joe: {foo},

compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/todo.lower-context-access-property-load.expect.md

@@ -2,7 +2,7 @@
 ## Input
 
 ```javascript
-// @enableLowerContextAccess
+// @lowerContextAccess
 function App() {
   const context = useContext(MyContext);
   const foo = context.foo;
@@ -15,7 +15,8 @@ function App() {
 ## Code
 
 ```javascript
-import { c as _c } from "react/compiler-runtime"; // @enableLowerContextAccess
+import { useContext_withSelector } from "react-compiler-runtime";
+import { c as _c } from "react/compiler-runtime"; // @lowerContextAccess
 function App() {
   const $ = _c(3);
   const context = useContext(MyContext);

compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/todo.lower-context-access-property-load.js

@@ -1,4 +1,4 @@
-// @enableLowerContextAccess
+// @lowerContextAccess
 function App() {
   const context = useContext(MyContext);
   const foo = context.foo;

compiler/packages/snap/src/compiler.ts

@@ -155,6 +155,14 @@ function makePluginOptions(
       .filter(s => s.length > 0);
   }
 
+  let lowerContextAccess = null;
+  if (firstLine.includes('@lowerContextAccess')) {
+    lowerContextAccess = {
+      source: 'react-compiler-runtime',
+      importSpecifierName: 'useContext_withSelector',
+    };
+  }
+
   let logs: Array<{filename: string | null; event: LoggerEvent}> = [];
   let logger: Logger | null = null;
   if (firstLine.includes('@logger')) {
@@ -207,6 +215,7 @@ function makePluginOptions(
       hookPattern,
       validatePreserveExistingMemoizationGuarantees,
       enableChangeDetectionForDebugging,
+      lowerContextAccess,
     },
     compilationMode,
     logger,