[compiler] Remove useState argument constraint. no-derived-computations-in-effects

Summary:
I missed this conditional messing things up for undefined useState() calls. We should be tracking them.

I also missed a test that expect an error was not throwing.

Test Plan:
Update broken test

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

@@ -105,7 +105,6 @@ import {inferMutationAliasingRanges} from '../Inference/InferMutationAliasingRan
 import {validateNoDerivedComputationsInEffects} from '../Validation/ValidateNoDerivedComputationsInEffects';
 import {validateNoDerivedComputationsInEffects_exp} from '../Validation/ValidateNoDerivedComputationsInEffects_exp';
 import {nameAnonymousFunctions} from '../Transform/NameAnonymousFunctions';
-import {optimizeForSSR} from '../Optimization/OptimizeForSSR';
 import {validateSourceLocations} from '../Validation/ValidateSourceLocations';
 
 export type CompilerPipelineValue =
@@ -238,11 +237,6 @@ function runWithEnvironment(
     }
   }
 
-  if (env.config.enableOptimizeForSSR) {
-    optimizeForSSR(hir);
-    log({kind: 'hir', name: 'OptimizeForSSR', value: hir});
-  }
-
   // Note: Has to come after infer reference effects because "dead" code may still affect inference
   deadCodeElimination(hir);
   log({kind: 'hir', name: 'DeadCodeElimination', value: hir});
@@ -320,10 +314,8 @@ function runWithEnvironment(
      * if inferred memoization is enabled. This makes all later passes which
      * transform reactive-scope labeled instructions no-ops.
      */
-    if (!env.config.enableOptimizeForSSR) {
-      inferReactiveScopeVariables(hir);
-      log({kind: 'hir', name: 'InferReactiveScopeVariables', value: hir});
-    }
+    inferReactiveScopeVariables(hir);
+    log({kind: 'hir', name: 'InferReactiveScopeVariables', value: hir});
   }
 
   const fbtOperands = memoizeFbtAndMacroOperandsInSameScope(hir);

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

@@ -691,8 +691,6 @@ export const EnvironmentConfigSchema = z.object({
    * by React to only execute in response to events, not during render.
    */
   enableInferEventHandlers: z.boolean().default(false),
-
-  enableOptimizeForSSR: z.boolean().default(false),
 });
 
 export type EnvironmentConfig = z.infer<typeof EnvironmentConfigSchema>;

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

@@ -1823,10 +1823,6 @@ export function isPrimitiveType(id: Identifier): boolean {
   return id.type.kind === 'Primitive';
 }
 
-export function isPlainObjectType(id: Identifier): boolean {
-  return id.type.kind === 'Object' && id.type.shapeId === 'BuiltInObject';
-}
-
 export function isArrayType(id: Identifier): boolean {
   return id.type.kind === 'Object' && id.type.shapeId === 'BuiltInArray';
 }

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

