[compiler] Stay in SSA form through entire pipeline

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