[playground] Add support for "use no memo" (#31561)

Fixes #31331

## Summary
There is a bug in
playground(https://github.com/facebook/react/issues/31331) which doesnt
support 'use memo' or 'use no memo' directives. Its misleading while
debugging components in the playground

## How did you test this change?
Ran test cases and added a few extra test cases as well

## Changes
1) Adds support for 'use memo' and 'use no memo'
2) Cleanup E2E test cases a bit
3) Adds test cases for use memo
4) Added documentation to run test cases

## Implementation
`parseFunctions` returns a set of functions to be compiled. But, it
doesnt filter out/handle memoized opted/un-opted functions using
directives.

ive just created a `compile` flag to enable/disable compiling
[here](https://github.com/facebook/react/pull/31561/files#diff-305de47a3fe3ce778e22d5c5cf438419a59de8e7f785b45f659e7b41b1e30b03R113)

Then I am just skipping those functions from getting compile
[here](https://github.com/facebook/react/pull/31561/files#diff-305de47a3fe3ce778e22d5c5cf438419a59de8e7f785b45f659e7b41b1e30b03R253)

compiler/apps/playground/README.md

@@ -26,6 +26,13 @@ $ npm run dev
 $ yarn
 ```
 
+## Testing
+```sh
+# Install playwright browser binaries
+$ npx playwright install --with-deps
+# Run tests
+$ yarn test
+```
 ## Deployment
 
 This project has been deployed using Vercel. Vercel does the exact same thing as we would

compiler/apps/playground/__tests__/e2e/__snapshots__/page.spec.ts/function-scope-no-directive-arrow-function-expression-output.txt

@@ -0,0 +1,20 @@
+function anonymous_1() {
+  "use no memo";
+  const Avatar = () => {
+    return <div>Avatar Content</div>;
+  };
+  const MemoizedAvatar = React.memo(Avatar);
+  const Bio = () => {
+    const handleBioUpdate = () => {
+      console.log("Bio updated");
+    };
+    return <div onClick={handleBioUpdate}>Bio Content</div>;
+  };
+  const MemoizedBio = React.memo(Bio);
+  return (
+    <div>
+      <MemoizedAvatar />
+      <MemoizedBio />
+    </div>
+  );
+}

compiler/apps/playground/__tests__/e2e/__snapshots__/page.spec.ts/function-scope-use-memo-arrow-function-expression-output.txt

@@ -0,0 +1,32 @@
+function anonymous_1() {
+  "use memo";
+  const $ = _c(3);
+  const Chart = _temp2;
+  let t0;
+  if ($[0] === Symbol.for("react.memo_cache_sentinel")) {
+    t0 = React.memo(Chart);
+    $[0] = t0;
+  } else {
+    t0 = $[0];
+  }
+  const MemoizedChart = t0;
+  const Graph = _temp3;
+  let t1;
+  if ($[1] === Symbol.for("react.memo_cache_sentinel")) {
+    t1 = React.memo(Graph);
+    $[1] = t1;
+  } else {
+    t1 = $[1];
+  }
+  const MemoizedGraph = t1;
+  let t2;
+  if ($[2] === Symbol.for("react.memo_cache_sentinel")) {
+    t2 = (
+      <div>
+        <MemoizedChart />
+        <MemoizedGraph />
+      </div>
+    );
+    $[2] = t2;
+  } else {
+    t2 = $[2];

compiler/apps/playground/__tests__/e2e/__snapshots__/page.spec.ts/function-scope-use-memo-function-declaration-output.txt

@@ -0,0 +1,32 @@
+function App() {
+  "use memo";
+  const $ = _c(3);
+  let t0;
+  if ($[0] === Symbol.for("react.memo_cache_sentinel")) {
+    const Sidebar = function Sidebar() {
+      const handleToggle = _temp;
+      return <aside onClick={handleToggle}>Sidebar Content</
+          aside>;
+    };
+    t0 = React.memo(Sidebar);
+    $[0] = t0;
+  } else {
+    t0 = $[0];
+  }
+  const MemoizedSidebar = t0;
+  let t1;
+  if ($[1] === Symbol.for("react.memo_cache_sentinel")) {
+    const Content = function Content() {
+      return <main>Main Content</main>;
+    };
+    t1 = React.memo(Content);
+    $[1] = t1;
+  } else {
+    t1 = $[1];
+  }
+  const MemoizedContent = t1;
+  let t2;
+  if ($[2] === Symbol.for("react.memo_cache_sentinel")) {
+    t2 = (
+      <div>
+        <MemoizedSidebar />

compiler/apps/playground/__tests__/e2e/__snapshots__/page.spec.ts/function-scope-use-no-memo-function-declaration-output.txt

@@ -0,0 +1,32 @@
+function Settings() {
+  "use memo";
+  const $ = _c(3);
+  let t0;
+  if ($[0] === Symbol.for("react.memo_cache_sentinel")) {
+    t0 = function Preferences() {
+      const handleSave = _temp;
+      return <div onClick={handleSave}>Preferences Content</
+          div>;
+    };
+    $[0] = t0;
+  } else {
+    t0 = $[0];
+  }
+  const Preferences = t0;
+  let t1;
+  if ($[1] === Symbol.for("react.memo_cache_sentinel")) {
+    t1 = function Notifications() {
+      return <div>Notifications Settings</div>;
+    };
+    $[1] = t1;
+  } else {
+    t1 = $[1];
+  }
+  const Notifications = t1;
+  let t2;
+  if ($[2] === Symbol.for("react.memo_cache_sentinel")) {
+    t2 = (
+      <div>
+        <Preferences />
+        <Notifications />
+      </div>

compiler/apps/playground/__tests__/e2e/__snapshots__/page.spec.ts/function-scope-use-no-memo-function-expression-output.txt

@@ -0,0 +1,18 @@
+function anonymous_1() {
+  "use no memo";
+  const Widget = function () {
+    const handleExpand = () => {
+      console.log("Widget expanded");
+    };
+    return <div onClick={handleExpand}>Widget Content</div>;
+  };
+  const Panel = function () {
+    return <section>Panel Information</section>;
+  };
+  return (
+    <div>
+      <Widget />
+      <Panel />
+    </div>
+  );
+}

compiler/apps/playground/__tests__/e2e/__snapshots__/page.spec.ts/module-scope-use-memo-output.txt

@@ -0,0 +1,15 @@
+function anonymous_1() {
+  const $ = _c(1);
+  const handleClick = _temp;
+  let t0;
+  if ($[0] === Symbol.for("react.memo_cache_sentinel")) {
+    t0 = <h1 onClick={handleClick}>Welcome to the App!</h1>;
+    $[0] = t0;
+  } else {
+    t0 = $[0];
+  }
+  return t0;
+}
+function _temp() {
+  console.log("Header clicked");
+}

compiler/apps/playground/__tests__/e2e/__snapshots__/page.spec.ts/module-scope-use-no-memo-function-expression-output.txt

@@ -0,0 +1,3 @@
+function anonymous_1() {
+  return <aside>Sidebar Information</aside>;
+}

compiler/apps/playground/__tests__/e2e/__snapshots__/page.spec.ts/module-scope-use-no-memo-output.txt

@@ -0,0 +1,7 @@
+function anonymous_1() {
+  const handleMouseOver = () => {
+    console.log("Footer hovered");
+  };
+  return <footer onMouseOver={handleMouseOver}>Footer 
+      Information</footer>;
+}

compiler/apps/playground/__tests__/e2e/page.spec.ts

@@ -8,42 +8,245 @@
 import {expect, test} from '@playwright/test';
 import {encodeStore, type Store} from '../../lib/stores';
 
-const STORE: Store = {
+test.describe.configure({mode: 'parallel'});
-  source: `export default function TestComponent({ x }) {
-  return <Button>{x}</Button>;
-}
-`,
-};
-const HASH = encodeStore(STORE);
 
 function concat(data: Array<string>): string {
   return data.join('');
 }
+const DIRECTIVE_TEST_CASES = [
+  {
+    name: 'module-scope-use-memo',
+    input: `'use memo';
 
-test('editor should compile successfully', async ({page}) => {
+const Header = () => {
-  await page.goto(`/#${HASH}`, {waitUntil: 'networkidle'});
+  const handleClick = () => {
+    console.log('Header clicked');
+  };
+
+  return <h1 onClick={handleClick}>Welcome to the App!</h1>;
+};`,
+  },
+  {
+    name: 'module-scope-use-no-memo',
+    input: `'use no memo';
+
+const Footer = () => {
+  const handleMouseOver = () => {
+    console.log('Footer hovered');
+  };
+
+  return <footer onMouseOver={handleMouseOver}>Footer Information</footer>;
+};
+`,
+  },
+  {
+    name: 'function-scope-use-memo-function-declaration',
+    input: `function App() {
+  'use memo';
+  
+  function Sidebar() {
+    const handleToggle = () => {
+      console.log('Sidebar toggled');
+    };
+
+    return <aside onClick={handleToggle}>Sidebar Content</aside>;
+  }
+
+  const MemoizedSidebar = React.memo(Sidebar);
+
+  function Content() {
+    return <main>Main Content</main>;
+  }
+
+  const MemoizedContent = React.memo(Content);
+
+  return (
+    <div>
+      <MemoizedSidebar />
+      <MemoizedContent />
+    </div>
+  );
+}`,
+  },
+  {
+    name: 'function-scope-use-no-memo-function-expression',
+    input: `const Dashboard = function() {
+  'use no memo';
+  const Widget = function() {
+    const handleExpand = () => {
+      console.log('Widget expanded');
+    };
+
+    return <div onClick={handleExpand}>Widget Content</div>;
+  };
+
+  const Panel = function() {
+    return <section>Panel Information</section>;
+  };
+
+  return (
+    <div>
+      <Widget />
+      <Panel />
+    </div>
+  );
+};`,
+  },
+  {
+    name: 'function-scope-use-memo-arrow-function-expression',
+    input: `const Analytics = () => {
+  'use memo';
+
+  const Chart = () => {
+    const handleRefresh = () => {
+      console.log('Chart refreshed');
+    };
+
+    return <div onClick={handleRefresh}>Chart Content</div>;
+  };
+
+  const MemoizedChart = React.memo(Chart);
+
+  const Graph = () => {
+    return <div>Graph Content</div>;
+  };
+
+  const MemoizedGraph = React.memo(Graph);
+
+  return (
+    <div>
+      <MemoizedChart />
+      <MemoizedGraph />
+    </div>
+  );
+};`,
+  },
+  {
+    name: 'module-scope-use-no-memo-function-expression',
+    input: `'use no memo';
+
+const Sidebar = function() {
+  return <aside>Sidebar Information</aside>;
+};`,
+  },
+  {
+    name: 'function-scope-no-directive-arrow-function-expression',
+    input: `
+const Profile = () => {
+'use no memo';
+  const Avatar = () => {
+    return <div>Avatar Content</div>;
+  };
+
+  const MemoizedAvatar = React.memo(Avatar);
+
+  const Bio = () => {
+    const handleBioUpdate = () => {
+      console.log('Bio updated');
+    };
+
+    return <div onClick={handleBioUpdate}>Bio Content</div>;
+  };
+
+  const MemoizedBio = React.memo(Bio);
+
+  return (
+    <div>
+      <MemoizedAvatar />
+      <MemoizedBio />
+    </div>
+  );
+};`,
+  },
+  {
+    name: 'function-scope-use-no-memo-function-declaration',
+    input: `'use no memo';
+
+function Settings() {
+  'use memo';
+
+  function Preferences() {
+    const handleSave = () => {
+      console.log('Preferences saved');
+    };
+
+    return <div onClick={handleSave}>Preferences Content</div>;
+  }
+
+  function Notifications() {
+    return <div>Notifications Settings</div>;
+  }
+
+  return (
+    <div>
+      <Preferences />
+      <Notifications />
+    </div>
+  );
+}`,
+  },
+];
+test('editor should open successfully', async ({page}) => {
+  await page.goto(`/`, {waitUntil: 'networkidle'});
   await page.screenshot({
     fullPage: true,
-    path: 'test-results/00-on-networkidle.png',
+    path: 'test-results/00-fresh-page.png',
   });
+});
+test('editor should compile from hash successfully', async ({page}) => {
+  const store: Store = {
+    source: `export default function TestComponent({ x }) {
+      return <Button>{x}</Button>;
+    }
+    `,
+  };
+  const hash = encodeStore(store);
+  await page.goto(`/#${hash}`, {waitUntil: 'networkidle'});
 
   // User input from hash compiles
   await page.screenshot({
     fullPage: true,
-    path: 'test-results/01-show-js-before.png',
+    path: 'test-results/01-compiles-from-hash.png',
   });
   const userInput =
     (await page.locator('.monaco-editor').nth(1).allInnerTexts()) ?? [];
-  expect(concat(userInput)).toMatchSnapshot('user-input.txt');
+  expect(concat(userInput)).toMatchSnapshot('01-user-output.txt');
+});
+test('reset button works', async ({page}) => {
+  const store: Store = {
+    source: `export default function TestComponent({ x }) {
+      return <Button>{x}</Button>;
+    }
+    `,
+  };
+  const hash = encodeStore(store);
+  await page.goto(`/#${hash}`, {waitUntil: 'networkidle'});
 
   // Reset button works
   page.on('dialog', dialog => dialog.accept());
   await page.getByRole('button', {name: 'Reset'}).click();
   await page.screenshot({
     fullPage: true,
-    path: 'test-results/02-show-js-after.png',
+    path: 'test-results/02-reset-button-works.png',
   });
   const defaultInput =
     (await page.locator('.monaco-editor').nth(1).allInnerTexts()) ?? [];
-  expect(concat(defaultInput)).toMatchSnapshot('default-input.txt');
+  expect(concat(defaultInput)).toMatchSnapshot('02-default-output.txt');
 });
+DIRECTIVE_TEST_CASES.forEach((t, idx) =>
+  test(`directives work: ${t.name}`, async ({page}) => {
+    const store: Store = {
+      source: t.input,
+    };
+    const hash = encodeStore(store);
+    await page.goto(`/#${hash}`, {waitUntil: 'networkidle'});
+    await page.screenshot({
+      fullPage: true,
+      path: `test-results/03-0${idx}-${t.name}.png`,
+    });
+
+    const useMemoOutput =
+      (await page.locator('.monaco-editor').nth(1).allInnerTexts()) ?? [];
+    expect(concat(useMemoOutput)).toMatchSnapshot(`${t.name}-output.txt`);
+  }),
+);

compiler/apps/playground/components/Editor/EditorImpl.tsx

@@ -18,6 +18,8 @@ import {
   ValueKind,
   runPlayground,
   type Hook,
+  findDirectiveDisablingMemoization,
+  findDirectiveEnablingMemoization,
 } from 'babel-plugin-react-compiler/src';
 import {type ReactFunctionType} from 'babel-plugin-react-compiler/src/HIR/Environment';
 import clsx from 'clsx';
@@ -43,6 +45,25 @@ import {
 import {printFunctionWithOutlined} from 'babel-plugin-react-compiler/src/HIR/PrintHIR';
 import {printReactiveFunctionWithOutlined} from 'babel-plugin-react-compiler/src/ReactiveScopes/PrintReactiveFunction';
 
+type FunctionLike =
+  | NodePath<t.FunctionDeclaration>
+  | NodePath<t.ArrowFunctionExpression>
+  | NodePath<t.FunctionExpression>;
+enum MemoizeDirectiveState {
+  Enabled = 'Enabled',
+  Disabled = 'Disabled',
+  Undefined = 'Undefined',
+}
+
+const MEMOIZE_ENABLED_OR_UNDEFINED_STATES = new Set([
+  MemoizeDirectiveState.Enabled,
+  MemoizeDirectiveState.Undefined,
+]);
+
+const MEMOIZE_ENABLED_OR_DISABLED_STATES = new Set([
+  MemoizeDirectiveState.Enabled,
+  MemoizeDirectiveState.Disabled,
+]);
 function parseInput(input: string, language: 'flow' | 'typescript'): any {
   // Extract the first line to quickly check for custom test directives
   if (language === 'flow') {
@@ -63,29 +84,36 @@ function parseInput(input: string, language: 'flow' | 'typescript'): any {
 function parseFunctions(
   source: string,
   language: 'flow' | 'typescript',
-): Array<
+): Array<{
-  | NodePath<t.FunctionDeclaration>
+  compilationEnabled: boolean;
-  | NodePath<t.ArrowFunctionExpression>
+  fn: FunctionLike;
-  | NodePath<t.FunctionExpression>
+}> {
-> {
+  const items: Array<{
-  const items: Array<
+    compilationEnabled: boolean;
-    | NodePath<t.FunctionDeclaration>
+    fn: FunctionLike;
-    | NodePath<t.ArrowFunctionExpression>
+  }> = [];
-    | NodePath<t.FunctionExpression>
-  > = [];
   try {
     const ast = parseInput(source, language);
     traverse(ast, {
       FunctionDeclaration(nodePath) {
-        items.push(nodePath);
+        items.push({
+          compilationEnabled: shouldCompile(nodePath),
+          fn: nodePath,
+        });
         nodePath.skip();
       },
       ArrowFunctionExpression(nodePath) {
-        items.push(nodePath);
+        items.push({
+          compilationEnabled: shouldCompile(nodePath),
+          fn: nodePath,
+        });
         nodePath.skip();
       },
       FunctionExpression(nodePath) {
-        items.push(nodePath);
+        items.push({
+          compilationEnabled: shouldCompile(nodePath),
+          fn: nodePath,
+        });
         nodePath.skip();
       },
     });
@@ -98,9 +126,48 @@ function parseFunctions(
       suggestions: null,
     });
   }
+
   return items;
 }
 
+function shouldCompile(fn: FunctionLike): boolean {
+  const {body} = fn.node;
+  if (t.isBlockStatement(body)) {
+    const selfCheck = checkExplicitMemoizeDirectives(body.directives);
+    if (selfCheck === MemoizeDirectiveState.Enabled) return true;
+    if (selfCheck === MemoizeDirectiveState.Disabled) return false;
+
+    const parentWithDirective = fn.findParent(parentPath => {
+      if (parentPath.isBlockStatement() || parentPath.isProgram()) {
+        const directiveCheck = checkExplicitMemoizeDirectives(
+          parentPath.node.directives,
+        );
+        return MEMOIZE_ENABLED_OR_DISABLED_STATES.has(directiveCheck);
+      }
+      return false;
+    });
+
+    if (!parentWithDirective) return true;
+    const parentDirectiveCheck = checkExplicitMemoizeDirectives(
+      (parentWithDirective.node as t.Program | t.BlockStatement).directives,
+    );
+    return MEMOIZE_ENABLED_OR_UNDEFINED_STATES.has(parentDirectiveCheck);
+  }
+  return false;
+}
+
+function checkExplicitMemoizeDirectives(
+  directives: Array<t.Directive>,
+): MemoizeDirectiveState {
+  if (findDirectiveEnablingMemoization(directives).length) {
+    return MemoizeDirectiveState.Enabled;
+  }
+  if (findDirectiveDisablingMemoization(directives).length) {
+    return MemoizeDirectiveState.Disabled;
+  }
+  return MemoizeDirectiveState.Undefined;
+}
+
 const COMMON_HOOKS: Array<[string, Hook]> = [
   [
     'useFragment',
@@ -209,18 +276,38 @@ function compile(source: string): [CompilerOutput, 'flow' | 'typescript'] {
     // Extract the first line to quickly check for custom test directives
     const pragma = source.substring(0, source.indexOf('\n'));
     const config = parseConfigPragmaForTests(pragma);
-
+    const parsedFunctions = parseFunctions(source, language);
-    for (const fn of parseFunctions(source, language)) {
+    for (const func of parsedFunctions) {
-      const id = withIdentifier(getFunctionIdentifier(fn));
+      const id = withIdentifier(getFunctionIdentifier(func.fn));
+      const fnName = id.name;
+      if (!func.compilationEnabled) {
+        upsert({
+          kind: 'ast',
+          fnName,
+          name: 'CodeGen',
+          value: {
+            type: 'FunctionDeclaration',
+            id:
+              func.fn.isArrowFunctionExpression() ||
+              func.fn.isFunctionExpression()
+                ? withIdentifier(null)
+                : func.fn.node.id,
+            async: func.fn.node.async,
+            generator: !!func.fn.node.generator,
+            body: func.fn.node.body as t.BlockStatement,
+            params: func.fn.node.params,
+          },
+        });
+        continue;
+      }
       for (const result of runPlayground(
-        fn,
+        func.fn,
         {
           ...config,
           customHooks: new Map([...COMMON_HOOKS]),
         },
         getReactFunctionType(id),
       )) {
-        const fnName = id.name;
         switch (result.kind) {
           case 'ast': {
             upsert({

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

@@ -42,10 +42,10 @@ export type CompilerPass = {
   comments: Array<t.CommentBlock | t.CommentLine>;
   code: string | null;
 };
-const OPT_IN_DIRECTIVES = new Set(['use forget', 'use memo']);
+export const OPT_IN_DIRECTIVES = new Set(['use forget', 'use memo']);
 export const OPT_OUT_DIRECTIVES = new Set(['use no forget', 'use no memo']);
 
-function findDirectiveEnablingMemoization(
+export function findDirectiveEnablingMemoization(
   directives: Array<t.Directive>,
 ): Array<t.Directive> {
   return directives.filter(directive =>
@@ -53,7 +53,7 @@ function findDirectiveEnablingMemoization(
   );
 }
 
-function findDirectiveDisablingMemoization(
+export function findDirectiveDisablingMemoization(
   directives: Array<t.Directive>,
 ): Array<t.Directive> {
   return directives.filter(directive =>

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

@@ -20,6 +20,9 @@ export {
   run,
   runPlayground,
   OPT_OUT_DIRECTIVES,
+  OPT_IN_DIRECTIVES,
+  findDirectiveEnablingMemoization,
+  findDirectiveDisablingMemoization,
   type CompilerPipelineValue,
   type PluginOptions,
 } from './Entrypoint';