[playground] Fixes #31331

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/simple-use-memo-output.txt

@@ -0,0 +1,19 @@
+function useFoo(props) {
+  "use no memo";
+  const x = () => {};
+  const y = function (a) {};
+  return foo(props.x, x);
+}
+function Component() {
+  const $ = _c(2);
+  const x = useFoo();
+  let t0;
+  if ($[0] !== x) {
+    t0 = <div>{x}</div>;
+    $[0] = x;
+    $[1] = t0;
+  } else {
+    t0 = $[1];
+  }
+  return t0;
+}

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

@@ -8,42 +8,79 @@
 import {expect, test} from '@playwright/test';
 import {encodeStore, type Store} from '../../lib/stores';
 
-const STORE: Store = {
-  source: `export default function TestComponent({ x }) {
-  return <Button>{x}</Button>;
-}
-`,
-};
-const HASH = encodeStore(STORE);
-
 function concat(data: Array<string>): string {
   return data.join('');
 }
 
-test('editor should compile successfully', async ({page}) => {
-  await page.goto(`/#${HASH}`, {waitUntil: 'networkidle'});
+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('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('default-output.txt');
+});
+test('directives work', async ({page}) => {
+  const store: Store = {
+    source: `function useFoo(props) {
+              'use no memo';
+              const x = () => { };
+              const y = function (a) { };
+              return foo(props.x, x);
+            }
+            function Component() {
+              const x = useFoo();
+              return <div>{x}</div>;
+            }
+    `,
+  };
+  const hash = encodeStore(store);
+  await page.goto(`/#${hash}`, {waitUntil: 'networkidle'});
+  await page.screenshot({
+    fullPage: true,
+    path: 'test-results/03-simple-use-memo.png',
+  });
+
+  const useMemoOutput =
+    (await page.locator('.monaco-editor').nth(1).allInnerTexts()) ?? [];
+  expect(concat(useMemoOutput)).toMatchSnapshot('simple-use-memo-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<
-  | NodePath<t.FunctionDeclaration>
-  | NodePath<t.ArrowFunctionExpression>
-  | NodePath<t.FunctionExpression>
-> {
-  const items: Array<
-    | NodePath<t.FunctionDeclaration>
-    | NodePath<t.ArrowFunctionExpression>
-    | NodePath<t.FunctionExpression>
-  > = [];
+): Array<{
+  compilationEnabled: boolean;
+  fn: FunctionLike;
+}> {
+  const items: Array<{
+    compilationEnabled: boolean;
+    fn: FunctionLike;
+  }> = [];
   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,34 @@ 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);
-
-    for (const fn of parseFunctions(source, language)) {
-      const id = withIdentifier(getFunctionIdentifier(fn));
+    const parsedFunctions = parseFunctions(source, language);
+    for (const func of parsedFunctions) {
+      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() ? 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,18 +42,18 @@ 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 =>
     OPT_IN_DIRECTIVES.has(directive.value.value),
   );
 }
 
-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';