[patch][dce] Patch dce to have separate mark and sweep phases

--- Previously, our logic was something like: ```js fixed-point-loop { foreach instruction { mark referenced identifiers // assume that usages are always visited before declarations if (instruction is decl) { prune(instruction); } } foreach instruction { if not referenced { delete(instruction); } } ``` This contained a bug, as not all usages of a variable are guaranteed to be visited before its declaration. ```js // input let x = 0; while(x < 10) { x += 2; } return x; // hir entry: x$0 = 0 goto loop-test loop-test: x$1 = phi(x$0, x$2) if ... goto loop-body else goto fallthrough loop-body: x$2 = x$1 ... goto loop-test fallthrough: return x$1 ``` In this example,`x$2` is defined by `loop-body` and used by `loop-test`. Similarly, `x$1` is defined by `loop-test` and used by `loop-body`. --- TODO: trying to come up with more test fixtures
2025-11-01 09:12:30 +00:00 · 2024-01-18 18:29:19 -05:00
parent a0e90065c6
commit 4aa60d32b9
3 changed files with 186 additions and 67 deletions
@@ -29,62 +29,17 @@ import { assertExhaustive, retainWhere } from "../Utils/utils";
 * Note that unreachable blocks are already pruned during HIR construction.
 */
 export function deadCodeElimination(fn: HIRFunction): void {
-  const state = new State();
-
-  /*
-   * If there are no back-edges the algorithm can terminate after a single iteration
-   * of the blocks
+  /**
+   * Phase 1: Find/mark all referenced identifiers
+   * Usages may be visited AFTER declarations if there are circular phi / data dependencies
+   * between blocks, so we wait to sweep until after fixed point iteration is complete
   */
-  const hasLoop = hasBackEdge(fn);
+  const state = findReferencedIdentifiers(fn);

-  const reversedBlocks = [...fn.body.blocks.values()].reverse();
-  let size = state.count;
-  do {
-    size = state.count;
-
-    /*
-     * Iterate blocks in postorder (successors before predecessors, excepting loops)
-     * to find usages before declarations
-     */
-    for (const block of reversedBlocks) {
-      for (const operand of eachTerminalOperand(block.terminal)) {
-        state.reference(operand.identifier);
-      }
-
-      for (let i = block.instructions.length - 1; i >= 0; i--) {
-        const instr = block.instructions[i]!;
-        if (
-          !state.isIdOrNameUsed(instr.lvalue.identifier) &&
-          pruneableValue(instr.value, state) &&
-          // Can't prune the last value of a value block, that's its value!
-          !(block.kind !== "block" && i === block.instructions.length - 1)
-        ) {
-          continue;
-        }
-        state.reference(instr.lvalue.identifier);
-
-        /*
-         * For the last value of a value block, if it's not pruneable we can't
-         * rewrite it. This is necessary to preserve unused value blocks
-         */
-        if (block.kind !== "block" && i === block.instructions.length - 1) {
-          for (const place of eachInstructionValueOperand(instr.value)) {
-            state.reference(place.identifier);
-          }
-          continue;
-        }
-        // Otherwise rewrite instructions to remove unused parts of them
-        visitInstruction(instr, state);
-      }
-      for (const phi of block.phis) {
-        if (state.isIdOrNameUsed(phi.id)) {
-          for (const [_pred, operand] of phi.operands) {
-            state.reference(operand);
-          }
-        }
-      }
-    }
-  } while (state.count > size && hasLoop);
+  /**
+   * Phase 2: Prune / sweep unreferenced identifiers and instructions
+   * as possible (subject to HIR structural constraints)
+   */
  for (const [, block] of fn.body.blocks) {
    for (const phi of block.phis) {
      if (!state.isIdOrNameUsed(phi.id)) {
@@ -94,6 +49,14 @@ export function deadCodeElimination(fn: HIRFunction): void {
    retainWhere(block.instructions, (instr) =>
      state.isIdOrNameUsed(instr.lvalue.identifier)
    );
+    // Rewrite retained instructions
+    for (let i = 0; i < block.instructions.length; i++) {
+      const isBlockValue =
+        block.kind !== "block" && i === block.instructions.length - 1;
+      if (!isBlockValue) {
+        rewriteInstruction(block.instructions[i], state);
+      }
+    }
  }
 }

@@ -134,10 +97,81 @@ class State {
  }
 }

-function visitInstruction(instr: Instruction, state: State): void {
+function findReferencedIdentifiers(fn: HIRFunction): State {
+  /*
+   * If there are no back-edges the algorithm can terminate after a single iteration
+   * of the blocks
+   */
+  const hasLoop = hasBackEdge(fn);
+  const reversedBlocks = [...fn.body.blocks.values()].reverse();
+
+  const state = new State();
+  let size = state.count;
+  do {
+    size = state.count;
+
+    /*
+     * Iterate blocks in postorder (successors before predecessors, excepting loops)
+     * to visit usages before declarations
+     */
+    for (const block of reversedBlocks) {
+      for (const operand of eachTerminalOperand(block.terminal)) {
+        state.reference(operand.identifier);
+      }
+
+      for (let i = block.instructions.length - 1; i >= 0; i--) {
+        const instr = block.instructions[i]!;
+        const isBlockValue =
+          block.kind !== "block" && i === block.instructions.length - 1;
+
+        if (isBlockValue) {
+          /**
+           * The last instr of a value block is never eligible for pruning,
+           * as that's the block's value. Pessimistically consider all operands
+           * as used to avoid rewriting the last instruction
+           */
+          state.reference(instr.lvalue.identifier);
+          for (const place of eachInstructionValueOperand(instr.value)) {
+            state.reference(place.identifier);
+          }
+        } else if (
+          state.isIdOrNameUsed(instr.lvalue.identifier) ||
+          !pruneableValue(instr.value, state)
+        ) {
+          state.reference(instr.lvalue.identifier);
+
+          if (instr.value.kind === "StoreLocal") {
+            /*
+             * If this is a Let/Const declaration, mark the initializer as referenced
+             * only if the ssa'ed lval is also referenced
+             */
+            if (
+              instr.value.lvalue.kind === InstructionKind.Reassign ||
+              state.isIdUsed(instr.value.lvalue.place.identifier)
+            ) {
+              state.reference(instr.value.value.identifier);
+            }
+          } else {
+            for (const operand of eachInstructionValueOperand(instr.value)) {
+              state.reference(operand.identifier);
+            }
+          }
+        }
+      }
+      for (const phi of block.phis) {
+        if (state.isIdOrNameUsed(phi.id)) {
+          for (const [_pred, operand] of phi.operands) {
+            state.reference(operand);
+          }
+        }
+      }
+    }
+  } while (state.count > size && hasLoop);
+  return state;
+}
+
+function rewriteInstruction(instr: Instruction, state: State): void {
  if (instr.value.kind === "Destructure") {
-    // Mark the value as used, not the lvalues
-    state.reference(instr.value.value.identifier);
    // Remove unused lvalues
    switch (instr.value.lvalue.pattern.kind) {
      case "ArrayPattern": {
@@ -219,16 +253,6 @@ function visitInstruction(instr: Instruction, state: State): void {
        lvalue: instr.value.lvalue,
        loc: instr.value.loc,
      };
-    } else {
-      /*
-       * Else we mark the initializer as referenced, since the variable itself is
-       * referenced
-       */
-      state.reference(instr.value.value.identifier);
-    }
-  } else {
-    for (const operand of eachInstructionValueOperand(instr.value)) {
-      state.reference(operand.identifier);
    }
  }
 }
@@ -0,0 +1,72 @@
+
+## Input
+
+```javascript
+import { identity } from "shared-runtime";
+
+function Component({ data }) {
+  let x = 0;
+  for (const item of data) {
+    const { current, other } = item;
+    x += current;
+    identity(other);
+  }
+  return [x];
+}
+
+export const FIXTURE_ENTRYPOINT = {
+  fn: Component,
+  params: [
+    {
+      data: [
+        { current: 2, other: 3 },
+        { current: 4, other: 5 },
+      ],
+    },
+  ],
+};
+
+```
+
+## Code
+
+```javascript
+import { unstable_useMemoCache as useMemoCache } from "react";
+import { identity } from "shared-runtime";
+
+function Component(t25) {
+  const $ = useMemoCache(2);
+  const { data } = t25;
+  let x = 0;
+  for (const item of data) {
+    const { current, other } = item;
+    x = x + current;
+    identity(other);
+  }
+  let t0;
+  if ($[0] !== x) {
+    t0 = [x];
+    $[0] = x;
+    $[1] = t0;
+  } else {
+    t0 = $[1];
+  }
+  return t0;
+}
+
+export const FIXTURE_ENTRYPOINT = {
+  fn: Component,
+  params: [
+    {
+      data: [
+        { current: 2, other: 3 },
+        { current: 4, other: 5 },
+      ],
+    },
+  ],
+};
+
+```
+      
+### Eval output
+(kind: ok) [6]
@@ -0,0 +1,23 @@
+import { identity } from "shared-runtime";
+
+function Component({ data }) {
+  let x = 0;
+  for (const item of data) {
+    const { current, other } = item;
+    x += current;
+    identity(other);
+  }
+  return [x];
+}
+
+export const FIXTURE_ENTRYPOINT = {
+  fn: Component,
+  params: [
+    {
+      data: [
+        { current: 2, other: 3 },
+        { current: 4, other: 5 },
+      ],
+    },
+  ],
+};