| 1 | /** |
| 2 | * Copyright (c) Meta Platforms, Inc. and affiliates. |
| 3 | * |
| 4 | * This source code is licensed under the MIT license found in the |
| 5 | * LICENSE file in the root directory of this source tree. |
| 6 | */ |
| 7 | |
| 8 | import {CompilerError} from '../CompilerError'; |
| 9 | import {GeneratedSource} from '../HIR/HIR'; |
| 10 | |
| 11 | // Represents items which form disjoint sets. |
| 12 | export default class DisjointSet<T> { |
| 13 | #entries: Map<T, T> = new Map(); |
| 14 | |
| 15 | /* |
| 16 | * Updates the graph to reflect that the given @param items form a set, |
| 17 | * linking any previous sets that the items were part of into a single |
| 18 | * set. |
| 19 | */ |
| 20 | union(items: Array<T>): void { |
| 21 | const first = items.shift(); |
| 22 | CompilerError.invariant(first != null, { |
| 23 | reason: 'Expected set to be non-empty', |
| 24 | loc: GeneratedSource, |
| 25 | }); |
| 26 | /* |
| 27 | * determine an arbitrary "root" for this set: if the first |
| 28 | * item already has a root then use that, otherwise the first item |
| 29 | * will be the new root. |
| 30 | */ |
| 31 | let root = this.find(first); |
| 32 | if (root == null) { |
| 33 | root = first; |
| 34 | this.#entries.set(first, first); |
| 35 | } |
| 36 | // update remaining items (which may already be part of other sets) |
| 37 | for (const item of items) { |
| 38 | let itemParent = this.#entries.get(item); |
| 39 | if (itemParent == null) { |
| 40 | // new item, no existing set to update |
| 41 | this.#entries.set(item, root); |
| 42 | continue; |
| 43 | } else if (itemParent === root) { |
| 44 | continue; |
| 45 | } else { |
| 46 | let current = item; |
| 47 | while (itemParent !== root) { |
| 48 | this.#entries.set(current, root); |
| 49 | current = itemParent; |
| 50 | itemParent = this.#entries.get(current)!; |
| 51 | } |
| 52 | } |
| 53 | } |
| 54 | } |
| 55 | |
| 56 | /* |
| 57 | * Finds the set to which the given @param item is associated, if @param item |
| 58 | * is present in this set. If item is not present, returns null. |
| 59 | * |
| 60 | * Note that the returned value may be any item in the set to which the input |
| 61 | * belongs: the only guarantee is that all items in a set will return the same |
| 62 | * value in between calls to `union()`. |
| 63 | */ |
| 64 | find(item: T): T | null { |
| 65 | if (!this.#entries.has(item)) { |
| 66 | return null; |
| 67 | } |
| 68 | const parent = this.#entries.get(item)!; |
| 69 | if (parent === item) { |
| 70 | // this is the root element |
| 71 | return item; |
| 72 | } |
| 73 | // Recurse to find the root (caching all elements along the path to the root) |
| 74 | const root = this.find(parent)!; |
| 75 | // Cache the element itself |
| 76 | this.#entries.set(item, root); |
| 77 | return root; |
| 78 | } |
| 79 | |
| 80 | has(item: T): boolean { |
| 81 | return this.#entries.has(item); |
| 82 | } |
| 83 | |
| 84 | /* |
| 85 | * Forces the set into canonical form, ie with all items pointing directly to |
| 86 | * their root, and returns a Map representing the mapping of items to their roots. |
| 87 | */ |
| 88 | canonicalize(): Map<T, T> { |
| 89 | const entries = new Map<T, T>(); |
| 90 | for (const item of this.#entries.keys()) { |
| 91 | const root = this.find(item)!; |
| 92 | entries.set(item, root); |
| 93 | } |
| 94 | return entries; |
| 95 | } |
| 96 | |
| 97 | /* |
| 98 | * Calls the provided callback once for each item in the disjoint set, |
| 99 | * passing the @param item and the @param group to which it belongs. |
| 100 | */ |
| 101 | forEach(fn: (item: T, group: T) => void): void { |
| 102 | for (const item of this.#entries.keys()) { |
| 103 | const group = this.find(item)!; |
| 104 | fn(item, group); |
| 105 | } |
| 106 | } |
| 107 | |
| 108 | buildSets(): Array<Set<T>> { |
| 109 | const ids: Map<T, number> = new Map(); |
| 110 | const sets: Map<number, Set<T>> = new Map(); |
| 111 | |
| 112 | this.forEach((identifier, groupIdentifier) => { |
| 113 | let id = ids.get(groupIdentifier); |
| 114 | if (id == null) { |
| 115 | id = ids.size; |
| 116 | ids.set(groupIdentifier, id); |
| 117 | } |
| 118 | |
| 119 | let set = sets.get(id); |
| 120 | if (set === undefined) { |
| 121 | set = new Set(); |
| 122 | sets.set(id, set); |
| 123 | } |
| 124 | set.add(identifier); |
| 125 | }); |
| 126 | |
| 127 | return [...sets.values()]; |
| 128 | } |
| 129 | |
| 130 | get size(): number { |
| 131 | return this.#entries.size; |
| 132 | } |
| 133 | } |