DisjointSet Class — react Architecture

Source Code

compiler/packages/babel-plugin-react-compiler/src/Utils/DisjointSet.ts lines 12–133

export default class DisjointSet<T> {
  #entries: Map<T, T> = new Map();

  /*
   * Updates the graph to reflect that the given @param items form a set,
   * linking any previous sets that the items were part of into a single
   * set.
   */
  union(items: Array<T>): void {
    const first = items.shift();
    CompilerError.invariant(first != null, {
      reason: 'Expected set to be non-empty',
      loc: GeneratedSource,
    });
    /*
     * 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.find(first);
    if (root == null) {
      root = first;
      this.#entries.set(first, first);
    }
    // update remaining items (which may already be part of other sets)
    for (const item of items) {
      let itemParent = this.#entries.get(item);
      if (itemParent == null) {
        // new item, no existing set to update
        this.#entries.set(item, root);
        continue;
      } else if (itemParent === root) {
        continue;
      } else {
        let current = item;
        while (itemParent !== root) {
          this.#entries.set(current, root);
          current = itemParent;
          itemParent = this.#entries.get(current)!;
        }
      }
    }
  }

  /*
   * Finds the set to which the given @param item is associated, if @param item
   * is present in this set. If item is not present, returns null.
   *
   * Note that the returned value may be any item in the set to which the input
   * belongs: the only guarantee is that all items in a set will return the same
   * value in between calls to `union()`.
   */
  find(item: T): T | null {
    if (!this.#entries.has(item)) {
      return null;
    }
    const parent = this.#entries.get(item)!;
    if (parent === item) {
      // this is the root element
      return item;
    }
    // Recurse to find the root (caching all elements along the path to the root)
    const root = this.find(parent)!;
    // Cache the element itself
    this.#entries.set(item, root);
    return root;
  }

  has(item: T): boolean {
    return this.#entries.has(item);
  }

  /*
   * Forces the set into canonical form, ie with all items pointing directly to
   * their root, and returns a Map representing the mapping of items to their roots.
   */
  canonicalize(): Map<T, T> {
    const entries = new Map<T, T>();
    for (const item of this.#entries.keys()) {
      const root = this.find(item)!;
      entries.set(item, root);