compiler: super early exploration of instruction reordering

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

@@ -41,6 +41,7 @@ import {
   deadCodeElimination,
   pruneMaybeThrows,
 } from "../Optimization";
+import { instructionReordering } from "../Optimization/InstructionReordering";
 import {
   CodegenFunction,
   alignObjectMethodScopes,
@@ -177,6 +178,9 @@ function* runWithEnvironment(
   inferTypes(hir);
   yield log({ kind: "hir", name: "InferTypes", value: hir });
 
+  instructionReordering(hir);
+  yield log({ kind: "hir", name: "InstructionReordering", value: hir });
+
   if (env.config.validateHooksUsage) {
     validateHooksUsage(hir);
   }

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

@@ -0,0 +1,224 @@
+import {
+  BasicBlock,
+  HIRFunction,
+  IdentifierId,
+  Instruction,
+  markInstructionIds,
+} from "../HIR";
+import { printInstruction } from "../HIR/PrintHIR";
+import {
+  eachInstructionValueLValue,
+  eachInstructionValueOperand,
+  eachTerminalOperand,
+} from "../HIR/visitors";
+import { getOrInsertDefault } from "../Utils/utils";
+
+/**
+ * WIP early exploration of instruction reordering. This is a fairly aggressive form and has
+ * some issues. The idea of what's implemented:
+ *
+ * The high-level approach is to build a dependency graph where nodes generally correspond
+ * either to instructions OR to particular lvalue assignments of an expresssion. So
+ * `Destructure [x, y] = z` creates 3 nodes: one for the instruction, and one each for x and y.
+ * The lvalue nodes depend on the instruction node that assigns them.
+ *
+ * We add dependency edges for all the rvalues/lvalues of each instruction. In addition, we
+ * implicitly add dependencies btw non-reorderable instructions (more on that criteria) to
+ * serialize any instruction where order might be observable.
+ *
+ * We then distinguish two types of instructions that are reorderable:
+ * - Primitives, JSXText, JSX elements, and globals can be *globally* reordered, ie across blocks.
+ *   We defer emitting them until they are first used globally.
+ * - Array and object expressions are reorderable within basic blocks. This could likely be relaxed to be global.
+ * - StoreLocal, LoadLocal, and Destructure are reorderable within basic blocks. However, we serialize all
+ *   references to each named variable (reads and writes) to ensure that we aren't changing the order of evaluation
+ *   of variable references.
+ *
+ * The variable reordering relies on the fact that any variables that could be reassigned via a function expression
+ * are promoted to "context" variables and use LoadContext/StoreContext, which are not reorderable.
+ *
+ * In theory it might even be safe to do this variable reordering globally, but i want to think through that more.
+ *
+ * With the above context, the algorithm is approximately:
+ * - For each basic block:
+ *   - Iterate the instructions to create the dependency graph
+ *   - Re-emit instructions, "pulling" from all the values that are depended upon by the block's terminal.
+ *   - Emit any remaining instructions that cannot be globally reordered, starting from later instructions first.
+ *   - Save any globally-reorderable instructions into a global map that is shared across blocks, so they can be
+ *     emitted by the first block that needs them.
+ *
+ * Emitting instructions is currently naive: we just iterate in the order that the dependencies were established.
+ * If instruction 4 depends on instructions 1, 2, and 3, we'll visit in depth-first order and emit 1, 2, 3, 4.
+ * That's true even if instruction 1 and 2 are simple instructions (for ex primitives) while instruction 3 has its
+ * own large dependency tree.
+ *
+ * ## Issues/things to explore:
+ *
+ * - An obvious improvement is to weight the nodes and emit dependencies based on weight. Alternatively, we could try to
+ * determine the reactive dependencies of each node, and try to emit nodes that have the same dependencies together.
+ * - Reordering destructure statements means that we also end up deferring the evaluation of its RHS. So i noticed some
+ *   `const [state, setState] = useState(...)` getting moved around. But i think i might have just messed up the bit that
+ *   ensures non-reorderable instructions (like the useState() call here) are serialized. So this should just be a simple fix,
+ *   if i didn't already fix it (need to go back through the fixture output changes)
+ * - I also noticed that destructuring being moved meant that some reactive scopes ended up with less precise input, because
+ *   the destructure moved into the reactive scope itself (so the scope depends on the rvalue of the destructure, not the lvalues).
+ *   This is weird, i need to debug.
+ * - Probably more things.
+ */
+export function instructionReordering(fn: HIRFunction): void {
+  const globalDependencies: Dependencies = new Map();
+  for (const [, block] of fn.body.blocks) {
+    reorderBlock(block, globalDependencies);
+  }
+  markInstructionIds(fn.body);
+}
+
+type Dependencies = Map<IdentifierId, Node>;
+type Node = {
+  instruction: Instruction | null;
+  dependencies: Array<IdentifierId>;
+};
+
+function reorderBlock(
+  block: BasicBlock,
+  globalDependencies: Dependencies
+): void {
+  const dependencies: Dependencies = new Map();
+  const locals = new Map<string, IdentifierId>();
+  let previousIdentifier: IdentifierId | null = null;
+  for (const instr of block.instructions) {
+    const node: Node = getOrInsertDefault(
+      dependencies,
+      instr.lvalue.identifier.id,
+      {
+        instruction: instr,
+        dependencies: [],
+      }
+    );
+    if (getReorderingLevel(instr) === ReorderingLevel.None) {
+      if (previousIdentifier !== null) {
+        node.dependencies.push(previousIdentifier);
+      }
+      previousIdentifier = instr.lvalue.identifier.id;
+    }
+    for (const operand of eachInstructionValueOperand(instr.value)) {
+      if (
+        operand.identifier.name !== null &&
+        operand.identifier.name.kind === "named"
+      ) {
+        const previous = locals.get(operand.identifier.name.value);
+        if (previous !== undefined) {
+          node.dependencies.push(previous);
+        } else {
+          locals.set(operand.identifier.name.value, instr.lvalue.identifier.id);
+          node.dependencies.push(operand.identifier.id);
+        }
+      } else {
+        if (dependencies.has(operand.identifier.id)) {
+          node.dependencies.push(operand.identifier.id);
+        }
+      }
+    }
+    dependencies.set(instr.lvalue.identifier.id, node);
+
+    for (const lvalue of eachInstructionValueLValue(instr.value)) {
+      const lvalueNode = getOrInsertDefault(
+        dependencies,
+        lvalue.identifier.id,
+        {
+          instruction: null,
+          dependencies: [],
+        }
+      );
+      lvalueNode.dependencies.push(instr.lvalue.identifier.id);
+      if (
+        lvalue.identifier.name !== null &&
+        lvalue.identifier.name.kind === "named"
+      ) {
+        const previous = locals.get(lvalue.identifier.name.value);
+        if (previous !== undefined) {
+          node.dependencies.push(previous);
+        }
+      }
+    }
+  }
+
+  const instructions: Array<Instruction> = [];
+
+  function emit(id: IdentifierId): void {
+    const node = dependencies.get(id) ?? globalDependencies.get(id);
+    if (node == null) {
+      return;
+    }
+    dependencies.delete(id);
+    globalDependencies.delete(id);
+    for (const dep of node.dependencies) {
+      emit(dep);
+    }
+    if (node.instruction !== null) {
+      instructions.push(node.instruction);
+    }
+  }
+
+  for (const operand of eachTerminalOperand(block.terminal)) {
+    emit(operand.identifier.id);
+  }
+  for (const id of Array.from(dependencies.keys()).reverse()) {
+    const node = dependencies.get(id);
+    if (node == null) {
+      continue;
+    }
+    if (
+      node.instruction !== null &&
+      getReorderingLevel(node.instruction) === ReorderingLevel.Global
+    ) {
+      globalDependencies.set(id, node);
+    } else {
+      emit(id);
+    }
+  }
+  block.instructions = instructions;
+}
+
+function printDeps(deps: Dependencies): string {
+  return (
+    "[\n" +
+    Array.from(deps)
+      .map(
+        ([id, dep]) =>
+          `$${id} ${
+            dep.instruction != null ? printInstruction(dep.instruction) : ""
+          } deps=[${dep.dependencies.map((x) => `$${x}`).join(", ")}]`
+      )
+      .join("\n") +
+    "\n]"
+  );
+}
+
+enum ReorderingLevel {
+  None = "none",
+  Local = "local",
+  Global = "global",
+}
+function getReorderingLevel(instr: Instruction): ReorderingLevel {
+  switch (instr.value.kind) {
+    case "JsxExpression":
+    case "JsxFragment":
+    case "JSXText":
+    case "LoadGlobal":
+    case "Primitive":
+    case "TemplateLiteral": {
+      return ReorderingLevel.Global;
+    }
+    case "ArrayExpression":
+    case "ObjectExpression":
+    case "LoadLocal":
+    case "Destructure":
+    case "StoreLocal": {
+      return ReorderingLevel.Local;
+    }
+    default: {
+      return ReorderingLevel.None;
+    }
+  }
+}

compiler/packages/babel-plugin-react-compiler/src/Utils/utils.ts

@@ -30,7 +30,9 @@ export function assertExhaustive(_: never, errorMsg: string): never {
   throw new Error(errorMsg);
 }
 
-// Modifies @param array in place, retaining only the items where the predicate returns true.
+/**
+ * Modifies @param array in place, retaining only the items where the predicate returns true.
+ */
 export function retainWhere<T>(
   array: Array<T>,
   predicate: (item: T) => boolean

compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/alias-capture-in-method-receiver-and-mutate.expect.md

@@ -34,13 +34,12 @@ function Component() {
   const $ = _c(1);
   let t0;
   if ($[0] === Symbol.for("react.memo_cache_sentinel")) {
+    const x = [];
     const a = makeObject_Primitives();
 
-    const x = [];
+    t0 = [x, a];
     x.push(a);
-
     mutate(x);
-    t0 = [x, a];
     $[0] = t0;
   } else {
     t0 = $[0];

compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/alias-capture-in-method-receiver.expect.md

@@ -20,26 +20,19 @@ function Component() {
 ```javascript
 import { c as _c } from "react/compiler-runtime";
 function Component() {
-  const $ = _c(2);
+  const $ = _c(1);
   let t0;
   if ($[0] === Symbol.for("react.memo_cache_sentinel")) {
-    t0 = someObj();
-    $[0] = t0;
-  } else {
-    t0 = $[0];
-  }
-  const a = t0;
-  let t1;
-  if ($[1] === Symbol.for("react.memo_cache_sentinel")) {
     const x = [];
-    x.push(a);
+    const a = someObj();
 
-    t1 = [x, a];
-    $[1] = t1;
+    t0 = [x, a];
+    x.push(a);
+    $[0] = t0;
   } else {
-    t1 = $[1];
+    t0 = $[0];
   }
-  return t1;
+  return t0;
 }
 
 ```

compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/alias-computed-load.expect.md

@@ -19,19 +19,20 @@ function component(a) {
 import { c as _c } from "react/compiler-runtime";
 function component(a) {
   const $ = _c(2);
-  let x;
+  let t0;
   if ($[0] !== a) {
-    x = { a };
-    const y = {};
+    const x = { a };
 
+    t0 = x;
+    const y = {};
     y.x = x.a;
     mutate(y);
     $[0] = a;
-    $[1] = x;
+    $[1] = t0;
   } else {
-    x = $[1];
+    t0 = $[1];
   }
-  return x;
+  return t0;
 }
 
 ```

compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/alias-nested-member-path-mutate.expect.md

@@ -20,19 +20,21 @@ function component() {
 import { c as _c } from "react/compiler-runtime";
 function component() {
   const $ = _c(1);
-  let x;
+  let t0;
   if ($[0] === Symbol.for("react.memo_cache_sentinel")) {
-    const z = [];
+    const x = {};
+
+    t0 = x;
     const y = {};
+    const z = [];
     y.z = z;
-    x = {};
     x.y = y;
     mutate(x.y.z);
-    $[0] = x;
+    $[0] = t0;
   } else {
-    x = $[0];
+    t0 = $[0];
   }
-  return x;
+  return t0;
 }
 
 ```

compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/alias-nested-member-path.expect.md

@@ -25,18 +25,20 @@ export const FIXTURE_ENTRYPOINT = {
 import { c as _c } from "react/compiler-runtime";
 function component() {
   const $ = _c(1);
-  let x;
+  let t0;
   if ($[0] === Symbol.for("react.memo_cache_sentinel")) {
-    const z = [];
+    const x = {};
+
+    t0 = x;
     const y = {};
+    const z = [];
     y.z = z;
-    x = {};
     x.y = y;
-    $[0] = x;
+    $[0] = t0;
   } else {
-    x = $[0];
+    t0 = $[0];
   }
-  return x;
+  return t0;
 }
 
 export const FIXTURE_ENTRYPOINT = {