[compiler][poc] Quick experiment with SSR-optimization pass

Just a quick poc:
* Inline useState when the initializer is known to not be a function. The heuristic could be improved but will handle a large number of cases already.
* Prune effects
* Prune useRef if the ref is unused, by pruning 'ref' props on primitive components. Then DCE does the rest of the work - with a small change to allow `useRef()` calls to be dropped since function calls aren't normally eligible for dropping.
* Prune event handlers, by pruning props whose names start w "on" from primitive components. Then DCE removes the functions themselves.

Per the fixture, this gets pretty far.

Author: josephsavona

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

@@ -105,6 +105,7 @@ 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';
 
 export type CompilerPipelineValue =
   | {kind: 'ast'; name: string; value: CodegenFunction}
@@ -236,6 +237,11 @@ function runWithEnvironment(
     }
   }
 
+  if (env.config.enableAllowSetStateFromRefsInEffects) {
+    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});
@@ -305,7 +311,7 @@ function runWithEnvironment(
     value: hir,
   });
 
-  if (env.isInferredMemoEnabled) {
+  if (env.isInferredMemoEnabled && !env.config.enableOptimizeForSSR) {
     if (env.config.validateStaticComponents) {
       env.logErrors(validateStaticComponents(hir));
     }

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

@@ -670,6 +670,8 @@ export const EnvironmentConfigSchema = z.object({
    * from refs need to be stored in state during mount.
    */
   enableAllowSetStateFromRefsInEffects: z.boolean().default(true),
+
+  enableOptimizeForSSR: z.boolean().default(false),
 });
 
 export type EnvironmentConfig = z.infer<typeof EnvironmentConfigSchema>;

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

@@ -1823,6 +1823,10 @@ 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,6 +7,8 @@
 
 import {
   BlockId,
+  Environment,
+  getHookKind,
   HIRFunction,
   Identifier,
   IdentifierId,
@@ -68,9 +70,14 @@ 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);
@@ -112,7 +119,7 @@ function findReferencedIdentifiers(fn: HIRFunction): State {
   const hasLoop = hasBackEdge(fn);
   const reversedBlocks = [...fn.body.blocks.values()].reverse();
 
-  const state = new State();
+  const state = new State(fn.env);
   let size = state.count;
   do {
     size = state.count;
@@ -310,12 +317,23 @@ function pruneableValue(value: InstructionValue, state: State): boolean {
       // explicitly retain debugger statements to not break debugging workflows
       return false;
     }
-    case 'Await':
     case 'CallExpression':
+    case 'MethodCall': {
+      const calleee =
+        value.kind === 'CallExpression' ? value.callee : value.property;
+      const hookKind = getHookKind(state.env, calleee.identifier);
+      switch (hookKind) {
+        case 'useRef': {
+          // unused refs can be removed
+          return true;
+        }
+      }
+      return false;
+    }
+    case 'Await':
     case 'ComputedDelete':
     case 'ComputedStore':
     case 'PropertyDelete':
-    case 'MethodCall':
     case 'PropertyStore':
     case 'StoreGlobal': {
       /*

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

@@ -0,0 +1,138 @@
+/**
+ * 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 {traverse} from '@babel/types';
+import {
+  getHookKind,
+  HIRFunction,
+  IdentifierId,
+  isArrayType,
+  isObjectType,
+  isPlainObjectType,
+  isPrimitiveType,
+  isUseStateType,
+  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
+ * - Inline useMemo/useCallback (happens in earlier passes)
+ * - Remove effects
+ * - Remove refs where known to be unused during render (eg directly passed to a dom node)
+ * - Remove event handlers
+ */
+export function optimizeForSSR(fn: HIRFunction): void {
+  const inlinedState = new Set<IdentifierId>();
+  for (const block of fn.body.blocks.values()) {
+    for (const instr of block.instructions) {
+      const {value} = instr;
+      switch (value.kind) {
+        case 'JsxExpression': {
+          if (
+            value.tag.kind === 'BuiltinTag' &&
+            value.tag.name.indexOf('-') === -1
+          ) {
+            retainWhere(value.props, prop => {
+              return (
+                prop.kind === 'JsxSpreadAttribute' ||
+                (prop.name !== 'ref' && !prop.name.startsWith('on'))
+              );
+            });
+          }
+          break;
+        }
+        case 'Destructure': {
+          if (
+            isUseStateType(value.value.identifier) &&
+            inlinedState.has(value.value.identifier.id) &&
+            value.lvalue.pattern.kind === 'ArrayPattern' &&
+            value.lvalue.pattern.items.length >= 1 &&
+            value.lvalue.pattern.items[0].kind === 'Identifier'
+          ) {
+            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 '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)
+                ) {
+                  instr.value = {
+                    kind: 'LoadLocal',
+                    place: arg,
+                    loc: arg.loc,
+                  };
+                  inlinedState.add(instr.lvalue.identifier.id);
+                }
+              }
+              break;
+            }
+          }
+        }
+      }
+    }
+  }
+}

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

@@ -0,0 +1,32 @@
+
+## 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} />;
+}
+
+```
+      
+### Eval output
+(kind: exception) Fixture not implemented

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

@@ -0,0 +1,12 @@
+// @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} />;
+}