From 770ca4ab3ee173ab975d61d5022728f0cef5c67f Mon Sep 17 00:00:00 2001
From: Sathya Gunasekaran <gsathya.ceg@gmail.com>
Date: Mon, 19 Feb 2024 17:48:49 +0000
Subject: [PATCH] [Babel] Add support for "use memo"

We want to start moving away from "Forget", so this PR adds support "use memo"
and "use no memo"

I've left "use forget" and "use no forget" directives unchanged for now, as we
need to migrate existing users first and then come back and delete support for
these directives.
---
 .../src/Entrypoint/Program.ts                 | 23 +++++---
 .../compiler/use-memo-simple.expect.md        | 54 +++++++++++++++++++
 .../fixtures/compiler/use-memo-simple.js      | 11 ++++
 .../use-no-memo-module-level.expect.md        | 29 ++++++++++
 .../compiler/use-no-memo-module-level.js      |  8 +++
 .../compiler/use-no-memo-simple.expect.md     | 37 +++++++++++++
 .../fixtures/compiler/use-no-memo-simple.js   | 11 ++++
 .../packages/sprout/src/SproutTodoFilter.ts   |  1 +
 8 files changed, 167 insertions(+), 7 deletions(-)
 create mode 100644 compiler/packages/babel-plugin-react-forget/src/__tests__/fixtures/compiler/use-memo-simple.expect.md
 create mode 100644 compiler/packages/babel-plugin-react-forget/src/__tests__/fixtures/compiler/use-memo-simple.js
 create mode 100644 compiler/packages/babel-plugin-react-forget/src/__tests__/fixtures/compiler/use-no-memo-module-level.expect.md
 create mode 100644 compiler/packages/babel-plugin-react-forget/src/__tests__/fixtures/compiler/use-no-memo-module-level.js
 create mode 100644 compiler/packages/babel-plugin-react-forget/src/__tests__/fixtures/compiler/use-no-memo-simple.expect.md
 create mode 100644 compiler/packages/babel-plugin-react-forget/src/__tests__/fixtures/compiler/use-no-memo-simple.js

diff --git a/compiler/packages/babel-plugin-react-forget/src/Entrypoint/Program.ts b/compiler/packages/babel-plugin-react-forget/src/Entrypoint/Program.ts
index f2572dd77c..7071754a9c 100644
--- a/compiler/packages/babel-plugin-react-forget/src/Entrypoint/Program.ts
+++ b/compiler/packages/babel-plugin-react-forget/src/Entrypoint/Program.ts
@@ -36,20 +36,27 @@ export type CompilerPass = {
   comments: (t.CommentBlock | t.CommentLine)[];
 };
 