@@ -7,8 +7,6 @@
 
 import {
   BlockId,
-  Environment,
-  getHookKind,
   HIRFunction,
   Identifier,
   IdentifierId,
@@ -70,14 +68,9 @@ export function deadCodeElimination(fn: HIRFunction): void {
 }
 
 class State {
-  env: Environment;
   named: Set<string> = new Set();
   identifiers: Set<IdentifierId> = new Set();
 
-  constructor(env: Environment) {
-    this.env = env;
-  }
-
   // Mark the identifier as being referenced (not dead code)
   reference(identifier: Identifier): void {
     this.identifiers.add(identifier.id);
@@ -119,7 +112,7 @@ function findReferencedIdentifiers(fn: HIRFunction): State {
   const hasLoop = hasBackEdge(fn);
   const reversedBlocks = [...fn.body.blocks.values()].reverse();
 
-  const state = new State(fn.env);
+  const state = new State();
   let size = state.count;
   do {
     size = state.count;
@@ -317,27 +310,12 @@ function pruneableValue(value: InstructionValue, state: State): boolean {
       // explicitly retain debugger statements to not break debugging workflows
       return false;
     }
-    case 'CallExpression':
-    case 'MethodCall': {
-      if (state.env.config.enableOptimizeForSSR) {
-        const calleee =
-          value.kind === 'CallExpression' ? value.callee : value.property;
-        const hookKind = getHookKind(state.env, calleee.identifier);
-        switch (hookKind) {
-          case 'useState':
-          case 'useReducer':
-          case 'useRef': {
-            // unused refs can be removed
-            return true;
-          }
-        }
-      }
-      return false;
-    }
     case 'Await':
+    case 'CallExpression':
     case 'ComputedDelete':
     case 'ComputedStore':
     case 'PropertyDelete':
+    case 'MethodCall':
     case 'PropertyStore':
     case 'StoreGlobal': {
       /*

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

@@ -1,269 +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} from '..';
-import {
-  CallExpression,
-  getHookKind,
-  HIRFunction,
-  IdentifierId,
-  InstructionValue,
-  isArrayType,
-  isPlainObjectType,
-  isPrimitiveType,
-  isSetStateType,
-  isStartTransitionType,
-  LoadLocal,
-  StoreLocal,
-} from '../HIR';
-import {
-  eachInstructionValueOperand,
-  eachTerminalOperand,
-} from '../HIR/visitors';
-import {retainWhere} from '../Utils/utils';
-
-/**
- * Optimizes the code for running specifically in an SSR environment. This optimization
- * asssumes that setState will not be called during render during initial mount, which
- * allows inlining useState/useReducer.
- *
- * Optimizations:
- * - Inline useState/useReducer
- * - Remove effects
- * - Remove refs where known to be unused during render (eg directly passed to a dom node)
- * - Remove event handlers
- *
- * Note that an earlier pass already inlines useMemo/useCallback
- */
-export function optimizeForSSR(fn: HIRFunction): void {
-  const inlinedState = new Map<IdentifierId, InstructionValue>();
-  /**
-   * First pass identifies useState/useReducer which can be safely inlined. Any use
-   * of the hook return other than destructuring (with a specific pattern) prevents
-   * inlining.
-   *
-   * Supported cases:
-   * - `const [state, ] = useState( <primitive-array-or-object> )`
-   * - `const [state, ] = useReducer(..., <value>)`
-   * - `const [state, ] = useReducer[..., <value>, <init>]`
-   */
-  for (const block of fn.body.blocks.values()) {
-    for (const instr of block.instructions) {
-      const {value} = instr;
-      switch (value.kind) {
-        case 'Destructure': {
-          if (
-            inlinedState.has(value.value.identifier.id) &&
-            value.lvalue.pattern.kind === 'ArrayPattern' &&
-            value.lvalue.pattern.items.length >= 1 &&
-            value.lvalue.pattern.items[0].kind === 'Identifier'
-          ) {
-            // Allow destructuring of inlined states
-            continue;
-          }
-          break;
-        }
-        case 'MethodCall':
-        case 'CallExpression': {
-          const calleee =
-            value.kind === 'CallExpression' ? value.callee : value.property;
-          const hookKind = getHookKind(fn.env, calleee.identifier);
-          switch (hookKind) {
-            case 'useReducer': {
-              if (
-                value.args.length === 2 &&
-                value.args[1].kind === 'Identifier'
-              ) {
-                const arg = value.args[1];
-                const replace: LoadLocal = {
-                  kind: 'LoadLocal',
-                  place: arg,
-                  loc: arg.loc,
-                };
-                inlinedState.set(instr.lvalue.identifier.id, replace);
-              } else if (
-                value.args.length === 3 &&
-                value.args[1].kind === 'Identifier' &&
-                value.args[2].kind === 'Identifier'
-              ) {
-                const arg = value.args[1];
-                const initializer = value.args[2];
-                const replace: CallExpression = {
-                  kind: 'CallExpression',
-                  callee: initializer,
-                  args: [arg],
-                  loc: value.loc,
-                };
-                inlinedState.set(instr.lvalue.identifier.id, replace);
-              }
-              break;
-            }
-            case 'useState': {
-              if (
-                value.args.length === 1 &&
-                value.args[0].kind === 'Identifier'
-              ) {
-                const arg = value.args[0];
-                if (
-                  isPrimitiveType(arg.identifier) ||
-                  isPlainObjectType(arg.identifier) ||
-                  isArrayType(arg.identifier)
-                ) {
-                  const replace: LoadLocal = {
-                    kind: 'LoadLocal',
-                    place: arg,
-                    loc: arg.loc,
-                  };
-                  inlinedState.set(instr.lvalue.identifier.id, replace);
-                }
-              }
-              break;
-            }
-          }
-        }
-      }
-      // Any use of useState/useReducer return besides destructuring prevents inlining
-      if (inlinedState.size !== 0) {
-        for (const operand of eachInstructionValueOperand(value)) {
-          inlinedState.delete(operand.identifier.id);
-        }
-      }
-    }
-    if (inlinedState.size !== 0) {
-      for (const operand of eachTerminalOperand(block.terminal)) {
-        inlinedState.delete(operand.identifier.id);
-      }
-    }
-  }
-  for (const block of fn.body.blocks.values()) {
-    for (const instr of block.instructions) {
-      const {value} = instr;
-      switch (value.kind) {
-        case 'FunctionExpression': {
-          if (hasKnownNonRenderCall(value.loweredFunc.func)) {
-            instr.value = {
-              kind: 'Primitive',
-              value: undefined,
-              loc: value.loc,
-            };
-          }
-          break;
-        }
-        case 'JsxExpression': {
-          if (
-            value.tag.kind === 'BuiltinTag' &&
-            value.tag.name.indexOf('-') === -1
-          ) {
-            const tag = value.tag.name;
-            retainWhere(value.props, prop => {
-              return (
-                prop.kind === 'JsxSpreadAttribute' ||
-                (!isKnownEventHandler(tag, prop.name) && prop.name !== 'ref')
-              );
-            });
-          }
-          break;
-        }
-        case 'Destructure': {
-          if (inlinedState.has(value.value.identifier.id)) {
-            // Canonical check is part of determining if state can inline, this is for TS
-            CompilerError.invariant(
-              value.lvalue.pattern.kind === 'ArrayPattern' &&
-                value.lvalue.pattern.items.length >= 1 &&
-                value.lvalue.pattern.items[0].kind === 'Identifier',
-              {
-                reason:
-                  'Expected a valid destructuring pattern for inlined state',
-                description: null,
-                details: [
-                  {
-                    kind: 'error',
-                    message: 'Expected a valid destructuring pattern',
-                    loc: value.loc,
-                  },
-                ],
-              },
-            );
-            const store: StoreLocal = {
-              kind: 'StoreLocal',
-              loc: value.loc,
-              type: null,
-              lvalue: {
-                kind: value.lvalue.kind,
-                place: value.lvalue.pattern.items[0],
-              },
-              value: value.value,
-            };
-            instr.value = store;
-          }
-          break;
-        }
-        case 'MethodCall':
-        case 'CallExpression': {
-          const calleee =
-            value.kind === 'CallExpression' ? value.callee : value.property;
-          const hookKind = getHookKind(fn.env, calleee.identifier);
-          switch (hookKind) {
-            case 'useEffectEvent': {
-              if (
-                value.args.length === 1 &&
-                value.args[0].kind === 'Identifier'
-              ) {
-                const load: LoadLocal = {
-                  kind: 'LoadLocal',
-                  place: value.args[0],
-                  loc: value.loc,
-                };
-                instr.value = load;
-              }
-              break;
-            }
-            case 'useEffect':
-            case 'useLayoutEffect':
-            case 'useInsertionEffect': {
-              // Drop effects
-              instr.value = {
-                kind: 'Primitive',
-                value: undefined,
-                loc: value.loc,
-              };
-              break;
-            }
-            case 'useReducer':
-            case 'useState': {
-              const replace = inlinedState.get(instr.lvalue.identifier.id);
-              if (replace != null) {
-                instr.value = replace;
-              }
-              break;
-            }
-          }
-        }
-      }
-    }
-  }
-}
-
-function hasKnownNonRenderCall(fn: HIRFunction): boolean {
-  for (const block of fn.body.blocks.values()) {
-    for (const instr of block.instructions) {
-      if (
-        instr.value.kind === 'CallExpression' &&
-        (isSetStateType(instr.value.callee.identifier) ||
-          isStartTransitionType(instr.value.callee.identifier))
-      ) {
-        return true;
-      }
-    }
-  }
-  return false;
-}
-
-const EVENT_HANDLER_PATTERN = /^on[A-Z]/;
-function isKnownEventHandler(_tag: string, prop: string): boolean {
-  return EVENT_HANDLER_PATTERN.test(prop);
-}

compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/ssr/optimize-ssr.expect.md

@@ -1,30 +0,0 @@
-
-## Input
-
-```javascript
-// @enableOptimizeForSSR
-function Component() {
-  const [state, setState] = useState(0);
-  const ref = useRef(null);
-  const onChange = e => {
-    setState(e.target.value);
-  };
-  useEffect(() => {
-    log(ref.current.value);
-  });
-  return <input value={state} onChange={onChange} ref={ref} />;
-}
-
-```
-
-## Code
-
-```javascript
-// @enableOptimizeForSSR
-function Component() {
-  const state = 0;
-  return <input value={state} />;
-}
-
-```
-      

compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/ssr/optimize-ssr.js

@@ -1,12 +0,0 @@
-// @enableOptimizeForSSR
-function Component() {
-  const [state, setState] = useState(0);
-  const ref = useRef(null);
-  const onChange = e => {
-    setState(e.target.value);
-  };
-  useEffect(() => {
-    log(ref.current.value);
-  });
-  return <input value={state} onChange={onChange} ref={ref} />;
-}

compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/ssr/ssr-infer-event-handlers-from-setState.expect.md

@@ -1,36 +0,0 @@
-
-## Input
-
-```javascript
-// @enableOptimizeForSSR
-function Component() {
-  const [state, setState] = useState(0);
-  const ref = useRef(null);
-  const onChange = e => {
-    // The known setState call allows us to infer this as an event handler
-    // and prune it
-    setState(e.target.value);
-  };
-  useEffect(() => {
-    log(ref.current.value);
-  });
-  return <CustomInput value={state} onChange={onChange} ref={ref} />;
-}
-
-```
-
-## Code
-
-```javascript
-// @enableOptimizeForSSR
-function Component() {
-  const state = 0;
-  const ref = useRef(null);
-  const onChange = undefined;
-  return <CustomInput value={state} onChange={onChange} ref={ref} />;
-}
-
-```
-      
-### Eval output
-(kind: exception) Fixture not implemented

compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/ssr/ssr-infer-event-handlers-from-setState.js

@@ -1,14 +0,0 @@
-// @enableOptimizeForSSR
-function Component() {
-  const [state, setState] = useState(0);
-  const ref = useRef(null);
-  const onChange = e => {
-    // The known setState call allows us to infer this as an event handler
-    // and prune it
-    setState(e.target.value);
-  };
-  useEffect(() => {
-    log(ref.current.value);
-  });
-  return <CustomInput value={state} onChange={onChange} ref={ref} />;
-}

compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/ssr/ssr-infer-event-handlers-from-startTransition.expect.md

@@ -1,40 +0,0 @@
-
-## Input
-
-```javascript
-// @enableOptimizeForSSR
-function Component() {
-  const [, startTransition] = useTransition();
-  const [state, setState] = useState(0);
-  const ref = useRef(null);
-  const onChange = e => {
-    // The known startTransition call allows us to infer this as an event handler
-    // and prune it
-    startTransition(() => {
-      setState.call(null, e.target.value);
-    });
-  };
-  useEffect(() => {
-    log(ref.current.value);
-  });
-  return <CustomInput value={state} onChange={onChange} ref={ref} />;
-}
-
-```
-
-## Code
-
-```javascript
-// @enableOptimizeForSSR
-function Component() {
-  useTransition();
-  const state = 0;
-  const ref = useRef(null);
-  const onChange = undefined;
-  return <CustomInput value={state} onChange={onChange} ref={ref} />;
-}
-
-```
-      
-### Eval output
-(kind: exception) Fixture not implemented

compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/ssr/ssr-infer-event-handlers-from-startTransition.js

@@ -1,17 +0,0 @@
-// @enableOptimizeForSSR
-function Component() {
-  const [, startTransition] = useTransition();
-  const [state, setState] = useState(0);
-  const ref = useRef(null);
-  const onChange = e => {
-    // The known startTransition call allows us to infer this as an event handler
-    // and prune it
-    startTransition(() => {
-      setState.call(null, e.target.value);
-    });
-  };
-  useEffect(() => {
-    log(ref.current.value);
-  });
-  return <CustomInput value={state} onChange={onChange} ref={ref} />;
-}

compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/ssr/ssr-use-reducer-initializer.expect.md

@@ -1,42 +0,0 @@
-
-## Input
-
-```javascript
-// @enableOptimizeForSSR
-
-import {useReducer} from 'react';
-
-const initializer = x => x;
-
-function Component() {
-  const [state, dispatch] = useReducer((_, next) => next, 0, initializer);
-  const ref = useRef(null);
-  const onChange = e => {
-    dispatch(e.target.value);
-  };
-  useEffect(() => {
-    log(ref.current.value);
-  });
-  return <input value={state} onChange={onChange} ref={ref} />;
-}
-
-```
-
-## Code
-
-```javascript
-// @enableOptimizeForSSR
-
-import { useReducer } from "react";
-
-const initializer = (x) => {
-  return x;
-};
-
-function Component() {
-  const state = initializer(0);
-  return <input value={state} />;
-}
-
-```
-      

compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/ssr/ssr-use-reducer-initializer.js

@@ -1,17 +0,0 @@
-// @enableOptimizeForSSR
-
-import {useReducer} from 'react';
-
-const initializer = x => x;
-
-function Component() {
-  const [state, dispatch] = useReducer((_, next) => next, 0, initializer);
-  const ref = useRef(null);
-  const onChange = e => {
-    dispatch(e.target.value);
-  };
-  useEffect(() => {
-    log(ref.current.value);
-  });
-  return <input value={state} onChange={onChange} ref={ref} />;
-}

compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/ssr/ssr-use-reducer.expect.md

@@ -1,36 +0,0 @@
-
-## Input
-
-```javascript
-// @enableOptimizeForSSR
-
-import {useReducer} from 'react';
-
-function Component() {
-  const [state, dispatch] = useReducer((_, next) => next, 0);
-  const ref = useRef(null);
-  const onChange = e => {
-    dispatch(e.target.value);
-  };
-  useEffect(() => {
-    log(ref.current.value);
-  });
-  return <input value={state} onChange={onChange} ref={ref} />;
-}
-
-```
-
-## Code
-
-```javascript
-// @enableOptimizeForSSR
-
-import { useReducer } from "react";
-
-function Component() {
-  const state = 0;
-  return <input value={state} />;
-}
-
-```
-      

compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/ssr/ssr-use-reducer.js

@@ -1,15 +0,0 @@
-// @enableOptimizeForSSR
-
-import {useReducer} from 'react';
-
-function Component() {
-  const [state, dispatch] = useReducer((_, next) => next, 0);
-  const ref = useRef(null);
-  const onChange = e => {
-    dispatch(e.target.value);
-  };
-  useEffect(() => {
-    log(ref.current.value);
-  });
-  return <input value={state} onChange={onChange} ref={ref} />;
-}

compiler/packages/snap/src/SproutTodoFilter.ts

@@ -487,13 +487,6 @@ const skipFilter = new Set([
   'lower-context-selector-simple',
   'lower-context-acess-multiple',
   'bug-separate-memoization-due-to-callback-capturing',
-
-  // SSR optimization rewrites files in a way that causes differences or warnings
-  'ssr/optimize-ssr',
-  'ssr/ssr-use-reducer',
-  'ssr/ssr-use-reducer-initializer',
-  'ssr/infer-event-handlers-from-setState',
-  'ssr/infer-event-handlers-from-startTransition',
 ]);
 
 export default skipFilter;