Force DisjointSet.union to always pick a root

@josephsavona had the intuition that we were picking the wrong root in the 
previous infinite loop test case, so the fix for this is to force a root to 
always be picked. This works because `find` implements path compression (if a <- 
b and b <- c, then we can just point a <- c to "flatten" the tree which makes 
subsequent `find` operations more efficient since we don't have to follow the 
ancestor chain each time), so we're forcing all those unions to pick one parent. 

From some googling it looks like the traditional way to implement union is to 
call `find` so this should be the "right" way to fix it (?). 

I'm also adding a basic unit test for DisjointSet, I think we could revisit 
later and see if property testing is worth it but for now I mainly wanted to 
capture the regression test as a unit test.

compiler/forget/src/Utils/DisjointSet.ts

@@ -24,7 +24,7 @@ export default class DisjointSet<T> {
     // determine an arbitrary "root" for this set: if the first
     // item already has a root then use that, otherwise the first item
     // will be the new root.
-    let root = this.#entries.get(first);
+    let root = this.find(first);
     if (root == null) {
       root = first;
       this.#entries.set(first, first);

compiler/forget/src/__tests__/DisjointSet-test.ts

@@ -0,0 +1,112 @@
+import DisjointSet from "../Utils/DisjointSet";
+
+type TestIdentifier = {
+  id: number;
+  name: string;
+};
+
+describe("DisjointSet", () => {
+  let identifierId = 0;
+  function makeIdentifier(name: string): TestIdentifier {
+    return {
+      id: identifierId++,
+      name,
+    };
+  }
+
+  function makeIdentifiers(...names: string[]): TestIdentifier[] {
+    return names.map((name) => makeIdentifier(name));
+  }
+
+  beforeEach(() => {
+    identifierId = 0;
+  });
+
+  it(".find - finds the correct group which the item is associated with", () => {
+    const identifiers = new DisjointSet<TestIdentifier>();
+    const [x, y, z] = makeIdentifiers("x", "y", "z");
+
+    identifiers.union([x]);
+    identifiers.union([y, x]);
+
+    expect(identifiers.find(x)).toBe(y);
+    expect(identifiers.find(y)).toBe(y);
+    expect(identifiers.find(z)).toBe(null);
+  });
+
+  it(".size - returns 0 when empty", () => {
+    const identifiers = new DisjointSet<TestIdentifier>();
+
+    expect(identifiers.size).toBe(0);
+  });
+
+  it(".size - returns the correct size when non-empty", () => {
+    const identifiers = new DisjointSet<TestIdentifier>();
+    const [x, y] = makeIdentifiers("x", "y", "z");
+
+    identifiers.union([x]);
+    identifiers.union([y, x]);
+
+    expect(identifiers.size).toBe(2);
+  });
+
+  it(".buildSets - returns non-overlapping sets", () => {
+    const identifiers = new DisjointSet<TestIdentifier>();
+    const [a, b, c, x, y, z] = makeIdentifiers("a", "b", "c", "x", "y", "z");
+
+    identifiers.union([a]);
+    identifiers.union([b, a]);
+    identifiers.union([c, b]);
+
+    identifiers.union([x]);
+    identifiers.union([y, x]);
+    identifiers.union([z, y]);
+    identifiers.union([x, z]);
+
+    expect(identifiers.buildSets()).toMatchInlineSnapshot(`
+      [
+        Set {
+          {
+            "id": 0,
+            "name": "a",
+          },
+          {
+            "id": 1,
+            "name": "b",
+          },
+          {
+            "id": 2,
+            "name": "c",
+          },
+        },
+        Set {
+          {
+            "id": 3,
+            "name": "x",
+          },
+          {
+            "id": 4,
+            "name": "y",
+          },
+          {
+            "id": 5,
+            "name": "z",
+          },
+        },
+      ]
+    `);
+  });
+
+  // Regression test for issue #933
+  it("`forEach` doesn't infinite loop when there are cycles", () => {
+    const identifiers = new DisjointSet<TestIdentifier>();
+    const [x, y, z] = makeIdentifiers("x", "y", "z");
+
+    identifiers.union([x]);
+    identifiers.union([y, x]);
+    identifiers.union([z, y]);
+    identifiers.union([x, z]);
+
+    identifiers.forEach((_, group) => expect(group).toBe(z));
+  });
+});

compiler/forget/src/__tests__/fixtures/hir/issue933-disjoint-set-infinite-loop.expect.md

@@ -2,7 +2,7 @@
 ## Input
 
 ```javascript
-// This causes an infinite loop in the compiler
+// This caused an infinite loop in the compiler
 function MyApp(props) {
   const y = makeObj();
   const tmp = y.a;