-function findUseForgetDirective(directives: t.Directive[]): t.Directive | null {
+function findDirectiveEnablingMemoization(
+  directives: t.Directive[]
+): t.Directive | null {
   for (const directive of directives) {
-    if (directive.value.value === "use forget") {
+    const directiveValue = directive.value.value;
+    if (directiveValue === "use forget" || directiveValue === "use memo") {
       return directive;
     }
   }
   return null;
 }
 
-function findUseNoForgetDirective(
+function findDirectiveDisablingMemoization(
   directives: t.Directive[]
 ): t.Directive | null {
   for (const directive of directives) {
-    if (directive.value.value === "use no forget") {
+    const directiveValue = directive.value.value;
+    if (
+      directiveValue === "use no forget" ||
+      directiveValue === "use no memo"
+    ) {
       return directive;
     }
   }
@@ -189,7 +196,7 @@ export function compileProgram(
   pass: CompilerPass
 ): void {
   // Top level "use no forget", skip this file entirely
-  if (findUseNoForgetDirective(program.node.directives) != null) {
+  if (findDirectiveDisablingMemoization(program.node.directives) != null) {
     return;
   }
 
@@ -372,7 +379,9 @@ export function shouldVisitNode(fn: BabelFn, pass: CompilerPass): boolean {
   }
   if (fn.node.body.type === "BlockStatement") {
     // Opt-outs disable compilation regardless of mode
-    const useNoForget = findUseNoForgetDirective(fn.node.body.directives);
+    const useNoForget = findDirectiveDisablingMemoization(
+      fn.node.body.directives
+    );
     if (useNoForget != null) {
       pass.opts.logger?.logEvent(pass.filename, {
         kind: "CompileError",
@@ -387,7 +396,7 @@ export function shouldVisitNode(fn: BabelFn, pass: CompilerPass): boolean {
       return false;
     }
     // Otherwise opt-ins enable compilation regardless of mode
-    if (findUseForgetDirective(fn.node.body.directives) != null) {
+    if (findDirectiveEnablingMemoization(fn.node.body.directives) != null) {
       return true;
     }
   }
diff --git a/compiler/packages/babel-plugin-react-forget/src/__tests__/fixtures/compiler/use-memo-simple.expect.md b/compiler/packages/babel-plugin-react-forget/src/__tests__/fixtures/compiler/use-memo-simple.expect.md
new file mode 100644
index 0000000000..d21466648d
--- /dev/null
+++ b/compiler/packages/babel-plugin-react-forget/src/__tests__/fixtures/compiler/use-memo-simple.expect.md
@@ -0,0 +1,54 @@
+
+## Input
+
+```javascript
+function Component(props) {
+  "use memo";
+  let x = [props.foo];
+  return <div x={x}>"foo"</div>;
+}
+
+export const FIXTURE_ENTRYPOINT = {
+  fn: Component,
+  params: [{ foo: 1 }],
+  isComponent: true,
+};
+
+```
+
+## Code
+
+```javascript
+import { unstable_useMemoCache as useMemoCache } from "react";
+function Component(props) {
+  const $ = useMemoCache(4);
+  let t0;
+  if ($[0] !== props.foo) {
+    t0 = [props.foo];
+    $[0] = props.foo;
+    $[1] = t0;
+  } else {
+    t0 = $[1];
+  }
+  const x = t0;
+  let t1;
+  if ($[2] !== x) {
+    t1 = <div x={x}>"foo"</div>;
+    $[2] = x;
+    $[3] = t1;
+  } else {
+    t1 = $[3];
+  }
+  return t1;
+}
+
+export const FIXTURE_ENTRYPOINT = {
+  fn: Component,
+  params: [{ foo: 1 }],
+  isComponent: true,
+};
+
+```
+      
+### Eval output
+(kind: ok) <div x="1">"foo"</div>
\ No newline at end of file
diff --git a/compiler/packages/babel-plugin-react-forget/src/__tests__/fixtures/compiler/use-memo-simple.js b/compiler/packages/babel-plugin-react-forget/src/__tests__/fixtures/compiler/use-memo-simple.js
new file mode 100644
index 0000000000..3166f45d95
--- /dev/null
+++ b/compiler/packages/babel-plugin-react-forget/src/__tests__/fixtures/compiler/use-memo-simple.js
@@ -0,0 +1,11 @@
+function Component(props) {
+  "use memo";
+  let x = [props.foo];
+  return <div x={x}>"foo"</div>;
+}
+
+export const FIXTURE_ENTRYPOINT = {
+  fn: Component,
+  params: [{ foo: 1 }],
+  isComponent: true,
+};
diff --git a/compiler/packages/babel-plugin-react-forget/src/__tests__/fixtures/compiler/use-no-memo-module-level.expect.md b/compiler/packages/babel-plugin-react-forget/src/__tests__/fixtures/compiler/use-no-memo-module-level.expect.md
new file mode 100644
index 0000000000..ef306c8291
--- /dev/null
+++ b/compiler/packages/babel-plugin-react-forget/src/__tests__/fixtures/compiler/use-no-memo-module-level.expect.md
@@ -0,0 +1,29 @@
+
+## Input
+
+```javascript
+"use no memo";
+
+export default function foo(x, y) {
+  if (x) {
+    return foo(false, y);
+  }
+  return [y * 10];
+}
+
+```
+
+## Code
+
+```javascript
+"use no memo";
+
+export default function foo(x, y) {
+  if (x) {
+    return foo(false, y);
+  }
+  return [y * 10];
+}
+
+```
+      
\ No newline at end of file
diff --git a/compiler/packages/babel-plugin-react-forget/src/__tests__/fixtures/compiler/use-no-memo-module-level.js b/compiler/packages/babel-plugin-react-forget/src/__tests__/fixtures/compiler/use-no-memo-module-level.js
new file mode 100644
index 0000000000..a7a740e8b9
--- /dev/null
+++ b/compiler/packages/babel-plugin-react-forget/src/__tests__/fixtures/compiler/use-no-memo-module-level.js
@@ -0,0 +1,8 @@
+"use no memo";
+
+export default function foo(x, y) {
+  if (x) {
+    return foo(false, y);
+  }
+  return [y * 10];
+}
diff --git a/compiler/packages/babel-plugin-react-forget/src/__tests__/fixtures/compiler/use-no-memo-simple.expect.md b/compiler/packages/babel-plugin-react-forget/src/__tests__/fixtures/compiler/use-no-memo-simple.expect.md
new file mode 100644
index 0000000000..f8ddf9af96
--- /dev/null
+++ b/compiler/packages/babel-plugin-react-forget/src/__tests__/fixtures/compiler/use-no-memo-simple.expect.md
@@ -0,0 +1,37 @@
+
+## Input
+
+```javascript
+function Component(props) {
+  "use no memo";
+  let x = [props.foo];
+  return <div x={x}>"foo"</div>;
+}
+
+export const FIXTURE_ENTRYPOINT = {
+  fn: Component,
+  params: [{ foo: 1 }],
+  isComponent: true,
+};
+
+```
+
+## Code
+
+```javascript
+function Component(props) {
+  "use no memo";
+  let x = [props.foo];
+  return <div x={x}>"foo"</div>;
+}
+
+export const FIXTURE_ENTRYPOINT = {
+  fn: Component,
+  params: [{ foo: 1 }],
+  isComponent: true,
+};
+
+```
+      
+### Eval output
+(kind: ok) <div x="1">"foo"</div>
\ No newline at end of file
diff --git a/compiler/packages/babel-plugin-react-forget/src/__tests__/fixtures/compiler/use-no-memo-simple.js b/compiler/packages/babel-plugin-react-forget/src/__tests__/fixtures/compiler/use-no-memo-simple.js
new file mode 100644
index 0000000000..2a525e275b
--- /dev/null
+++ b/compiler/packages/babel-plugin-react-forget/src/__tests__/fixtures/compiler/use-no-memo-simple.js
@@ -0,0 +1,11 @@
+function Component(props) {
+  "use no memo";
+  let x = [props.foo];
+  return <div x={x}>"foo"</div>;
+}
+
+export const FIXTURE_ENTRYPOINT = {
+  fn: Component,
+  params: [{ foo: 1 }],
+  isComponent: true,
+};
diff --git a/compiler/packages/sprout/src/SproutTodoFilter.ts b/compiler/packages/sprout/src/SproutTodoFilter.ts
index d7ce84f7dc..33ab9597ae 100644
--- a/compiler/packages/sprout/src/SproutTodoFilter.ts
+++ b/compiler/packages/sprout/src/SproutTodoFilter.ts
@@ -385,6 +385,7 @@ const skipFilter = new Set([
   "useMemo-return-empty",
   "useMemo-simple",
   "use-no-forget-module-level",
+  "use-no-memo-module-level",
   // defines multiple functions
   "alias-while",
   "babel-existing-react-import",