[compiler] Promote temporaries when necessary to prevent codegen reordering over side-effectful operations

ghstack-source-id: 639191e63a
Pull Request resolved: https://github.com/facebook/react/pull/30554

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

@@ -452,13 +452,6 @@ function* runWithEnvironment(
     value: reactiveFunction,
   });
 
-  promoteUsedTemporaries(reactiveFunction);
-  yield log({
-    kind: 'reactive',
-    name: 'PromoteUsedTemporaries',
-    value: reactiveFunction,
-  });
-
   pruneUnusedLValues(reactiveFunction);
   yield log({
     kind: 'reactive',
@@ -466,6 +459,13 @@ function* runWithEnvironment(
     value: reactiveFunction,
   });
 
+  promoteUsedTemporaries(reactiveFunction);
+  yield log({
+    kind: 'reactive',
+    name: 'PromoteUsedTemporaries',
+    value: reactiveFunction,
+  });
+
   extractScopeDeclarationsFromDestructuring(reactiveFunction);
   yield log({
     kind: 'reactive',

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,10 +1219,10 @@ 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)',
@@ -1230,7 +1230,7 @@ function codegenInstructionNullable(
           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(),

compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/bug-codegen-inline-iife.expect.md

@@ -1,92 +0,0 @@
-
-## Input
-
-```javascript
-import {makeArray, print} from 'shared-runtime';
-
-/**
- * Exposes bug involving iife inlining + codegen.
- * We currently inline iifes to labeled blocks (not value-blocks).
- *
- * Here, print(1) and the evaluation of makeArray(...) get the same scope
- * as the compiler infers that the makeArray call may mutate its arguments.
- * Since print(1) does not get its own scope (and is thus not a declaration
- * or dependency), it does not get promoted.
- * As a result, print(1) gets reordered across the labeled-block instructions
- * to be inlined at the makeArray callsite.
- *
- * Current evaluator results:
- *   Found differences in evaluator results
- *   Non-forget (expected):
- *   (kind: ok) [null,2]
- *   logs: [1,2]
- *   Forget:
- *   (kind: ok) [null,2]
- *   logs: [2,1]
- */
-function useTest() {
-  return makeArray<number | void>(
-    print(1),
-    (function foo() {
-      print(2);
-      return 2;
-    })(),
-  );
-}
-
-export const FIXTURE_ENTRYPOINT = {
-  fn: useTest,
-  params: [],
-};
-
-```
-
-## Code
-
-```javascript
-import { c as _c } from "react/compiler-runtime";
-import { makeArray, print } from "shared-runtime";
-
-/**
- * Exposes bug involving iife inlining + codegen.
- * We currently inline iifes to labeled blocks (not value-blocks).
- *
- * Here, print(1) and the evaluation of makeArray(...) get the same scope
- * as the compiler infers that the makeArray call may mutate its arguments.
- * Since print(1) does not get its own scope (and is thus not a declaration
- * or dependency), it does not get promoted.
- * As a result, print(1) gets reordered across the labeled-block instructions
- * to be inlined at the makeArray callsite.
- *
- * Current evaluator results:
- *   Found differences in evaluator results
- *   Non-forget (expected):
- *   (kind: ok) [null,2]
- *   logs: [1,2]
- *   Forget:
- *   (kind: ok) [null,2]
- *   logs: [2,1]
- */
-function useTest() {
-  const $ = _c(1);
-  let t0;
-  if ($[0] === Symbol.for("react.memo_cache_sentinel")) {
-    let t1;
-
-    print(2);
-    t1 = 2;
-    t0 = makeArray(print(1), t1);
-    $[0] = t0;
-  } else {
-    t0 = $[0];
-  }
-  return t0;
-}
-
-export const FIXTURE_ENTRYPOINT = {
-  fn: useTest,
-  params: [],
-};
-
-```
-      

compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/bug-codegen-inline-iife.ts

@@ -1,36 +0,0 @@
-import {makeArray, print} from 'shared-runtime';
-
-/**
- * Exposes bug involving iife inlining + codegen.
- * We currently inline iifes to labeled blocks (not value-blocks).
- *
- * Here, print(1) and the evaluation of makeArray(...) get the same scope
- * as the compiler infers that the makeArray call may mutate its arguments.
- * Since print(1) does not get its own scope (and is thus not a declaration
- * or dependency), it does not get promoted.
- * As a result, print(1) gets reordered across the labeled-block instructions
- * to be inlined at the makeArray callsite.
- *
- * Current evaluator results:
- *   Found differences in evaluator results
- *   Non-forget (expected):
- *   (kind: ok) [null,2]
- *   logs: [1,2]
- *   Forget:
- *   (kind: ok) [null,2]
- *   logs: [2,1]
- */
-function useTest() {
-  return makeArray<number | void>(
-    print(1),
-    (function foo() {
-      print(2);
-      return 2;
-    })(),
-  );
-}
-
-export const FIXTURE_ENTRYPOINT = {
-  fn: useTest,
-  params: [],
-};

compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/codegen-inline-iife-reassign.expect.md

@@ -0,0 +1,62 @@
+
+## Input
+
+```javascript
+import {makeArray, print} from 'shared-runtime';
+
+function useTest() {
+  let w = {};
+  return makeArray(
+    (w = 42),
+    w,
+    (function foo() {
+      w = 999;
+      return 2;
+    })(),
+  );
+}
+
+export const FIXTURE_ENTRYPOINT = {
+  fn: useTest,
+  params: [],
+};
+
+```
+
+## Code
+
+```javascript
+import { c as _c } from "react/compiler-runtime";
+import { makeArray, print } from "shared-runtime";
+
+function useTest() {
+  const $ = _c(1);
+  let t0;
+  if ($[0] === Symbol.for("react.memo_cache_sentinel")) {
+    let w;
+    w = {};
+
+    const t1 = (w = 42);
+    const t2 = w;
+
+    w;
+    let t3;
+    w = 999;
+    t3 = 2;
+    t0 = makeArray(t1, t2, t3);
+    $[0] = t0;
+  } else {
+    t0 = $[0];
+  }
+  return t0;
+}
+
+export const FIXTURE_ENTRYPOINT = {
+  fn: useTest,
+  params: [],
+};
+
+```
+      
+### Eval output
+(kind: ok) [42,42,2]

compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/codegen-inline-iife-reassign.ts

@@ -0,0 +1,18 @@
+import {makeArray, print} from 'shared-runtime';
+
+function useTest() {
+  let w = {};
+  return makeArray(
+    (w = 42),
+    w,
+    (function foo() {
+      w = 999;
+      return 2;
+    })(),
+  );
+}
+
+export const FIXTURE_ENTRYPOINT = {
+  fn: useTest,
+  params: [],
+};

compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/codegen-inline-iife-storeprop.expect.md

@@ -0,0 +1,60 @@
+
+## Input
+
+```javascript
+import {makeArray, print} from 'shared-runtime';
+
+function useTest() {
+  let w = {};
+  return makeArray(
+    (w.x = 42),
+    w.x,
+    (function foo() {
+      w.x = 999;
+      return 2;
+    })(),
+  );
+}
+
+export const FIXTURE_ENTRYPOINT = {
+  fn: useTest,
+  params: [],
+};
+
+```
+
+## Code
+
+```javascript
+import { c as _c } from "react/compiler-runtime";
+import { makeArray, print } from "shared-runtime";
+
+function useTest() {
+  const $ = _c(1);
+  let t0;
+  if ($[0] === Symbol.for("react.memo_cache_sentinel")) {
+    const w = {};
+
+    const t1 = (w.x = 42);
+    const t2 = w.x;
+    let t3;
+
+    w.x = 999;
+    t3 = 2;
+    t0 = makeArray(t1, t2, t3);
+    $[0] = t0;
+  } else {
+    t0 = $[0];
+  }
+  return t0;
+}
+
+export const FIXTURE_ENTRYPOINT = {
+  fn: useTest,
+  params: [],
+};
+
+```
+      
+### Eval output
+(kind: ok) [42,42,2]

compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/codegen-inline-iife-storeprop.ts

@@ -0,0 +1,18 @@
+import {makeArray, print} from 'shared-runtime';
+
+function useTest() {
+  let w = {};
+  return makeArray(
+    (w.x = 42),
+    w.x,
+    (function foo() {
+      w.x = 999;
+      return 2;
+    })(),
+  );
+}
+
+export const FIXTURE_ENTRYPOINT = {
+  fn: useTest,
+  params: [],
+};

compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/codegen-inline-iife.expect.md

@@ -0,0 +1,56 @@
+
+## Input
+
+```javascript
+import {makeArray, print} from 'shared-runtime';
+
+function useTest() {
+  return makeArray<number | void>(
+    print(1),
+    (function foo() {
+      print(2);
+      return 2;
+    })(),
+  );
+}
+
+export const FIXTURE_ENTRYPOINT = {
+  fn: useTest,
+  params: [],
+};
+
+```
+
+## Code
+
+```javascript
+import { c as _c } from "react/compiler-runtime";
+import { makeArray, print } from "shared-runtime";
+
+function useTest() {
+  const $ = _c(1);
+  let t0;
+  if ($[0] === Symbol.for("react.memo_cache_sentinel")) {
+    const t1 = print(1);
+    let t2;
+
+    print(2);
+    t2 = 2;
+    t0 = makeArray(t1, t2);
+    $[0] = t0;
+  } else {
+    t0 = $[0];
+  }
+  return t0;
+}
+
+export const FIXTURE_ENTRYPOINT = {
+  fn: useTest,
+  params: [],
+};
+
+```
+      
+### Eval output
+(kind: ok) [null,2]
+logs: [1,2]

compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/codegen-inline-iife.ts

@@ -0,0 +1,16 @@
+import {makeArray, print} from 'shared-runtime';
+
+function useTest() {
+  return makeArray<number | void>(
+    print(1),
+    (function foo() {
+      print(2);
+      return 2;
+    })(),
+  );
+}
+
+export const FIXTURE_ENTRYPOINT = {
+  fn: useTest,
+  params: [],
+};