[babel] Allow userspace useMemoCache

Not every consumer of Forget will be able to run an experimental version of 
React. In the meantime before useMemoCache is stable, provide a way for OSS to 
pass in a userspace impl.

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

@@ -145,3 +145,21 @@ export function insertUseMemoCacheImportDeclaration(
     )
   );
 }
+
+export function insertUserspaceUseMemoCacheImportDeclaration(
+  program: NodePath<t.Program>,
+  moduleName: string
+): void {
+  program.unshiftContainer(
+    "body",
+    t.importDeclaration(
+      [
+        t.importSpecifier(
+          t.identifier("useMemoCache"),
+          t.identifier("unstable_useMemoCache")
+        ),
+      ],
+      t.stringLiteral(moduleName)
+    )
+  );
+}

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

@@ -109,6 +109,17 @@ export type PluginOptions = {
    * ```
    */
   compilationMode: CompilationMode;
+
+  /**
+   * If specified, Forget will import `useMemoCache` from this module instead of React. Use this if
+   * you are for whatever reason unable to use an experimental version of React.
+   *
+   * ```
+   * // If specified:
+   * import {unstable_useMemoCache} from 'useMemoCache_DO_NOT_USE';
+   * ```
+   */
+  useMemoCacheSource: string | null;
 };
 
 export type CompilationMode =
@@ -156,6 +167,7 @@ export const defaultOptions: PluginOptions = {
   gating: null,
   instrumentForget: null,
   noEmit: false,
+  useMemoCacheSource: null,
 } as const;
 
 export function parsePluginOptions(obj: unknown): PluginOptions {

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

@@ -21,10 +21,11 @@ import {
   addImportsToProgram,
   findExistingImports,
   insertUseMemoCacheImportDeclaration,
+  insertUserspaceUseMemoCacheImportDeclaration,
   updateExistingReactImportDeclaration,
 } from "./Imports";
 import { addInstrumentForget } from "./Instrumentation";
-import { PluginOptions, parsePluginOptions } from "./Options";
+import { ExternalFunction, PluginOptions, parsePluginOptions } from "./Options";
 import { compileFn } from "./Pipeline";
 
 export type CompilerPass = {
@@ -325,18 +326,25 @@ export function compileProgram(
     // `import {unstable_useMemoCache as useMemoCache} from 'react'` and rename
     // `React.unstable_useMemoCache(n)` to `useMemoCache(n)`;
     if (didInsertUseMemoCache) {
-      if (hasExistingReactImport) {
-        const didUpdateImport = updateExistingReactImportDeclaration(program);
-        if (didUpdateImport === false) {
-          throw new Error(
-            "Expected an ImportDeclaration of react in order to update ImportSpecifiers with useMemoCache"
-          );
+      if (options.useMemoCacheSource === null) {
+        if (hasExistingReactImport) {
+          const didUpdateImport = updateExistingReactImportDeclaration(program);
+          if (didUpdateImport === false) {
+            throw new Error(
+              "Expected an ImportDeclaration of react in order to update ImportSpecifiers with useMemoCache"
+            );
+          }
+        } else {
+          insertUseMemoCacheImportDeclaration(program);
         }
-      } else {
-        insertUseMemoCacheImportDeclaration(program);
+      } else if (typeof options.useMemoCacheSource === "string") {
+        insertUserspaceUseMemoCacheImportDeclaration(
+          program,
+          options.useMemoCacheSource
+        );
       }
     }
-    const externalFunctions = [];
+    const externalFunctions: ExternalFunction[] = [];
     // TODO: check for duplicate import specifiers
     if (options.gating != null) {
       externalFunctions.push(options.gating);

compiler/packages/babel-plugin-react-forget/src/__tests__/fixtures/compiler/userspace-use-memo-cache.expect.md

@@ -0,0 +1,73 @@
+
+## Input
+
+```javascript
+// @useMemoCacheSource
+function Component(props) {
+  const [x, setX] = useState(1);
+  let y;
+  if (props.cond) {
+    y = x * 2;
+  }
+  return (
+    <Button
+      onClick={() => {
+        setX(10 * y);
+      }}
+    ></Button>
+  );
+}
+
+export const FIXTURE_ENTRYPOINT = {
+  fn: Component,
+  params: [true],
+  isComponent: true,
+};
+
+```
+
+## Code
+
+```javascript
+import { unstable_useMemoCache as useMemoCache } from "shared-runtime"; // @useMemoCacheSource
+function Component(props) {
+  const $ = useMemoCache(5);
+  const [x, setX] = useState(1);
+  let y;
+  if ($[0] !== props.cond || $[1] !== x) {
+    if (props.cond) {
+      y = x * 2;
+    }
+    $[0] = props.cond;
+    $[1] = x;
+    $[2] = y;
+  } else {
+    y = $[2];
+  }
+
+  const t0 = y;
+  let t1;
+  if ($[3] !== t0) {
+    t1 = (
+      <Button
+        onClick={() => {
+          setX(10 * y);
+        }}
+      />
+    );
+    $[3] = t0;
+    $[4] = t1;
+  } else {
+    t1 = $[4];
+  }
+  return t1;
+}
+
+export const FIXTURE_ENTRYPOINT = {
+  fn: Component,
+  params: [true],
+  isComponent: true,
+};
+
+```
+      

compiler/packages/babel-plugin-react-forget/src/__tests__/fixtures/compiler/userspace-use-memo-cache.js

@@ -0,0 +1,21 @@
+// @useMemoCacheSource
+function Component(props) {
+  const [x, setX] = useState(1);
+  let y;
+  if (props.cond) {
+    y = x * 2;
+  }
+  return (
+    <Button
+      onClick={() => {
+        setX(10 * y);
+      }}
+    ></Button>
+  );
+}
+
+export const FIXTURE_ENTRYPOINT = {
+  fn: Component,
+  params: [true],
+  isComponent: true,
+};

compiler/packages/fixture-test-utils/src/compiler-utils.ts

@@ -24,6 +24,7 @@ export function transformFixtureInput(
   let instrumentForget = null;
   let enableEmitFreeze = null;
   let compilationMode: CompilationMode = "all";
+  let useMemoCacheSource = null;
 
   if (firstLine.indexOf("@compilationMode(annotation)") !== -1) {
     assert(
@@ -58,6 +59,9 @@ export function transformFixtureInput(
       importSpecifierName: "makeReadOnly",
     };
   }
+  if (firstLine.includes("@useMemoCacheSource")) {
+    useMemoCacheSource = "shared-runtime";
+  }
   const config = parseConfigPragmaFn(firstLine);
   const result = pluginFn(
     input,
@@ -104,6 +108,7 @@ export function transformFixtureInput(
       instrumentForget,
       panicThreshold: "ALL_ERRORS",
       noEmit: false,
+      useMemoCacheSource,
     },
     includeAst
   );

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

@@ -191,3 +191,13 @@ export const ObjectWithHooks = {
     return 0;
   },
 };
+
+const $empty = Symbol.for("react.memo_cache_sentinel");
+export function unstable_useMemoCache(size: number) {
+  "use no forget";
+  const $ = new Array(size);
+  for (let ii = 0; ii < size; ii++) {
+    $[ii] = $empty;
+  }
+  return React.useRef($).current;
+}