[hir] Add a buildAliasSets pass

Moves the existing alias set building logic from inferMutableLifetimes to a 
separate pass. 

Additionally maintain abstract state to refine aliasing to not include 
primitives.

compiler/forget/src/HIR/BuildAliasSets.ts

@@ -0,0 +1,102 @@
+import invariant from "invariant";
+import DisjointSet from "./DisjointSet";
+import { HIRFunction, Identifier, Instruction, Place, LValue } from "./HIR";
+import { printInstructionValue } from "./PrintHIR";
+
+type AbstractValue = AbstractObject | AbstractPrimitive;
+type AbstractObject = {
+  kind: "Object";
+  values: Map<string, AbstractValue>;
+};
+type AbstractPrimitive = {
+  kind: "Primitive";
+  value: number | boolean | string | null | undefined;
+};
+
+class AbstractState {
+  aliases = new DisjointSet<Identifier>();
+  #values = new Map<Identifier, AbstractValue>();
+
+  // Simple lvalue:
+  //   lvalue = alias;
+  //   lvalue = alias.memberPath;
+  alias(lvalue: LValue, alias: Place) {
+    // Simple alias:
+    //    lvalue = alias;
+    if (alias.memberPath === null) {
+      let value = this.#values.get(alias.identifier);
+
+      // Don't know what this, let's default to an Object conservatively.
+      if (value === undefined) {
+        value = { kind: "Object", values: new Map() };
+      }
+
+      this.#values.set(lvalue.place.identifier, value);
+
+      // No need to alias Primitives
+      if (value.kind !== "Primitive") {
+        this.aliases.union([lvalue.place.identifier, alias.identifier]);
+      }
+    }
+  }
+
+  buildAliasSets(): Array<Set<Identifier>> {
+    const aliasIds: Map<Identifier, number> = new Map();
+    const aliasSets: Map<number, Set<Identifier>> = new Map();
+
+    this.aliases.forEach((identifier, groupIdentifier) => {
+      let aliasId = aliasIds.get(groupIdentifier);
+      if (aliasId == null) {
+        aliasId = aliasIds.size;
+        aliasIds.set(groupIdentifier, aliasId);
+      }
+
+      let aliasSet = aliasSets.get(aliasId);
+      if (aliasSet === undefined) {
+        aliasSet = new Set();
+        aliasSets.set(aliasId, aliasSet);
+      }
+      aliasSet.add(identifier);
+    });
+
+    return [...aliasSets.values()];
+  }
+}
+
+export function buildAliasSets(func: HIRFunction): Array<Set<Identifier>> {
+  const state = new AbstractState();
+  for (const [_, block] of func.body.blocks) {
+    for (const instr of block.instructions) {
+      inferInstr(instr, state);
+    }
+  }
+  return state.buildAliasSets();
+}
+
+function inferInstr(instr: Instruction, state: AbstractState) {
+  const { lvalue, value: instrValue } = instr;
+  let alias: Place | null = null;
+  switch (instrValue.kind) {
+    case "Identifier": {
+      alias = instrValue;
+      break;
+    }
+    default:
+      return;
+  }
+
+  invariant(
+    alias !== null,
+    `expected ${printInstructionValue(instrValue)} to have an alias`
+  );
+
+  // TODO(gsn): handle this.
+  if (lvalue === null) {
+    return;
+  }
+
+  // simple aliasing
+  if (lvalue.place.memberPath === null) {
+    state.alias(lvalue, alias);
+  }
+}

compiler/forget/src/HIR/InferMutableLifetimes.ts

@@ -7,6 +7,7 @@
 
 import invariant from "invariant";
 import { assertExhaustive } from "../Common/utils";
+import { buildAliasSets } from "./BuildAliasSets";
 import DisjointSet from "./DisjointSet";
 import {
   Effect,
@@ -85,8 +86,6 @@ function inferPlace(place: Place, instr: Instruction) {
 }
 
 export function inferMutableRanges(func: HIRFunction) {
-  const aliases = new DisjointSet<Identifier>();
-
   for (const [_, block] of func.body.blocks) {
     for (const phi of block.phis) {
       let start = Number.MAX_SAFE_INTEGER;
@@ -111,20 +110,6 @@ export function inferMutableRanges(func: HIRFunction) {
       }
 
       if (instr.lvalue !== null) {
-        if (instr.value.kind === "Identifier") {
-          // TODO(gsn): Handle complex aliasing.
-          if (
-            instr.value.memberPath === null &&
-            instr.lvalue.place.memberPath === null
-          ) {
-            // direct aliasing: `a = b`;
-            aliases.union([
-              instr.lvalue.place.identifier,
-              instr.value.identifier,
-            ]);
-          }
-        }
-
         if (instr.lvalue.place.memberPath === null) {
           const lvalueId = instr.lvalue.place.identifier;
 
@@ -142,45 +127,22 @@ export function inferMutableRanges(func: HIRFunction) {
     }
   }
 
-  const aliasIds: Map<Identifier, number> = new Map();
-  // Store the mutable range and set of identifiers for each scope
-  const aliasIndentifiers: Map<
-    number,
-    { end: InstructionId; identifiers: Set<Identifier> }
-  > = new Map();
-
-  aliases.forEach((identifier, groupIdentifier) => {
-    let aliasId = aliasIds.get(groupIdentifier);
-    if (aliasId == null) {
-      aliasId = aliasIds.size;
-      aliasIds.set(groupIdentifier, aliasId);
-    }
-
-    let alias = aliasIndentifiers.get(aliasId);
-    if (alias === undefined) {
-      alias = {
-        end: identifier.mutableRange.end,
-        identifiers: new Set(),
-      };
-      aliasIndentifiers.set(aliasId, alias);
-    } else {
-      alias.end = makeInstructionId(
-        Math.max(alias.end, identifier.mutableRange.end)
-      );
-    }
-    alias.identifiers.add(identifier);
-  });
-
-  for (const [_, alias] of aliasIndentifiers) {
+  const aliasSets = buildAliasSets(func);
+  for (const aliasSet of aliasSets) {
     // Update mutableRange.end only if the identifiers have actually been
     // mutated.
-    const haveIdentifiersBeenMutated = [...alias.identifiers].some(
+    const haveIdentifiersBeenMutated = [...aliasSet].some(
       (id) => id.mutableRange.end > id.mutableRange.start
     );
 
     if (haveIdentifiersBeenMutated) {
-      for (const identifier of alias.identifiers) {
-        identifier.mutableRange.end = alias.end;
+      // Find final instruction which mutates this alias set.
+      const mutableRangeEnds = [...aliasSet].map((id) => id.mutableRange.end);
+      const maxMutableRangeEnd = Math.max(...mutableRangeEnds) as InstructionId;
+
+      // Update mutableRange.end for all aliases in this set.
+      for (const alias of aliasSet) {
+        alias.mutableRange.end = maxMutableRangeEnd;
       }
     }
   }

compiler/forget/src/HIR/PrintHIR.ts

@@ -175,7 +175,7 @@ export function printTerminal(terminal: Terminal): Array<string> | string {
   return value;
 }
 
-function printInstructionValue(instrValue: InstructionValue): string {
+export function printInstructionValue(instrValue: InstructionValue): string {
   let value = "";
   switch (instrValue.kind) {
     case "ArrayExpression": {