[compiler] Stay in SSA form through entire pipeline

This PR updates to use SSA form through the entire compilation pipeline. This means that in both HIR form and ReactiveFunction form, `Identifier` instances map 1:1 to `IdentifierId` values. If two identifiers have the same IdentifierId, they are the same instance. What this means is that all our passes can use this more precise information to determine if two particular identifiers are not just the same variable, but the same SSA "version" of that variable.

However, some parts of our analysis really care about program variables as opposed to SSA versions, and were relying on LeaveSSA to reset identifiers such that all Identifier instances for a particular program variable would have the same IdentifierId (though not necessarily the same Identifier instance). With LeaveSSA removed, those analysis passes can now use DeclarationId instead to uniquely identify a program variable.

Note that this PR surfaces some opportunties to improve edge-cases around reassigned values being declared/reassigned/depended-upon across multiple scopes. Several passes could/should use IdentifierId to more precisely identify exactly which values are accessed - for example, a scope that reassigns `x` but doesn't use `x` prior to reassignment doesn't have to take a dependency on `x`. But today we take a dependnecy.

My approach for these cases was to add a "TODO LeaveSSA" comment with notes and the name of the fixture demonstrating the difference, but to intentionally preserve the existing behavior (generally, switching to use DeclarationId when IdentifierId would have been more precise).

Beyond updating passes to use DeclarationId instead of Identifier/IdentifierId, the other change here is to extract out the remaining necessary bits of LeaveSSA into a new pass that rewrites InstructionKind (const/let/reassign/etc) based on whether a value is actually const or has reassignments and should be let.

ghstack-source-id: 4b9b1d5
Pull Request resolved: #30573

Author: josephsavona

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

@@ -77,7 +77,11 @@ import {flattenScopesWithHooksOrUseHIR} from '../ReactiveScopes/FlattenScopesWit
 import {pruneAlwaysInvalidatingScopes} from '../ReactiveScopes/PruneAlwaysInvalidatingScopes';
 import pruneInitializationDependencies from '../ReactiveScopes/PruneInitializationDependencies';
 import {stabilizeBlockIds} from '../ReactiveScopes/StabilizeBlockIds';
