From e2e5e389af35ae9223a99acc09cd8da651240467 Mon Sep 17 00:00:00 2001
From: Sathya Gunasekaran <gsathya.ceg@gmail.com>
Date: Mon, 12 Dec 2022 20:35:23 +0000
Subject: [PATCH] [be] Move buildAliasSets to DisjointSet

---
 compiler/forget/src/HIR/DisjointSet.ts        | 22 +++++++++++++++++
 compiler/forget/src/HIR/InferAlias.ts         | 24 -------------------
 .../src/HIR/InferMutableRangesForAlias.ts     |  3 +--
 compiler/forget/src/HIR/PrintHIR.ts           |  3 +--
 4 files changed, 24 insertions(+), 28 deletions(-)

diff --git a/compiler/forget/src/HIR/DisjointSet.ts b/compiler/forget/src/HIR/DisjointSet.ts
index 6710777232..6b554dace7 100644
--- a/compiler/forget/src/HIR/DisjointSet.ts
+++ b/compiler/forget/src/HIR/DisjointSet.ts
@@ -84,6 +84,28 @@ export default class DisjointSet<T> {
     }
   }
 
+  buildSets(): Array<Set<T>> {
+    const ids: Map<T, number> = new Map();
+    const sets: Map<number, Set<T>> = new Map();
+
+    this.forEach((identifier, groupIdentifier) => {
+      let id = ids.get(groupIdentifier);
+      if (id == null) {
+        id = ids.size;
+        ids.set(groupIdentifier, id);
+      }
+
+      let set = sets.get(id);
+      if (set === undefined) {
+        set = new Set();
+        sets.set(id, set);
+      }
+      set.add(identifier);
+    });
+
+    return [...sets.values()];
+  }
+
   get size(): number {
     return this.#entries.size;
   }
diff --git a/compiler/forget/src/HIR/InferAlias.ts b/compiler/forget/src/HIR/InferAlias.ts
index f5fbee0b8f..274217c2ff 100644
--- a/compiler/forget/src/HIR/InferAlias.ts
+++ b/compiler/forget/src/HIR/InferAlias.ts
@@ -15,30 +15,6 @@ type AbstractPrimitive = {
 
 export type AliasSet = Set<Identifier>;
 
-export function buildAliasSets(
-  aliases: DisjointSet<Identifier>
-): Array<AliasSet> {
-  const aliasIds: Map<Identifier, number> = new Map();
-  const aliasSets: Map<number, Set<Identifier>> = new Map();
-
-  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()];
-}
-
 class AbstractState {
   aliases = new DisjointSet<Identifier>();
   // NOTE(gsn): Should this be a part of AbstractObject? No, because this has
diff --git a/compiler/forget/src/HIR/InferMutableRangesForAlias.ts b/compiler/forget/src/HIR/InferMutableRangesForAlias.ts
index de0c7855ad..023ea8cd42 100644
--- a/compiler/forget/src/HIR/InferMutableRangesForAlias.ts
+++ b/compiler/forget/src/HIR/InferMutableRangesForAlias.ts
@@ -1,9 +1,8 @@
 import DisjointSet from "./DisjointSet";
 import { Identifier, InstructionId } from "./HIR";
-import { buildAliasSets } from "./InferAlias";
 
 export function inferMutableRangesForAlias(aliases: DisjointSet<Identifier>) {
-  const aliasSets = buildAliasSets(aliases);
+  const aliasSets = aliases.buildSets();
   for (const aliasSet of aliasSets) {
     // Update mutableRange.end only if the identifiers have actually been
     // mutated.
diff --git a/compiler/forget/src/HIR/PrintHIR.ts b/compiler/forget/src/HIR/PrintHIR.ts
index 2063885c9b..8b69b9e7b5 100644
--- a/compiler/forget/src/HIR/PrintHIR.ts
+++ b/compiler/forget/src/HIR/PrintHIR.ts
@@ -23,7 +23,6 @@ import {
   SourceLocation,
   Terminal,
 } from "./HIR";
-import { buildAliasSets } from "./InferAlias";
 import { eachReactiveScope } from "./visitors";
 
 export type Options = {
@@ -335,7 +334,7 @@ export function printSourceLocation(loc: SourceLocation): string {
 }
 
 export function printAliases(aliases: DisjointSet<Identifier>): string {
-  const aliasSets = buildAliasSets(aliases);
+  const aliasSets = aliases.buildSets();
 
   const items = [];
   for (const aliasSet of aliasSets) {