[compiler] Add new ValidateNoVoidUseMemo pass

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

@@ -82,6 +82,7 @@ import {
 import {inferTypes} from '../TypeInference';
 import {
   validateContextVariableLValues,
+  validateNoVoidUseMemo,
   validateHooksUsage,
   validateMemoizedEffectDependencies,
   validateNoCapitalizedCalls,
@@ -167,6 +168,9 @@ function runWithEnvironment(
 
   validateContextVariableLValues(hir);
   validateUseMemo(hir).unwrap();
+  if (env.config.validateNoVoidUseMemo) {
+    validateNoVoidUseMemo(hir).unwrap();
+  }
 
   if (
     env.isInferredMemoEnabled &&

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

@@ -189,6 +189,7 @@ export function lower(
     const fallthrough = builder.reserve('block');
     const terminal: ReturnTerminal = {
       kind: 'return',
+      returnVariant: 'Implicit',
       loc: GeneratedSource,
       value: lowerExpressionToTemporary(builder, body),
       id: makeInstructionId(0),
@@ -219,6 +220,7 @@ export function lower(
   builder.terminate(
     {
       kind: 'return',
+      returnVariant: 'Void',
       loc: GeneratedSource,
       value: lowerValueToTemporary(builder, {
         kind: 'Primitive',
@@ -302,6 +304,7 @@ function lowerStatement(
       }
       const terminal: ReturnTerminal = {
         kind: 'return',
+        returnVariant: 'Explicit',
         loc: stmt.node.loc ?? GeneratedSource,
         value,
         id: makeInstructionId(0),

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

@@ -631,6 +631,17 @@ export const EnvironmentConfigSchema = z.object({
    * ```
    */
   lowerContextAccess: ExternalFunctionSchema.nullable().default(null),
+
+  /**
+   * If enabled, will validate useMemos that don't return any values:
+   *
+   * Valid:
+   *   useMemo(() => foo, [foo]);
+   *   useMemo(() => { return foo }, [foo]);
+   * Invalid:
+   *   useMemo(() => { ... }, [...]);
+   */
+  validateNoVoidUseMemo: z.boolean().default(false),
 });
 
 export type EnvironmentConfig = z.infer<typeof EnvironmentConfigSchema>;

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

@@ -446,8 +446,20 @@ export type ThrowTerminal = {
 };
 export type Case = {test: Place | null; block: BlockId};
 
+export type ReturnVariant = 'Void' | 'Implicit' | 'Explicit';
 export type ReturnTerminal = {
   kind: 'return';
+  /**
+   * Void:
+   *   () => { ... }
+   *   function() { ... }
+   * Implicit (ArrowFunctionExpression only):
+   *   () => foo
+   * Explicit:
+   *   () => { return ... }
+   *   function () { return ... }
+   */
+  returnVariant: ReturnVariant;
   loc: SourceLocation;
   value: Place;
   id: InstructionId;

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

@@ -215,7 +215,7 @@ export function printTerminal(terminal: Terminal): Array<string> | string {
       break;
     }
     case 'return': {
-      value = `[${terminal.id}] Return${
+      value = `[${terminal.id}] Return ${terminal.returnVariant}${
         terminal.value != null ? ' ' + printPlace(terminal.value) : ''
       }`;
       if (terminal.effects != null) {

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

@@ -777,6 +777,7 @@ export function mapTerminalSuccessors(
     case 'return': {
       return {
         kind: 'return',
+        returnVariant: terminal.returnVariant,
         loc: terminal.loc,
         value: terminal.value,
         id: makeInstructionId(0),

compiler/packages/babel-plugin-react-compiler/src/Optimization/LowerContextAccess.ts

@@ -237,6 +237,7 @@ function emitSelectorFn(env: Environment, keys: Array<string>): Instruction {
     terminal: {
       id: makeInstructionId(0),
       kind: 'return',
+      returnVariant: 'Explicit',
       loc: GeneratedSource,
       value: arrayInstr.lvalue,
       effects: null,

compiler/packages/babel-plugin-react-compiler/src/Optimization/OutlineJsx.ts

@@ -352,6 +352,7 @@ function emitOutlinedFn(
     terminal: {
       id: makeInstructionId(0),
       kind: 'return',
+      returnVariant: 'Explicit',
       loc: GeneratedSource,
       value: instructions.at(-1)!.lvalue,
       effects: null,

compiler/packages/babel-plugin-react-compiler/src/Validation/ValidateNoVoidUseMemo.ts

@@ -0,0 +1,156 @@
+/**
+ * Copyright (c) Meta Platforms, Inc. and affiliates.
+ *
+ * This source code is licensed under the MIT license found in the
+ * LICENSE file in the root directory of this source tree.
+ */
+
+import {CompilerError, ErrorSeverity} from '../CompilerError';
+import {
+  HIRFunction,
+  IdentifierId,
+  FunctionExpression,
+  SourceLocation,
+  Environment,
+  Instruction,
+  getHookKindForType,
+} from '../HIR';
+import {Result} from '../Utils/Result';
+
+type TemporariesSidemap = {
+  useMemoHooks: Map<IdentifierId, {name: string; loc: SourceLocation}>;
+  funcExprs: Map<IdentifierId, FunctionExpression>;
+  react: Set<IdentifierId>;
+};
+
+/**
+ * Validates that useMemo has at least one explicit return statement.
+ *
+ * Valid cases:
+ * - useMemo(() => value)                      // implicit arrow function return
+ * - useMemo(() => { return value; })          // explicit return
+ * - useMemo(() => { return; })                // explicit undefined
+ * - useMemo(() => { if (cond) return val; })  // at least one return
+ *
+ * Invalid cases:
+ * - useMemo(() => { console.log(); })         // no return statement at all
+ */
+export function validateNoVoidUseMemo(
+  fn: HIRFunction,
+): Result<void, CompilerError> {
+  const errors = new CompilerError();
+  const sidemap: TemporariesSidemap = {
+    useMemoHooks: new Map(),
+    funcExprs: new Map(),
+    react: new Set(),
+  };
+
+  for (const [, block] of fn.body.blocks) {
+    for (const instr of block.instructions) {
+      collectTemporaries(instr, fn.env, sidemap);
+    }
+  }
+
+  for (const [, block] of fn.body.blocks) {
+    for (const instr of block.instructions) {
+      if (instr.value.kind === 'CallExpression') {
+        const callee = instr.value.callee.identifier;
+        const useMemoHook = sidemap.useMemoHooks.get(callee.id);
+
+        if (useMemoHook !== undefined && instr.value.args.length > 0) {
+          const firstArg = instr.value.args[0];
+          if (firstArg.kind !== 'Identifier') {
+            continue;
+          }
+
+          let funcToCheck = sidemap.funcExprs.get(firstArg.identifier.id);
+
+          if (!funcToCheck) {
+            for (const [, searchBlock] of fn.body.blocks) {
+              for (const searchInstr of searchBlock.instructions) {
+                if (
+                  searchInstr.lvalue &&
+                  searchInstr.lvalue.identifier.id === firstArg.identifier.id &&
+                  searchInstr.value.kind === 'FunctionExpression'
+                ) {
+                  funcToCheck = searchInstr.value;
+                  break;
+                }
+              }
+              if (funcToCheck) break;
+            }
+          }
+
+          if (funcToCheck) {
+            const hasReturn = checkFunctionHasNonVoidReturn(
+              funcToCheck.loweredFunc.func,
+            );
+
+            if (!hasReturn) {
+              errors.push({
+                severity: ErrorSeverity.InvalidReact,
+                reason: `React Compiler has skipped optimizing this component because ${useMemoHook.name} doesn't return a value. ${useMemoHook.name} should only be used for memoizing values, not running arbitrary side effects.`,
+                loc: useMemoHook.loc,
+                suggestions: null,
+                description: null,
+              });
+            }
+          }
+        }
+      }
+    }
+  }
+  return errors.asResult();
+}
+
+function checkFunctionHasNonVoidReturn(func: HIRFunction): boolean {
+  for (const [, block] of func.body.blocks) {
+    if (block.terminal.kind === 'return') {
+      if (
+        block.terminal.returnVariant === 'Explicit' ||
+        block.terminal.returnVariant === 'Implicit'
+      ) {
+        return true;
+      }
+    }
+  }
+  return false;
+}
+
+function collectTemporaries(
+  instr: Instruction,
+  env: Environment,
+  sidemap: TemporariesSidemap,
+): void {
+  const {value, lvalue} = instr;
+  switch (value.kind) {
+    case 'FunctionExpression': {
+      sidemap.funcExprs.set(lvalue.identifier.id, value);
+      break;
+    }
+    case 'LoadGlobal': {
+      const global = env.getGlobalDeclaration(value.binding, value.loc);
+      const hookKind = global !== null ? getHookKindForType(env, global) : null;
+      if (hookKind === 'useMemo') {
+        sidemap.useMemoHooks.set(lvalue.identifier.id, {
+          name: value.binding.name,
+          loc: instr.loc,
+        });
+      } else if (value.binding.name === 'React') {
+        sidemap.react.add(lvalue.identifier.id);
+      }
+      break;
+    }
+    case 'PropertyLoad': {
+      if (sidemap.react.has(value.object.identifier.id)) {
+        if (value.property === 'useMemo') {
+          sidemap.useMemoHooks.set(lvalue.identifier.id, {
+            name: value.property,
+            loc: instr.loc,
+          });
+        }
+      }
+      break;
+    }
+  }
+}

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

@@ -6,6 +6,7 @@
  */
 
 export {validateContextVariableLValues} from './ValidateContextVariableLValues';
+export {validateNoVoidUseMemo} from './ValidateNoVoidUseMemo';
 export {validateHooksUsage} from './ValidateHooksUsage';
 export {validateMemoizedEffectDependencies} from './ValidateMemoizedEffectDependencies';
 export {validateNoCapitalizedCalls} from './ValidateNoCapitalizedCalls';

compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/error.useMemo-no-return-value.expect.md

@@ -0,0 +1,41 @@
+
+## Input
+
+```javascript
+// @validateNoVoidUseMemo
+function Component() {
+  const value = useMemo(() => {
+    console.log('computing');
+  }, []);
+  const value2 = React.useMemo(() => {
+    console.log('computing');
+  }, []);
+  return (
+    <div>
+      {value}
+      {value2}
+    </div>
+  );
+}
+
+```
+
+
+## Error
+
+```
+Found 1 error:
+
+Error: React Compiler has skipped optimizing this component because useMemo doesn't return a value. useMemo should only be used for memoizing values, not running arbitrary side effects.
+
+error.useMemo-no-return-value.ts:3:16
+  1 | // @validateNoVoidUseMemo
+  2 | function Component() {
+> 3 |   const value = useMemo(() => {
+    |                 ^^^^^^^ React Compiler has skipped optimizing this component because useMemo doesn't return a value. useMemo should only be used for memoizing values, not running arbitrary side effects.
+  4 |     console.log('computing');
+  5 |   }, []);
+  6 |   const value2 = React.useMemo(() => {
+```
+
+

compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/error.useMemo-no-return-value.js

@@ -0,0 +1,15 @@
+// @validateNoVoidUseMemo
+function Component() {
+  const value = useMemo(() => {
+    console.log('computing');
+  }, []);
+  const value2 = React.useMemo(() => {
+    console.log('computing');
+  }, []);
+  return (
+    <div>
+      {value}
+      {value2}
+    </div>
+  );
+}

compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/useMemo-arrow-implicit-return.expect.md

@@ -0,0 +1,40 @@
+
+## Input
+
+```javascript
+// @validateNoVoidUseMemo
+function Component() {
+  const value = useMemo(() => computeValue(), []);
+  return <div>{value}</div>;
+}
+
+```
+
+## Code
+
+```javascript
+import { c as _c } from "react/compiler-runtime"; // @validateNoVoidUseMemo
+function Component() {
+  const $ = _c(2);
+  let t0;
+  if ($[0] === Symbol.for("react.memo_cache_sentinel")) {
+    t0 = computeValue();
+    $[0] = t0;
+  } else {
+    t0 = $[0];
+  }
+  const value = t0;
+  let t1;
+  if ($[1] === Symbol.for("react.memo_cache_sentinel")) {
+    t1 = <div>{value}</div>;
+    $[1] = t1;
+  } else {
+    t1 = $[1];
+  }
+  return t1;
+}
+
+```
+
+### Eval output
+(kind: exception) Fixture not implemented

compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/useMemo-arrow-implicit-return.js

@@ -0,0 +1,5 @@
+// @validateNoVoidUseMemo
+function Component() {
+  const value = useMemo(() => computeValue(), []);
+  return <div>{value}</div>;
+}