Repository: facebook/react
Pull request: [compiler] Reuse DropManualMemoization for ValidateNoVoidUseMemo


Much of the logic in the new validation pass is already implemented in DropManualMemoization, so let's combine them. I opted to keep the environment flag so we can more precisely control the rollout.

---
[//]: # (BEGIN SAPLING FOOTER)
Stack created with [Sapling](https://sapling-scm.com). Best reviewed with [ReviewStack](https://reviewstack.dev/facebook/react/pull/34001).
* #34022
* #34002
* __->__ #34001

Changed files:
filename: compiler/packages/babel-plugin-react-compiler/src/Entrypoint/Pipeline.ts
patch: @@ -82,7 +82,6 @@ import {
 import {inferTypes} from '../TypeInference';
 import {
   validateContextVariableLValues,
-  validateNoVoidUseMemo,
   validateHooksUsage,
   validateMemoizedEffectDependencies,
   validateNoCapitalizedCalls,
@@ -168,17 +167,14 @@ function runWithEnvironment(
 
   validateContextVariableLValues(hir);
   validateUseMemo(hir).unwrap();
-  if (env.config.validateNoVoidUseMemo) {
-    validateNoVoidUseMemo(hir).unwrap();
-  }
 
   if (
     env.isInferredMemoEnabled &&
     !env.config.enablePreserveExistingManualUseMemo &&
     !env.config.disableMemoizationForDebugging &&
     !env.config.enableChangeDetectionForDebugging
   ) {
-    dropManualMemoization(hir);
+    dropManualMemoization(hir).unwrap();
     log({kind: 'hir', name: 'DropManualMemoization', value: hir});
   }
 

filename: compiler/packages/babel-plugin-react-compiler/src/Inference/DropManualMemoization.ts
patch: @@ -5,7 +5,12 @@
  * LICENSE file in the root directory of this source tree.
  */
 
-import {CompilerError, SourceLocation} from '..';
+import {
+  CompilerDiagnostic,
+  CompilerError,
+  ErrorSeverity,
+  SourceLocation,
+} from '..';
 import {
   CallExpression,
   Effect,
@@ -30,6 +35,7 @@ import {
   makeInstructionId,
 } from '../HIR';
 import {createTemporaryPlace, markInstructionIds} from '../HIR/HIRBuilder';
+import {Result} from '../Utils/Result';
 
 type ManualMemoCallee = {
   kind: 'useMemo' | 'useCallback';
@@ -341,8 +347,14 @@ function extractManualMemoizationArgs(
  * rely on type inference to find useMemo/useCallback invocations, and instead does basic tracking
  * of globals and property loads to find both direct calls as well as usage via the React namespace,
  * eg `React.useMemo()`.
+ *
+ * This pass also validates that useMemo callbacks return a value (not void), ensuring that useMemo
+ * is only used for memoizing values and not for running arbitrary side effects.
  */
-export function dropManualMemoization(func: HIRFunction): void {
+export function dropManualMemoization(
+  func: HIRFunction,
+): Result<void, CompilerError> {
+  const errors = new CompilerError();
   const isValidationEnabled =
     func.env.config.validatePreserveExistingMemoizationGuarantees ||
     func.env.config.validateNoSetStateInRender ||
@@ -390,6 +402,41 @@ export function dropManualMemoization(func: HIRFunction): void {
             manualMemo.kind,
             sidemap,
           );
+
+          /**
+           * Bailout on void return useMemos. This is an anti-pattern where code might be using
+           * useMemo like useEffect: running arbirtary side-effects synced to changes in specific
+           * values.
+           */
+          if (
+            func.env.config.validateNoVoidUseMemo &&
+            manualMemo.kind === 'useMemo'
+          ) {
+            const funcToCheck = sidemap.functions.get(
+              fnPlace.identifier.id,
+            )?.value;
+            if (funcToCheck !== undefined && funcToCheck.loweredFunc.func) {
+              if (!hasNonVoidReturn(funcToCheck.loweredFunc.func)) {
+                errors.pushDiagnostic(
+                  CompilerDiagnostic.create({
+                    severity: ErrorSeverity.InvalidReact,
+                    category: 'useMemo() callbacks must return a value',
+                    description: `This ${
+                      manualMemo.loadInstr.value.kind === 'PropertyLoad'
+                        ? 'React.useMemo'
+                        : 'useMemo'
+                    } callback doesn't return a value. useMemo is for computing and caching values, not for arbitrary side effects.`,
+                    suggestions: null,
+                  }).withDetail({
+                    kind: 'error',
+                    loc: instr.value.loc,
+                    message: 'useMemo() callbacks must return a value',
+                  }),
+                );
+              }
+            }
+          }
+
           instr.value = getManualMemoizationReplacement(
             fnPlace,
             instr.value.loc,
@@ -486,6 +533,8 @@ export function dropManualMemoization(func: HIRFunction): void {
       markInstructionIds(func.body);
     }
   }
+
+  return errors.asResult();
 }
 
 function findOptionalPlaces(fn: HIRFunction): Set<IdentifierId> {
@@ -530,3 +579,17 @@ function findOptionalPlaces(fn: HIRFunction): Set<IdentifierId> {
   }
   return optionals;
 }
+
+function hasNonVoidReturn(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;
+}

filename: compiler/packages/babel-plugin-react-compiler/src/Validation/ValidateNoVoidUseMemo.ts
patch: @@ -1,156 +0,0 @@
-/**
- * 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;
-    }
-  }
-}

filename: compiler/packages/babel-plugin-react-compiler/src/Validation/index.ts
patch: @@ -6,7 +6,6 @@
  */
 
 export {validateContextVariableLValues} from './ValidateContextVariableLValues';
-export {validateNoVoidUseMemo} from './ValidateNoVoidUseMemo';
 export {validateHooksUsage} from './ValidateHooksUsage';
 export {validateMemoizedEffectDependencies} from './ValidateMemoizedEffectDependencies';
 export {validateNoCapitalizedCalls} from './ValidateNoCapitalizedCalls';

filename: compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/error.useMemo-no-return-value.expect.md
patch: @@ -24,18 +24,41 @@ function Component() {
 ## Error
 
 ```
-Found 1 error:
+Found 2 errors:
 
-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() callbacks must return a value
+
+This useMemo callback doesn't return a value. useMemo is for computing and caching values, not for 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 |   }, []);
+    |                 ^^^^^^^^^^^^^^^
+> 4 |     console.log('computing');
+    | ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+> 5 |   }, []);
+    | ^^^^^^^^^ useMemo() callbacks must return a value
   6 |   const value2 = React.useMemo(() => {
+  7 |     console.log('computing');
+  8 |   }, []);
+
+Error: useMemo() callbacks must return a value
+
+This React.useMemo callback doesn't return a value. useMemo is for computing and caching values, not for arbitrary side effects.
+
+error.useMemo-no-return-value.ts:6:17
+   4 |     console.log('computing');
+   5 |   }, []);
+>  6 |   const value2 = React.useMemo(() => {
+     |                  ^^^^^^^^^^^^^^^^^^^^^
+>  7 |     console.log('computing');
+     | ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+>  8 |   }, []);
+     | ^^^^^^^^^ useMemo() callbacks must return a value
+   9 |   return (
+  10 |     <div>
+  11 |       {value}
 ```
           
       
\ No newline at end of file