From 8246956331b690452a3d3690c70f45b714f9683a Mon Sep 17 00:00:00 2001
From: Lauren Tan <poteto@users.noreply.github.com>
Date: Thu, 22 Dec 2022 16:16:16 -0500
Subject: [PATCH] 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      |   2 +-
 .../forget/src/__tests__/DisjointSet-test.ts  | 112 ++++++++++++++++++
 ...ue933-disjoint-set-infinite-loop.expect.md |   2 +-
 3 files changed, 114 insertions(+), 2 deletions(-)
 create mode 100644 compiler/forget/src/__tests__/DisjointSet-test.ts

diff --git a/compiler/forget/src/Utils/DisjointSet.ts b/compiler/forget/src/Utils/DisjointSet.ts
index c9c4bc56f3..717eb67f25 100644
--- a/compiler/forget/src/Utils/DisjointSet.ts
+++ b/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);
diff --git a/compiler/forget/src/__tests__/DisjointSet-test.ts b/compiler/forget/src/__tests__/DisjointSet-test.ts
new file mode 100644
index 0000000000..69a44789a3
--- /dev/null
+++ b/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));
+  });
+});
diff --git a/compiler/forget/src/__tests__/fixtures/hir/issue933-disjoint-set-infinite-loop.expect.md b/compiler/forget/src/__tests__/fixtures/hir/issue933-disjoint-set-infinite-loop.expect.md
index 0350ec3cf1..14d39faa8b 100644
--- a/compiler/forget/src/__tests__/fixtures/hir/issue933-disjoint-set-infinite-loop.expect.md
+++ b/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;