-import {eliminateRedundantPhi, enterSSA, leaveSSA} from '../SSA';
+import {
+  eliminateRedundantPhi,
+  enterSSA,
+  rewriteInstructionKindsBasedOnReassignment,
+} from '../SSA';
 import {inferTypes} from '../TypeInference';
 import {
   logCodegenFunction,
@@ -98,6 +102,7 @@ import {
 } from '../Validation';
 import {validateLocalsNotReassignedAfterRender} from '../Validation/ValidateLocalsNotReassignedAfterRender';
 import {outlineFunctions} from '../Optimization/OutlineFunctions';
+import {propagatePhiTypes} from '../TypeInference/PropagatePhiTypes';
 
 export type CompilerPipelineValue =
   | {kind: 'ast'; name: string; value: CodegenFunction}
@@ -237,8 +242,19 @@ function* runWithEnvironment(
   inferReactivePlaces(hir);
   yield log({kind: 'hir', name: 'InferReactivePlaces', value: hir});
 
-  leaveSSA(hir);
-  yield log({kind: 'hir', name: 'LeaveSSA', value: hir});
+  rewriteInstructionKindsBasedOnReassignment(hir);
+  yield log({
+    kind: 'hir',
+    name: 'RewriteInstructionKindsBasedOnReassignment',
+    value: hir,
+  });
+
+  propagatePhiTypes(hir);
+  yield log({
+    kind: 'hir',
+    name: 'PropagatePhiTypes',
+    value: hir,
+  });
 
   inferReactiveScopeVariables(hir);
   yield log({kind: 'hir', name: 'InferReactiveScopeVariables', value: hir});

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

@@ -1248,6 +1248,9 @@ export function makeIdentifierName(name: string): ValidatedIdentifier {
 
 /**
  * Given an unnamed identifier, promote it to a named identifier.
+ *
+ * Note: this uses the identifier's DeclarationId to ensure that all
+ * instances of the same declaration will have the same name.
  */
 export function promoteTemporary(identifier: Identifier): void {
   CompilerError.invariant(identifier.name === null, {
@@ -1258,7 +1261,7 @@ export function promoteTemporary(identifier: Identifier): void {
   });
   identifier.name = {
     kind: 'promoted',
-    value: `#t${identifier.id}`,
+    value: `#t${identifier.declarationId}`,
   };
 }
 
@@ -1269,6 +1272,9 @@ export function isPromotedTemporary(name: string): boolean {
 /**
  * Given an unnamed identifier, promote it to a named identifier, distinguishing
  * it as a value that needs to be capitalized since it appears in JSX element tag position
+ *
+ * Note: this uses the identifier's DeclarationId to ensure that all
+ * instances of the same declaration will have the same name.
  */
 export function promoteTemporaryJsxTag(identifier: Identifier): void {
   CompilerError.invariant(identifier.name === null, {
@@ -1279,7 +1285,7 @@ export function promoteTemporaryJsxTag(identifier: Identifier): void {
   });
   identifier.name = {
     kind: 'promoted',
-    value: `#T${identifier.id}`,
+    value: `#T${identifier.declarationId}`,
   };
 }
 

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

@@ -908,3 +908,16 @@ export function createTemporaryPlace(
     loc: GeneratedSource,
   };
 }
+
+/**
+ * Clones an existing Place, returning a new temporary Place that shares the
+ * same metadata properties as the original place (effect, reactive flag, type)
+ * but has a new, temporary Identifier.
+ */
+export function clonePlaceToTemporary(env: Environment, place: Place): Place {
+  const temp = createTemporaryPlace(env, place.loc);
+  temp.effect = place.effect;
+  temp.identifier.type = place.identifier.type;
+  temp.reactive = place.reactive;
+  return temp;
+}

compiler/packages/babel-plugin-react-compiler/src/Inference/AnalyseFunctions.ts

@@ -20,7 +20,7 @@ import {
 } from '../HIR';
 import {deadCodeElimination} from '../Optimization';
 import {inferReactiveScopeVariables} from '../ReactiveScopes';
-import {leaveSSA} from '../SSA';
+import {rewriteInstructionKindsBasedOnReassignment} from '../SSA';
 import {logHIRFunction} from '../Utils/logger';
 import {inferMutableContextVariables} from './InferMutableContextVariables';
 import {inferMutableRanges} from './InferMutableRanges';
@@ -108,7 +108,7 @@ function lower(func: HIRFunction): void {
   inferReferenceEffects(func, {isFunctionExpression: true});
   deadCodeElimination(func);
   inferMutableRanges(func);
-  leaveSSA(func);
+  rewriteInstructionKindsBasedOnReassignment(func);
   inferReactiveScopeVariables(func);
   inferMutableContextVariables(func);
   logHIRFunction('AnalyseFunction (inner)', func);

compiler/packages/babel-plugin-react-compiler/src/ReactiveScopes/CodegenReactiveFunction.ts

@@ -18,6 +18,7 @@ import {Environment, EnvironmentConfig, ExternalFunction} from '../HIR';
 import {
   ArrayPattern,
   BlockId,
+  DeclarationId,
   GeneratedSource,
   Identifier,
   IdentifierId,
@@ -309,9 +310,9 @@ function codegenReactiveFunction(
 ): Result<CodegenFunction, CompilerError> {
   for (const param of fn.params) {
     if (param.kind === 'Identifier') {
-      cx.temp.set(param.identifier.id, null);
+      cx.temp.set(param.identifier.declarationId, null);
     } else {
-      cx.temp.set(param.place.identifier.id, null);
+      cx.temp.set(param.place.identifier.declarationId, null);
     }
   }
 
@@ -392,7 +393,11 @@ class Context {
   env: Environment;
   fnName: string;
   #nextCacheIndex: number = 0;
-  #declarations: Set<IdentifierId> = new Set();
+  /**
+   * Tracks which named variables have been declared to dedupe declarations,
+   * so this uses DeclarationId instead of IdentifierId
+   */
+  #declarations: Set<DeclarationId> = new Set();
   temp: Temporaries;
   errors: CompilerError = new CompilerError();
   objectMethods: Map<IdentifierId, ObjectMethod> = new Map();
@@ -418,11 +423,11 @@ class Context {
   }
 
   declare(identifier: Identifier): void {
-    this.#declarations.add(identifier.id);
+    this.#declarations.add(identifier.declarationId);
   }
 
   hasDeclared(identifier: Identifier): boolean {
-    return this.#declarations.has(identifier.id);
+    return this.#declarations.has(identifier.declarationId);
   }
 
   synthesizeName(name: string): ValidIdentifierName {
@@ -1147,7 +1152,7 @@ function codegenTerminal(
       let catchParam = null;
       if (terminal.handlerBinding !== null) {
         catchParam = convertIdentifier(terminal.handlerBinding.identifier);
-        cx.temp.set(terminal.handlerBinding.identifier.id, null);
+        cx.temp.set(terminal.handlerBinding.identifier.declarationId, null);
       }
       return t.tryStatement(
         codegenBlock(cx, terminal.block),
@@ -1205,7 +1210,7 @@ function codegenInstructionNullable(
           kind !== InstructionKind.Reassign &&
           place.identifier.name === null
         ) {
-          cx.temp.set(place.identifier.id, null);
+          cx.temp.set(place.identifier.declarationId, null);
         }
         const isDeclared = cx.hasDeclared(place.identifier);
         hasReasign ||= isDeclared;
@@ -1261,7 +1266,7 @@ function codegenInstructionNullable(
         );
         if (instr.lvalue !== null) {
           if (instr.value.kind !== 'StoreContext') {
-            cx.temp.set(instr.lvalue.identifier.id, expr);
+            cx.temp.set(instr.lvalue.identifier.declarationId, expr);
             return null;
           } else {
             // Handle chained reassignments for context variables
@@ -1530,7 +1535,7 @@ function createCallExpression(
   }
 }
 
-type Temporaries = Map<IdentifierId, t.Expression | t.JSXText | null>;
+type Temporaries = Map<DeclarationId, t.Expression | t.JSXText | null>;
 
 function codegenLabel(id: BlockId): string {
   return `bb${id}`;
@@ -1549,7 +1554,7 @@ function codegenInstruction(
   }
   if (instr.lvalue.identifier.name === null) {
     // temporary
-    cx.temp.set(instr.lvalue.identifier.id, value);
+    cx.temp.set(instr.lvalue.identifier.declarationId, value);
     return t.emptyStatement();
   } else {
     const expressionValue = convertValueToExpression(value);
@@ -2498,7 +2503,7 @@ function codegenPlaceToExpression(cx: Context, place: Place): t.Expression {
 }
 
 function codegenPlace(cx: Context, place: Place): t.Expression | t.JSXText {
-  let tmp = cx.temp.get(place.identifier.id);
+  let tmp = cx.temp.get(place.identifier.declarationId);
   if (tmp != null) {
     return tmp;
   }

compiler/packages/babel-plugin-react-compiler/src/ReactiveScopes/ExtractScopeDeclarationsFromDestructuring.ts

@@ -6,6 +6,7 @@
  */
 
 import {
+  DeclarationId,
   Destructure,
   Environment,
   IdentifierId,
@@ -17,6 +18,7 @@ import {
   ReactiveStatement,
   promoteTemporary,
 } from '../HIR';
+import {clonePlaceToTemporary} from '../HIR/HIRBuilder';
 import {eachPatternOperand, mapPatternOperands} from '../HIR/visitors';
 import {
   ReactiveFunctionTransform,
@@ -82,7 +84,11 @@ export function extractScopeDeclarationsFromDestructuring(
 
 class State {
   env: Environment;
-  declared: Set<IdentifierId> = new Set();
+  /**
+   * We need to track which program variables are already declared to convert
+   * declarations into reassignments, so we use DeclarationId
+   */
+  declared: Set<DeclarationId> = new Set();
 
   constructor(env: Environment) {
     this.env = env;
@@ -92,7 +98,7 @@ class State {
 class Visitor extends ReactiveFunctionTransform<State> {
   override visitScope(scope: ReactiveScopeBlock, state: State): void {
     for (const [, declaration] of scope.scope.declarations) {
-      state.declared.add(declaration.identifier.id);
+      state.declared.add(declaration.identifier.declarationId);
     }
     this.traverseScope(scope, state);
   }
@@ -131,7 +137,7 @@ function transformDestructuring(
   let reassigned: Set<IdentifierId> = new Set();
   let hasDeclaration = false;
   for (const place of eachPatternOperand(destructure.lvalue.pattern)) {
-    const isDeclared = state.declared.has(place.identifier.id);
+    const isDeclared = state.declared.has(place.identifier.declarationId);
     if (isDeclared) {
       reassigned.add(place.identifier.id);
     }
@@ -150,15 +156,7 @@ function transformDestructuring(
     if (!reassigned.has(place.identifier.id)) {
       return place;
     }
-    const tempId = state.env.nextIdentifierId;
-    const temporary = {
-      ...place,
-      identifier: {
-        ...place.identifier,
-        id: tempId,
-        name: null, // overwritten below
-      },
-    };
+    const temporary = clonePlaceToTemporary(state.env, place);
     promoteTemporary(temporary.identifier);
     renamed.set(place, temporary);
     return temporary;

compiler/packages/babel-plugin-react-compiler/src/ReactiveScopes/InferReactiveScopeVariables.ts

@@ -8,6 +8,7 @@
 import {CompilerError, SourceLocation} from '..';
 import {Environment} from '../HIR';
 import {
+  DeclarationId,
   GeneratedSource,
   HIRFunction,
   Identifier,
@@ -257,21 +258,34 @@ export function findDisjointMutableValues(
   fn: HIRFunction,
 ): DisjointSet<Identifier> {
   const scopeIdentifiers = new DisjointSet<Identifier>();
+
+  const declarations = new Map<DeclarationId, Identifier>();
+  function declareIdentifier(lvalue: Place): void {
+    if (!declarations.has(lvalue.identifier.declarationId)) {
+      declarations.set(lvalue.identifier.declarationId, lvalue.identifier);
+    }
+  }
+
   for (const [_, block] of fn.body.blocks) {
     /*
      * If a phi is mutated after creation, then we need to alias all of its operands such that they
      * are assigned to the same scope.
      */
     for (const phi of block.phis) {
       if (
-        // The phi was reset because it was not mutated after creation
         phi.id.mutableRange.start + 1 !== phi.id.mutableRange.end &&
         phi.id.mutableRange.end >
           (block.instructions.at(0)?.id ?? block.terminal.id)
       ) {
-        for (const [, phiId] of phi.operands) {
-          scopeIdentifiers.union([phi.id, phiId]);
+        const operands = [phi.id];
+        const declaration = declarations.get(phi.id.declarationId);
+        if (declaration !== undefined) {
+          operands.push(declaration);
+        }
+        for (const [_, phiId] of phi.operands) {
+          operands.push(phiId);
         }
+        scopeIdentifiers.union(operands);
       } else if (fn.env.config.enableForest) {
         for (const [, phiId] of phi.operands) {
           scopeIdentifiers.union([phi.id, phiId]);
@@ -286,9 +300,15 @@ export function findDisjointMutableValues(
         operands.push(instr.lvalue!.identifier);
       }
       if (
+        instr.value.kind === 'DeclareLocal' ||
+        instr.value.kind === 'DeclareContext'
+      ) {
+        declareIdentifier(instr.value.lvalue.place);
+      } else if (
         instr.value.kind === 'StoreLocal' ||
         instr.value.kind === 'StoreContext'
       ) {
+        declareIdentifier(instr.value.lvalue.place);
         if (
           instr.value.lvalue.place.identifier.mutableRange.end >
           instr.value.lvalue.place.identifier.mutableRange.start + 1
@@ -303,6 +323,7 @@ export function findDisjointMutableValues(
         }
       } else if (instr.value.kind === 'Destructure') {
         for (const place of eachPatternOperand(instr.value.lvalue.pattern)) {
+          declareIdentifier(place);
           if (
             place.identifier.mutableRange.end >
             place.identifier.mutableRange.start + 1