Implement HIR visitors (#747)

compiler/forget/src/HIR/EliminateRedundantPhi.ts

@@ -7,8 +7,8 @@
 
 import invariant from "invariant";
 import { assertExhaustive } from "../Common/utils";
-import { BlockId, HIRFunction, Identifier, Place } from "./HIR";
-import { eachInstructionOperand } from "./HIRBuilder";
+import { BlockId, HIRFunction, Identifier, Place, Terminal } from "./HIR";
+import { eachInstructionOperand, eachTerminalOperand } from "./visitors";
 
 /**
  * Pass to eliminate redundant phi nodes:
@@ -91,38 +91,8 @@ export function eliminateRedundantPhi(fn: HIRFunction) {
 
       // Rewrite all terminal operands
       const { terminal } = block;
-      switch (terminal.kind) {
-        case "if": {
-          rewritePlace(terminal.test, rewrites);
-          break;
-        }
-        case "switch": {
-          rewritePlace(terminal.test, rewrites);
-          for (const case_ of terminal.cases) {
-            if (case_.test === null) {
-              continue;
-            }
-            rewritePlace(case_.test, rewrites);
-          }
-          break;
-        }
-        case "return":
-        case "throw": {
-          if (terminal.value !== null) {
-            rewritePlace(terminal.value, rewrites);
-          }
-          break;
-        }
-        case "goto": {
-          // no-op
-          break;
-        }
-        default: {
-          assertExhaustive(
-            terminal,
-            `Unexpected terminal kind '${(terminal as any).kind}'`
-          );
-        }
+      for (const place of eachTerminalOperand(terminal)) {
+        rewritePlace(place, rewrites);
       }
     }
     // We only need to loop if there were newly eliminated phis in this iteration

compiler/forget/src/HIR/HIRBuilder.ts

@@ -19,6 +19,7 @@ import {
   Terminal,
 } from "./HIR";
 import { printInstruction } from "./PrintHIR";
+import { mapTerminalSuccessors } from "./visitors";
 
 // *******************************************************************************************
 // *******************************************************************************************
@@ -553,127 +554,3 @@ function getTargetIfIndirection(block: BasicBlock): number | null {
     ? block.terminal.block
     : null;
 }
-
-/**
- * Maps a terminal node's block assignments using the provided function.
- *
- * TODO: this visits successors in reverse ordering to facilitate shrink()'s
- * goal of producing a reverse postorder graph where siblings are in-order.
- */
-export function mapTerminalSuccessors(
-  terminal: Terminal,
-  fn: (block: BlockId, isFallthrough: boolean) => BlockId
-): Terminal {
-  switch (terminal.kind) {
-    case "goto": {
-      const target = fn(terminal.block, false);
-      return {
-        kind: "goto",
-        block: target,
-      };
-    }
-    case "if": {
-      const consequent = fn(terminal.consequent, false);
-      const alternate = fn(terminal.alternate, false);
-      const fallthrough =
-        terminal.fallthrough !== null ? fn(terminal.fallthrough, true) : null;
-      return {
-        kind: "if",
-        test: terminal.test,
-        consequent,
-        alternate,
-        fallthrough,
-      };
-    }
-    case "switch": {
-      const cases = terminal.cases.map((case_) => {
-        const target = fn(case_.block, false);
-        return {
-          test: case_.test,
-          block: target,
-        };
-      });
-      const fallthrough =
-        terminal.fallthrough !== null ? fn(terminal.fallthrough, true) : null;
-      return {
-        kind: "switch",
-        test: terminal.test,
-        cases,
-        fallthrough,
-      };
-    }
-    case "return": {
-      return {
-        kind: "return",
-        value: terminal.value,
-      };
-    }
-    case "throw": {
-      return terminal;
-    }
-    default: {
-      assertExhaustive(
-        terminal,
-        `Unexpected terminal kind '${(terminal as any as Terminal).kind}'`
-      );
-    }
-  }
-}
-
-export function* eachInstructionOperand(instr: Instruction) {
-  const instrValue = instr.value;
-  switch (instrValue.kind) {
-    case "NewExpression":
-    case "CallExpression": {
-      yield instrValue.callee;
-      yield* instrValue.args;
-      break;
-    }
-    case "BinaryExpression": {
-      yield instrValue.left;
-      yield instrValue.right;
-      break;
-    }
-    case "Identifier": {
-      yield instrValue;
-      break;
-    }
-    case "UnaryExpression": {
-      yield instrValue.value;
-      break;
-    }
-    case "JsxExpression": {
-      yield instrValue.tag;
-      yield* instrValue.props.values();
-      if (instrValue.children) {
-        yield* instrValue.children;
-      }
-      break;
-    }
-    case "JsxFragment": {
-      yield* instrValue.children;
-      break;
-    }
-    case "ObjectExpression": {
-      if (instrValue.properties !== null) {
-        yield* instrValue.properties.values();
-      }
-      break;
-    }
-    case "ArrayExpression": {
-      yield* instrValue.elements;
-      break;
-    }
-    case "OtherStatement":
-    case "Primitive":
-    case "JSXText": {
-      break;
-    }
-    default: {
-      assertExhaustive(
-        instrValue,
-        `Unexpected instruction kind '${(instrValue as any).kind}'`
-      );
-    }
-  }
-}

compiler/forget/src/HIR/InferMutableLifetimes.ts

@@ -7,7 +7,7 @@
 
 import { assertExhaustive } from "../Common/utils";
 import { Effect, HIRFunction, Instruction, Place } from "./HIR";
-import { eachInstructionOperand } from "./HIRBuilder";
+import { eachInstructionOperand } from "./visitors";
 import { printInstruction, printPlace } from "./PrintHIR";
 
 /**

compiler/forget/src/HIR/InferReferenceEffects.ts

@@ -19,7 +19,11 @@ import {
   Terminal,
   ValueKind,
 } from "./HIR";
-import { mapTerminalSuccessors } from "./HIRBuilder";
+import {
+  eachInstructionOperand,
+  eachTerminalOperand,
+  mapTerminalSuccessors,
+} from "./visitors";
 import { printMixedHIR, printPlace, printSourceLocation } from "./PrintHIR";
 
 /**
@@ -503,57 +507,44 @@ function inferBlock(env: Environment, block: BasicBlock) {
 
   for (const instr of block.instructions) {
     const instrValue = instr.value;
+    let effectKind: Effect | null = null;
     let valueKind: ValueKind;
     switch (instrValue.kind) {
       case "BinaryExpression": {
         valueKind = ValueKind.Immutable;
-        env.reference(instrValue.left, Effect.Read);
-        env.reference(instrValue.right, Effect.Read);
+        effectKind = Effect.Read;
         break;
       }
       case "ArrayExpression": {
         valueKind = ValueKind.Mutable;
-        for (const element of instrValue.elements) {
-          env.reference(element, Effect.Read);
-        }
+        effectKind = Effect.Read;
         break;
       }
       case "NewExpression": {
         valueKind = ValueKind.Mutable;
-        env.reference(instrValue.callee, Effect.Mutate);
-        for (const arg of instrValue.args) {
-          env.reference(arg, Effect.Mutate);
-        }
+        effectKind = Effect.Mutate;
         break;
       }
       case "CallExpression": {
-        let effectKind = Effect.Mutate;
         valueKind = ValueKind.Mutable;
+        effectKind = Effect.Mutate;
         const hook = parseHookCall(instrValue.callee);
         if (hook !== null) {
           effectKind = hook.effectKind;
           valueKind = hook.valueKind;
         }
-        env.reference(instrValue.callee, effectKind);
-        for (const arg of instrValue.args) {
-          env.reference(arg, effectKind);
-        }
         break;
       }
       case "ObjectExpression": {
         valueKind = ValueKind.Mutable;
         // Object construction captures but does not modify the key/property values
-        if (instrValue.properties !== null) {
-          for (const [_key, value] of instrValue.properties) {
-            env.reference(value, Effect.Read);
-          }
-        }
+        effectKind = Effect.Read;
         break;
       }
       case "UnaryExpression": {
         // TODO check that value must be a primitive, or make conditional based on the operator
         valueKind = ValueKind.Immutable;
-        env.reference(instrValue.value, Effect.Read);
+        effectKind = Effect.Read;
         break;
       }
       case "OtherStatement": {
@@ -563,22 +554,12 @@ function inferBlock(env: Environment, block: BasicBlock) {
       }
       case "JsxExpression": {
         valueKind = ValueKind.Frozen;
-        env.reference(instrValue.tag, Effect.Freeze);
-        for (const [_prop, value] of instrValue.props) {
-          env.reference(value, Effect.Freeze);
-        }
-        if (instrValue.children !== null) {
-          for (const child of instrValue.children) {
-            env.reference(child, Effect.Freeze);
-          }
-        }
+        effectKind = Effect.Freeze;
         break;
       }
       case "JsxFragment": {
         valueKind = ValueKind.Frozen;
-        for (const child of instrValue.children) {
-          env.reference(child, Effect.Freeze);
-        }
+        effectKind = Effect.Freeze;
         break;
       }
       case "JSXText":
@@ -615,6 +596,16 @@ function inferBlock(env: Environment, block: BasicBlock) {
         assertExhaustive(instrValue, "Unexpected instruction kind");
       }
     }
+
+    for (const operand of eachInstructionOperand(instr)) {
+      invariant(
+        effectKind != null,
+        "effectKind must be set for instruction value `%s`",
+        instrValue.kind
+      );
+      env.reference(operand, effectKind);
+    }
+
     env.initialize(instrValue, valueKind);
     if (instr.lvalue !== null) {
       if (instr.lvalue.place.memberPath === null) {
@@ -625,39 +616,13 @@ function inferBlock(env: Environment, block: BasicBlock) {
       instr.lvalue.place.effect = Effect.Mutate;
     }
   }
-  switch (block.terminal.kind) {
-    case "throw": {
-      env.reference(block.terminal.value, Effect.Freeze);
-      break;
-    }
-    case "return": {
-      if (block.terminal.value !== null) {
-        env.reference(block.terminal.value, Effect.Freeze);
-      }
-      break;
-    }
-    case "if": {
-      env.reference(block.terminal.test, Effect.Read);
-      break;
-    }
-    case "switch": {
-      env.reference(block.terminal.test, Effect.Read);
-      for (const case_ of block.terminal.cases) {
-        if (case_.test !== null) {
-          env.reference(case_.test, Effect.Read);
-        }
-      }
-      break;
-    }
-    case "goto": {
-      break;
-    }
-    default: {
-      assertExhaustive(
-        block.terminal,
-        `Unexpected terminal kind '${(block.terminal as any as Terminal).kind}'`
-      );
-    }
+
+  const effect =
+    block.terminal.kind === "return" || block.terminal.kind === "throw"
+      ? Effect.Freeze
+      : Effect.Read;
+  for (const operand of eachTerminalOperand(block.terminal)) {
+    env.reference(operand, effect);
   }
 }
 

compiler/forget/src/HIR/visitors.ts

@@ -0,0 +1,166 @@
+/**
+ * Copyright (c) Meta Platforms, Inc. and affiliates.
+ *
+ * This source code is licensed under the MIT license found in the
+ * LICENSE file in the root directory of this source tree.
+ */
+
+import { assertExhaustive } from "../Common/utils";
+import { BlockId, Instruction, Place, Terminal } from "./HIR";
+
+export function* eachInstructionOperand(instr: Instruction): Iterable<Place> {
+  const instrValue = instr.value;
+  switch (instrValue.kind) {
+    case "NewExpression":
+    case "CallExpression": {
+      yield instrValue.callee;
+      yield* instrValue.args;
+      break;
+    }
+    case "BinaryExpression": {
+      yield instrValue.left;
+      yield instrValue.right;
+      break;
+    }
+    case "Identifier": {
+      yield instrValue;
+      break;
+    }
+    case "UnaryExpression": {
+      yield instrValue.value;
+      break;
+    }
+    case "JsxExpression": {
+      yield instrValue.tag;
+      yield* instrValue.props.values();
+      if (instrValue.children) {
+        yield* instrValue.children;
+      }
+      break;
+    }
+    case "JsxFragment": {
+      yield* instrValue.children;
+      break;
+    }
+    case "ObjectExpression": {
+      if (instrValue.properties !== null) {
+        yield* instrValue.properties.values();
+      }
+      break;
+    }
+    case "ArrayExpression": {
+      yield* instrValue.elements;
+      break;
+    }
+    case "OtherStatement":
+    case "Primitive":
+    case "JSXText": {
+      break;
+    }
+    default: {
+      assertExhaustive(
+        instrValue,
+        `Unexpected instruction kind '${(instrValue as any).kind}'`
+      );
+    }
+  }
+}
+
+/**
+ * Maps a terminal node's block assignments using the provided function.
+ */
+export function mapTerminalSuccessors(
+  terminal: Terminal,
+  fn: (block: BlockId, isFallthrough: boolean) => BlockId
+): Terminal {
+  switch (terminal.kind) {
+    case "goto": {
+      const target = fn(terminal.block, false);
+      return {
+        kind: "goto",
+        block: target,
+      };
+    }
+    case "if": {
+      const consequent = fn(terminal.consequent, false);
+      const alternate = fn(terminal.alternate, false);
+      const fallthrough =
+        terminal.fallthrough !== null ? fn(terminal.fallthrough, true) : null;
+      return {
+        kind: "if",
+        test: terminal.test,
+        consequent,
+        alternate,
+        fallthrough,
+      };
+    }
+    case "switch": {
+      const cases = terminal.cases.map((case_) => {
+        const target = fn(case_.block, false);
+        return {
+          test: case_.test,
+          block: target,
+        };
+      });
+      const fallthrough =
+        terminal.fallthrough !== null ? fn(terminal.fallthrough, true) : null;
+      return {
+        kind: "switch",
+        test: terminal.test,
+        cases,
+        fallthrough,
+      };
+    }
+    case "return": {
+      return {
+        kind: "return",
+        value: terminal.value,
+      };
+    }
+    case "throw": {
+      return terminal;
+    }
+    default: {
+      assertExhaustive(
+        terminal,
+        `Unexpected terminal kind '${(terminal as any as Terminal).kind}'`
+      );
+    }
+  }
+}
+
+export function* eachTerminalOperand(terminal: Terminal): Iterable<Place> {
+  switch (terminal.kind) {
+    case "if": {
+      yield terminal.test;
+      break;
+    }
+    case "switch": {
+      yield terminal.test;
+      for (const case_ of terminal.cases) {
+        if (case_.test === null) {
+          continue;
+        }
+        yield case_.test;
+      }
+      break;
+    }
+    case "return":
+    case "throw": {
+      if (terminal.value !== null) {
+        yield terminal.value;
+      }
+      break;
+    }
+    case "goto": {
+      // no-op
+      break;
+    }
+    default: {
+      assertExhaustive(
+        terminal,
+        `Unexpected terminal kind '${(terminal as any).kind}'`
+      );
+    }
+  }
+}