[compiler] Promote temporaries when necessary to prevent codegen reor…

…dering over side-effectful operations

ghstack-source-id: 639191e63a0d2b4290d1265a2da12fb17de750d9
Pull Request resolved: #30554

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

@@ -452,17 +452,17 @@ function* runWithEnvironment(
     value: reactiveFunction,
   });
 
-  promoteUsedTemporaries(reactiveFunction);
+  pruneUnusedLValues(reactiveFunction);
   yield log({
     kind: 'reactive',
-    name: 'PromoteUsedTemporaries',
+    name: 'PruneUnusedLValues',
     value: reactiveFunction,
   });
 
-  pruneUnusedLValues(reactiveFunction);
+  promoteUsedTemporaries(reactiveFunction);
   yield log({
     kind: 'reactive',
-    name: 'PruneUnusedLValues',
+    name: 'PromoteUsedTemporaries',
     value: reactiveFunction,
   });
 

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

@@ -1209,7 +1209,7 @@ function codegenInstructionNullable(
       value = null;
     } else {
       lvalue = instr.value.lvalue.pattern;
-      let hasReasign = false;
+      let hasReassign = false;
       let hasDeclaration = false;
       for (const place of eachPatternOperand(lvalue)) {
         if (
@@ -1219,18 +1219,18 @@ function codegenInstructionNullable(
           cx.temp.set(place.identifier.declarationId, null);
         }
         const isDeclared = cx.hasDeclared(place.identifier);
-        hasReasign ||= isDeclared;
+        hasReassign ||= isDeclared;
         hasDeclaration ||= !isDeclared;
       }
-      if (hasReasign && hasDeclaration) {
+      if (hasReassign && hasDeclaration) {
         CompilerError.invariant(false, {
           reason:
             'Encountered a destructuring operation where some identifiers are already declared (reassignments) but others are not (declarations)',
           description: null,
           loc: instr.loc,
           suggestions: null,
         });
-      } else if (hasReasign) {
+      } else if (hasReassign) {
         kind = InstructionKind.Reassign;
       }
       value = codegenPlaceToExpression(cx, instr.value.value);

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

@@ -15,13 +15,17 @@ import {
   PrunedReactiveScopeBlock,
   ReactiveFunction,
   ReactiveScope,
+  ReactiveInstruction,
   ReactiveScopeBlock,
   ReactiveValue,
   ScopeId,
+  SpreadPattern,
   promoteTemporary,
   promoteTemporaryJsxTag,
+  IdentifierId,
 } from '../HIR/HIR';
 import {ReactiveFunctionVisitor, visitReactiveFunction} from './visitors';
+import {eachInstructionValueLValue, eachPatternOperand} from '../HIR/visitors';
 
 /**
  * Phase 2: Promote identifiers which are used in a place that requires a named variable.
@@ -225,6 +229,199 @@ class CollectPromotableTemporaries extends ReactiveFunctionVisitor<State> {
   }
 }
 
+type InterState = Map<IdentifierId, [Identifier, boolean]>;
+class PromoteInterposedTemporaries extends ReactiveFunctionVisitor<InterState> {
+  #promotable: State;
+  #consts: Set<IdentifierId> = new Set();
+  #globals: Set<IdentifierId> = new Set();
+
+  /*
+   * Unpromoted temporaries will be emitted at their use sites rather than as separate
+   * declarations. However, this causes errors if an interposing temporary has been
+   * promoted, or if an interposing instruction has had its lvalues deleted, because such
+   * temporaries will be emitted as separate statements, which can effectively cause
+   * code to be reordered, and when that code has side effects that changes program behavior.
+   * This visitor promotes temporarties that have such interposing instructions to preserve
+   * source ordering.
+   */
+  constructor(promotable: State, params: Array<Place | SpreadPattern>) {
+    super();
+    params.forEach(param => {
+      switch (param.kind) {
+        case 'Identifier':
+          this.#consts.add(param.identifier.id);
+          break;
+        case 'Spread':
+          this.#consts.add(param.place.identifier.id);
+          break;
+      }
+    });
+    this.#promotable = promotable;
+  }
+
+  override visitPlace(
+    _id: InstructionId,
+    place: Place,
+    state: InterState,
+  ): void {
+    const promo = state.get(place.identifier.id);
+    if (promo) {
+      const [identifier, needsPromotion] = promo;
+      if (
+        needsPromotion &&
+        identifier.name === null &&
+        !this.#consts.has(identifier.id)
+      ) {
+        /*
+         * If the identifier hasn't been promoted but is marked as needing
+         * promotion by the logic in `visitInstruction`, and we've seen a
+         * use of it after said marking, promote it
+         */
+        promoteIdentifier(identifier, this.#promotable);
+      }
+    }
+  }
+
+  override visitInstruction(
+    instruction: ReactiveInstruction,
+    state: InterState,
+  ): void {
+    for (const lval of eachInstructionValueLValue(instruction.value)) {
+      CompilerError.invariant(lval.identifier.name != null, {
+        reason:
+          'PromoteInterposedTemporaries: Assignment targets not expected to be temporaries',
+        loc: instruction.loc,
+      });
+    }
+
+    switch (instruction.value.kind) {
+      case 'CallExpression':
+      case 'MethodCall':
+      case 'Await':
+      case 'PropertyStore':
+      case 'PropertyDelete':
+      case 'ComputedStore':
+      case 'ComputedDelete':
+      case 'PostfixUpdate':
+      case 'PrefixUpdate':
+      case 'StoreLocal':
+      case 'StoreContext':
+      case 'StoreGlobal':
+      case 'Destructure': {
+        let constStore = false;
+
+        if (
+          (instruction.value.kind === 'StoreContext' ||
+            instruction.value.kind === 'StoreLocal') &&
+          (instruction.value.lvalue.kind === 'Const' ||
+            instruction.value.lvalue.kind === 'HoistedConst')
+        ) {
+          /*
+           * If an identifier is const, we don't need to worry about it
+           * being mutated between being loaded and being used
+           */
+          this.#consts.add(instruction.value.lvalue.place.identifier.id);
+          constStore = true;
+        }
+        if (
+          instruction.value.kind === 'Destructure' &&
+          (instruction.value.lvalue.kind === 'Const' ||
+            instruction.value.lvalue.kind === 'HoistedConst')
+        ) {
+          [...eachPatternOperand(instruction.value.lvalue.pattern)].forEach(
+            ident => this.#consts.add(ident.identifier.id),
+          );
+          constStore = true;
+        }
+        if (instruction.value.kind === 'MethodCall') {
+          // Treat property of method call as constlike so we don't promote it.
+          this.#consts.add(instruction.value.property.identifier.id);
+        }
+
+        super.visitInstruction(instruction, state);
+        if (
+          !constStore &&
+          (instruction.lvalue == null ||
+            instruction.lvalue.identifier.name != null)
+        ) {
+          /*
+           * If we've stripped the lvalue or promoted the lvalue, then we will emit this instruction
+           * as a statement in codegen.
+           *
+           * If this instruction will be emitted directly as a statement rather than as a temporary
+           * during codegen, then it can interpose between the defs and the uses of other temporaries.
+           * Since this instruction could potentially mutate those defs, it's not safe to relocate
+           * the definition of those temporaries to after this instruction. Mark all those temporaries
+           * as needing promotion, but don't promote them until we actually see them being used.
+           */
+          for (const [key, [ident, _]] of state.entries()) {
+            state.set(key, [ident, true]);
+          }
+        }
+        if (instruction.lvalue && instruction.lvalue.identifier.name === null) {
+          // Add this instruction's lvalue to the state, initially not marked as needing promotion
+          state.set(instruction.lvalue.identifier.id, [
+            instruction.lvalue.identifier,
+            false,
+          ]);
+        }
+        break;
+      }
+      case 'DeclareContext':
+      case 'DeclareLocal': {
+        if (
+          instruction.value.lvalue.kind === 'Const' ||
+          instruction.value.lvalue.kind === 'HoistedConst'
+        ) {
+          this.#consts.add(instruction.value.lvalue.place.identifier.id);
+        }
+        super.visitInstruction(instruction, state);
+        break;
+      }
+      case 'LoadContext':
+      case 'LoadLocal': {
+        if (instruction.lvalue && instruction.lvalue.identifier.name === null) {
+          if (this.#consts.has(instruction.value.place.identifier.id)) {
+            this.#consts.add(instruction.lvalue.identifier.id);
+          }
+          state.set(instruction.lvalue.identifier.id, [
+            instruction.lvalue.identifier,
+            false,
+          ]);
+        }
+        super.visitInstruction(instruction, state);
+        break;
+      }
+      case 'PropertyLoad':
+      case 'ComputedLoad': {
+        if (instruction.lvalue) {
+          if (this.#globals.has(instruction.value.object.identifier.id)) {
+            this.#globals.add(instruction.lvalue.identifier.id);
+            this.#consts.add(instruction.lvalue.identifier.id);
+          }
+          if (instruction.lvalue.identifier.name === null) {
+            state.set(instruction.lvalue.identifier.id, [
+              instruction.lvalue.identifier,
+              false,
+            ]);
+          }
+        }
+        super.visitInstruction(instruction, state);
+        break;
+      }
+      case 'LoadGlobal': {
+        instruction.lvalue &&
+          this.#globals.add(instruction.lvalue.identifier.id);
+        super.visitInstruction(instruction, state);
+        break;
+      }
+      default: {
+        super.visitInstruction(instruction, state);
+      }
+    }
+  }
+}
+
 export function promoteUsedTemporaries(fn: ReactiveFunction): void {
   const state: State = {
     tags: new Set(),
@@ -239,6 +436,12 @@ export function promoteUsedTemporaries(fn: ReactiveFunction): void {
     }
   }
   visitReactiveFunction(fn, new PromoteTemporaries(), state);
+
+  visitReactiveFunction(
+    fn,
+    new PromoteInterposedTemporaries(state, fn.params),
+    new Map(),
+  );
   visitReactiveFunction(
     fn,
     new PromoteAllInstancedOfPromotedTemporaries(),