Support use operator

Implements support for use: 

* Teaches InferReactivePlaces to treat use() result as reactive 

* Teaches FlattenScopesWithHooks to also flatten scopes with use() 

Handles both `use()` and `React.use()`.

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

@@ -47,7 +47,7 @@ import {
   codegenFunction,
   extractScopeDeclarationsFromDestructuring,
   flattenReactiveLoops,
-  flattenScopesWithHooks,
+  flattenScopesWithHooksOrUse,
   inferReactiveScopeVariables,
   memoizeFbtOperandsInSameScope,
   mergeOverlappingReactiveScopes,
@@ -277,7 +277,7 @@ function* runWithEnvironment(
 
   assertScopeInstructionsWithinScopes(reactiveFunction);
 
-  flattenScopesWithHooks(reactiveFunction);
+  flattenScopesWithHooksOrUse(reactiveFunction);
   yield log({
     kind: "reactive",
     name: "FlattenScopesWithHooks",

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

@@ -572,11 +572,6 @@ export class Environment {
 
 // From https://github.com/facebook/react/blob/main/packages/eslint-plugin-react-hooks/src/RulesOfHooks.js#LL18C1-L23C2
 export function isHookName(name: string): boolean {
-  /*
-   * if (__EXPERIMENTAL__) {
-   *   return name === 'use' || /^use[A-Z0-9]/.test(name);
-   * }
-   */
   return /^use[A-Z0-9]/.test(name);
 }
 

compiler/packages/babel-plugin-react-forget/src/HIR/Globals.ts

@@ -12,6 +12,7 @@ import {
   BuiltInUseEffectHookId,
   BuiltInUseInsertionEffectHookId,
   BuiltInUseLayoutEffectHookId,
+  BuiltInUseOperatorId,
   BuiltInUseRefId,
   BuiltInUseStateId,
   ShapeRegistry,
@@ -239,7 +240,7 @@ const TYPED_GLOBALS: Array<[string, BuiltInType]> = [
  * now that FeatureFlag `enableTreatHooksAsFunctions` is removed we can
  * use positional params too (?)
  */
-const BUILTIN_HOOKS: Array<[string, BuiltInType]> = [
+const REACT_APIS: Array<[string, BuiltInType]> = [
   [
     "useContext",
     addHook(DEFAULT_SHAPES, {
@@ -342,13 +343,28 @@ const BUILTIN_HOOKS: Array<[string, BuiltInType]> = [
       BuiltInUseInsertionEffectHookId
     ),
   ],
+  [
+    "use",
+    addFunction(
+      DEFAULT_SHAPES,
+      [],
+      {
+        positionalParams: [],
+        restParam: Effect.Freeze,
+        returnType: { kind: "Poly" },
+        calleeEffect: Effect.Read,
+        returnValueKind: ValueKind.Frozen,
+      },
+      BuiltInUseOperatorId
+    ),
+  ],
 ];
 
 TYPED_GLOBALS.push(
   [
     "React",
     addObject(DEFAULT_SHAPES, null, [
-      ...BUILTIN_HOOKS,
+      ...REACT_APIS,
       [
         "createElement",
         addFunction(DEFAULT_SHAPES, [], {
@@ -395,7 +411,7 @@ TYPED_GLOBALS.push(
 
 export type Global = BuiltInType | PolyType;
 export type GlobalRegistry = Map<string, Global>;
-export const DEFAULT_GLOBALS: GlobalRegistry = new Map(BUILTIN_HOOKS);
+export const DEFAULT_GLOBALS: GlobalRegistry = new Map(REACT_APIS);
 
 // Hack until we add ObjectShapes for all globals
 for (const name of UNTYPED_GLOBALS) {

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

@@ -1416,6 +1416,12 @@ export function getHookKind(env: Environment, id: Identifier): HookKind | null {
   return getHookKindForType(env, id.type);
 }
 
+export function isUseOperator(id: Identifier): boolean {
+  return (
+    id.type.kind === "Function" && id.type.shapeId === "BuiltInUseOperator"
+  );
+}
+
 export function getHookKindForType(
   env: Environment,
   type: Type

compiler/packages/babel-plugin-react-forget/src/HIR/ObjectShape.ts

@@ -197,6 +197,7 @@ export const BuiltInMixedReadonlyId = "BuiltInMixedReadonly";
 export const BuiltInUseEffectHookId = "BuiltInUseEffectHook";
 export const BuiltInUseLayoutEffectHookId = "BuiltInUseLayoutEffectHook";
 export const BuiltInUseInsertionEffectHookId = "BuiltInUseInsertionEffectHook";
+export const BuiltInUseOperatorId = "BuiltInUseOperator";
 
 // ShapeRegistry with default definitions for built-ins.
 export const BUILTIN_SHAPES: ShapeRegistry = new Map();

compiler/packages/babel-plugin-react-forget/src/Inference/InferReactivePlaces.ts

@@ -16,6 +16,7 @@ import {
   computePostDominatorTree,
   getHookKind,
   isSetStateType,
+  isUseOperator,
 } from "../HIR";
 import { PostDominator } from "../HIR/Dominator";
 import {
@@ -195,18 +196,23 @@ export function inferReactivePlaces(fn: HIRFunction): void {
           hasReactiveInput ||= reactive;
         }
 
-        /*
-         * Hooks may always return a reactive variable, even if their inputs are
-         * non-reactive, because they can access state or context.
+        /**
+         * Hooks and the 'use' operator are sources of reactivity because
+         * they can access state (for hooks) or context (for hooks/use).
+         *
+         * Technically, `use` could be used to await a non-reactive promise,
+         * but we are conservative and assume that the value could be reactive.
          */
         if (
           value.kind === "CallExpression" &&
-          getHookKind(fn.env, value.callee.identifier) != null
+          (getHookKind(fn.env, value.callee.identifier) != null ||
+            isUseOperator(value.callee.identifier))
         ) {
           hasReactiveInput = true;
         } else if (
           value.kind === "MethodCall" &&
-          getHookKind(fn.env, value.property.identifier) != null
+          (getHookKind(fn.env, value.property.identifier) != null ||
+            isUseOperator(value.property.identifier))
         ) {
           hasReactiveInput = true;
         }

compiler/packages/babel-plugin-react-forget/src/ReactiveScopes/FlattenScopesWithHooksOrUse.ts

@@ -13,6 +13,7 @@ import {
   ReactiveStatement,
   ReactiveValue,
   getHookKind,
+  isUseOperator,
 } from "../HIR";
 import {
   ReactiveFunctionTransform,
@@ -20,18 +21,28 @@ import {
   visitReactiveFunction,
 } from "./visitors";
 
-/*
- * Most parts of compilation do not treat hooks specially, because there is no guarantee that custom
- * hooks obey any particular contract. For example, we can't assume that custom hooks won't modify
- * their arguments, and we can't assume that hooks return immutable or memoized values. Therefore
- * earlier passes largely ignore hooks, and may end up creating reactive scopes that contain hook calls.
+/**
+ * For simplicity the majority of compiler passes do not treat hooks specially. However, hooks are different
+ * from regular functions in two key ways:
+ * - They can introduce reactivity even when their arguments are non-reactive (accounted for in InferReactivePlaces)
+ * - They cannot be called conditionally
  *
- * This pass then finds and removes any scopes that transitively contain a hook call. By running all
+ * The `use` operator is similar:
+ * - It can access context, and therefore introduce reactivity
+ * - It can be called conditionally, but _it must be called if the component needs the return value_. This is because
+ *   React uses the fact that use was called to remember that the component needs the value, and that changes to the
+ *   input should invalidate the component itself.
+ *
+ * This pass accounts for the "can't call conditionally" aspect of both hooks and use. Though the reasoning is slightly
+ * different for reach, the result is that we can't memoize scopes that call hooks or use since this would make them
+ * called conditionally in the output.
+ *
+ * The pass finds and removes any scopes that transitively contain a hook or use call. By running all
  * the reactive scope inference first, agnostic of hooks, we know that the reactive scopes accurately
  * describe the set of values which "construct together", and remove _all_ that memoization in order
  * to ensure the hook call does not inadvertently become conditional.
  */
-export function flattenScopesWithHooks(fn: ReactiveFunction): void {
+export function flattenScopesWithHooksOrUse(fn: ReactiveFunction): void {
   visitReactiveFunction(fn, new Transform(), {
     env: fn.env,
     hasHook: false,
@@ -69,13 +80,19 @@ class Transform extends ReactiveFunctionTransform<State> {
     this.traverseValue(id, value, state);
     switch (value.kind) {
       case "CallExpression": {
-        if (getHookKind(state.env, value.callee.identifier) != null) {
+        if (
+          getHookKind(state.env, value.callee.identifier) != null ||
+          isUseOperator(value.callee.identifier)
+        ) {
           state.hasHook = true;
         }
         break;
       }
       case "MethodCall": {
-        if (getHookKind(state.env, value.property.identifier) != null) {
+        if (
+          getHookKind(state.env, value.property.identifier) != null ||
+          isUseOperator(value.property.identifier)
+        ) {
           state.hasHook = true;
         }
         break;

compiler/packages/babel-plugin-react-forget/src/ReactiveScopes/index.ts

@@ -16,7 +16,7 @@ export {
 } from "./CodegenReactiveFunction";
 export { extractScopeDeclarationsFromDestructuring } from "./ExtractScopeDeclarationsFromDestructuring";
 export { flattenReactiveLoops } from "./FlattenReactiveLoops";
-export { flattenScopesWithHooks } from "./FlattenScopesWithHooks";
+export { flattenScopesWithHooksOrUse } from "./FlattenScopesWithHooksOrUse";
 export { inferReactiveScopeVariables } from "./InferReactiveScopeVariables";
 export { memoizeFbtOperandsInSameScope } from "./MemoizeFbtOperandsInSameScope";
 export { mergeOverlappingReactiveScopes } from "./MergeOverlappingReactiveScopes";

compiler/packages/babel-plugin-react-forget/src/__tests__/fixtures/compiler/use-operator-call-expression.expect.md

@@ -0,0 +1,129 @@
+
+## Input
+
+```javascript
+import { ValidateMemoization } from "shared-runtime";
+import { use, useMemo } from "react";
+
+const FooContext = React.createContext(null);
+function Component(props) {
+  return (
+    <FooContext.Provider value={props.value}>
+      <Inner />
+    </FooContext.Provider>
+  );
+}
+
+function Inner(props) {
+  const input = use(FooContext);
+  const output = useMemo(() => [input], [input]);
+  return <ValidateMemoization inputs={[input]} output={output} />;
+}
+
+export const FIXTURE_ENTRYPOINT = {
+  fn: Component,
+  params: [{ value: 42 }],
+  sequentialRenders: [
+    { value: null },
+    { value: 42 },
+    { value: 42 },
+    { value: null },
+    { value: null },
+    { value: 42 },
+    { value: null },
+    { value: 42 },
+    { value: null },
+  ],
+};
+
+```
+
+## Code
+
+```javascript
+import { ValidateMemoization } from "shared-runtime";
+import { use, useMemo, unstable_useMemoCache as useMemoCache } from "react";
+
+const FooContext = React.createContext(null);
+function Component(props) {
+  const $ = useMemoCache(3);
+  let t0;
+  if ($[0] === Symbol.for("react.memo_cache_sentinel")) {
+    t0 = <Inner />;
+    $[0] = t0;
+  } else {
+    t0 = $[0];
+  }
+  let t1;
+  if ($[1] !== props.value) {
+    t1 = <FooContext.Provider value={props.value}>{t0}</FooContext.Provider>;
+    $[1] = props.value;
+    $[2] = t1;
+  } else {
+    t1 = $[2];
+  }
+  return t1;
+}
+
+function Inner(props) {
+  const $ = useMemoCache(7);
+  const input = use(FooContext);
+  let t0;
+  let t1;
+  if ($[0] !== input) {
+    t1 = [input];
+    $[0] = input;
+    $[1] = t1;
+  } else {
+    t1 = $[1];
+  }
+  t0 = t1;
+  const output = t0;
+  let t2;
+  if ($[2] !== input) {
+    t2 = [input];
+    $[2] = input;
+    $[3] = t2;
+  } else {
+    t2 = $[3];
+  }
+  let t3;
+  if ($[4] !== t2 || $[5] !== output) {
+    t3 = <ValidateMemoization inputs={t2} output={output} />;
+    $[4] = t2;
+    $[5] = output;
+    $[6] = t3;
+  } else {
+    t3 = $[6];
+  }
+  return t3;
+}
+
+export const FIXTURE_ENTRYPOINT = {
+  fn: Component,
+  params: [{ value: 42 }],
+  sequentialRenders: [
+    { value: null },
+    { value: 42 },
+    { value: 42 },
+    { value: null },
+    { value: null },
+    { value: 42 },
+    { value: null },
+    { value: 42 },
+    { value: null },
+  ],
+};
+
+```
+      
+### Eval output
+(kind: ok) <div>{"inputs":[null],"output":["[[ cyclic ref *2 ]]"]}</div>
+<div>{"inputs":[42],"output":[42]}</div>
+<div>{"inputs":[42],"output":[42]}</div>
+<div>{"inputs":[null],"output":["[[ cyclic ref *2 ]]"]}</div>
+<div>{"inputs":[null],"output":["[[ cyclic ref *2 ]]"]}</div>
+<div>{"inputs":[42],"output":[42]}</div>
+<div>{"inputs":[null],"output":["[[ cyclic ref *2 ]]"]}</div>
+<div>{"inputs":[42],"output":[42]}</div>
+<div>{"inputs":[null],"output":["[[ cyclic ref *2 ]]"]}</div>

compiler/packages/babel-plugin-react-forget/src/__tests__/fixtures/compiler/use-operator-call-expression.js

@@ -0,0 +1,33 @@
+import { ValidateMemoization } from "shared-runtime";
+import { use, useMemo } from "react";
+
+const FooContext = React.createContext(null);
+function Component(props) {
+  return (
+    <FooContext.Provider value={props.value}>
+      <Inner />
+    </FooContext.Provider>
+  );
+}
+
+function Inner(props) {
+  const input = use(FooContext);
+  const output = useMemo(() => [input], [input]);
+  return <ValidateMemoization inputs={[input]} output={output} />;
+}
+
+export const FIXTURE_ENTRYPOINT = {
+  fn: Component,
+  params: [{ value: 42 }],
+  sequentialRenders: [
+    { value: null },
+    { value: 42 },
+    { value: 42 },
+    { value: null },
+    { value: null },
+    { value: 42 },
+    { value: null },
+    { value: 42 },
+    { value: null },
+  ],
+};

compiler/packages/babel-plugin-react-forget/src/__tests__/fixtures/compiler/use-operator-method-call.expect.md

@@ -0,0 +1,131 @@
+
+## Input
+
+```javascript
+import { ValidateMemoization } from "shared-runtime";
+import { useMemo } from "react";
+import * as React from "react";
+
+const FooContext = React.createContext(null);
+function Component(props) {
+  return (
+    <FooContext.Provider value={props.value}>
+      <Inner />
+    </FooContext.Provider>
+  );
+}
+
+function Inner(props) {
+  const input = React.use(FooContext);
+  const output = useMemo(() => [input], [input]);
+  return <ValidateMemoization inputs={[input]} output={output} />;
+}
+
+export const FIXTURE_ENTRYPOINT = {
+  fn: Component,
+  params: [{ value: 42 }],
+  sequentialRenders: [
+    { value: null },
+    { value: 42 },
+    { value: 42 },
+    { value: null },
+    { value: null },
+    { value: 42 },
+    { value: null },
+    { value: 42 },
+    { value: null },
+  ],
+};
+
+```
+
+## Code
+
+```javascript
+import { ValidateMemoization } from "shared-runtime";
+import { useMemo, unstable_useMemoCache as useMemoCache } from "react";
+import * as React from "react";
+
+const FooContext = React.createContext(null);
+function Component(props) {
+  const $ = useMemoCache(3);
+  let t0;
+  if ($[0] === Symbol.for("react.memo_cache_sentinel")) {
+    t0 = <Inner />;
+    $[0] = t0;
+  } else {
+    t0 = $[0];
+  }
+  let t1;
+  if ($[1] !== props.value) {
+    t1 = <FooContext.Provider value={props.value}>{t0}</FooContext.Provider>;
+    $[1] = props.value;
+    $[2] = t1;
+  } else {
+    t1 = $[2];
+  }
+  return t1;
+}
+
+function Inner(props) {
+  const $ = useMemoCache(7);
+  const input = React.use(FooContext);
+  let t0;
+  let t1;
+  if ($[0] !== input) {
+    t1 = [input];
+    $[0] = input;
+    $[1] = t1;
+  } else {
+    t1 = $[1];
+  }
+  t0 = t1;
+  const output = t0;
+  let t2;
+  if ($[2] !== input) {
+    t2 = [input];
+    $[2] = input;
+    $[3] = t2;
+  } else {
+    t2 = $[3];
+  }
+  let t3;
+  if ($[4] !== t2 || $[5] !== output) {
+    t3 = <ValidateMemoization inputs={t2} output={output} />;
+    $[4] = t2;
+    $[5] = output;
+    $[6] = t3;
+  } else {
+    t3 = $[6];
+  }
+  return t3;
+}
+
+export const FIXTURE_ENTRYPOINT = {
+  fn: Component,
+  params: [{ value: 42 }],
+  sequentialRenders: [
+    { value: null },
+    { value: 42 },
+    { value: 42 },
+    { value: null },
+    { value: null },
+    { value: 42 },
+    { value: null },
+    { value: 42 },
+    { value: null },
+  ],
+};
+
+```
+      
+### Eval output
+(kind: ok) <div>{"inputs":[null],"output":["[[ cyclic ref *2 ]]"]}</div>
+<div>{"inputs":[42],"output":[42]}</div>
+<div>{"inputs":[42],"output":[42]}</div>
+<div>{"inputs":[null],"output":["[[ cyclic ref *2 ]]"]}</div>
+<div>{"inputs":[null],"output":["[[ cyclic ref *2 ]]"]}</div>
+<div>{"inputs":[42],"output":[42]}</div>
+<div>{"inputs":[null],"output":["[[ cyclic ref *2 ]]"]}</div>
+<div>{"inputs":[42],"output":[42]}</div>
+<div>{"inputs":[null],"output":["[[ cyclic ref *2 ]]"]}</div>

compiler/packages/babel-plugin-react-forget/src/__tests__/fixtures/compiler/use-operator-method-call.js

@@ -0,0 +1,34 @@
+import { ValidateMemoization } from "shared-runtime";
+import { useMemo } from "react";
+import * as React from "react";
+
+const FooContext = React.createContext(null);
+function Component(props) {
+  return (
+    <FooContext.Provider value={props.value}>
+      <Inner />
+    </FooContext.Provider>
+  );
+}
+
+function Inner(props) {
+  const input = React.use(FooContext);
+  const output = useMemo(() => [input], [input]);
+  return <ValidateMemoization inputs={[input]} output={output} />;
+}
+
+export const FIXTURE_ENTRYPOINT = {
+  fn: Component,
+  params: [{ value: 42 }],
+  sequentialRenders: [
+    { value: null },
+    { value: 42 },
+    { value: 42 },
+    { value: null },
+    { value: null },
+    { value: 42 },
+    { value: null },
+    { value: 42 },
+    { value: null },
+  ],
+};