Class DisjointDict

mypy/checker.py:9315–9407

An variation of the union-find algorithm/data structure where instead of keeping track of just disjoint sets, we keep track of disjoint dicts -- keep track of multiple Set[Key] -> Set[Value] mappings, where each mapping's keys are guaranteed to be disjoint. This data structure is curren

class DisjointDict(Generic[TKey, TValue]):
    """An variation of the union-find algorithm/data structure where instead of keeping
    track of just disjoint sets, we keep track of disjoint dicts -- keep track of multiple
    Set[Key] -> Set[Value] mappings, where each mapping's keys are guaranteed to be disjoint.

    This data structure is currently used exclusively by 'group_comparison_operands' below
    to merge chains of '==' and 'is' comparisons when two or more chains use the same expression
    in best-case O(n), where n is the number of operands.

    Specifically, the `add_mapping()` function and `items()` functions will take on average
    O(k + v) and O(n) respectively, where k and v are the number of keys and values we're adding
    for a given chain. Note that k <= n and v <= n.

    We hit these average/best-case scenarios for most user code: e.g. when the user has just
    a single chain like 'a == b == c == d == ...' or multiple disjoint chains like
    'a==b < c==d < e==f < ...'. (Note that a naive iterative merging would be O(n^2) for
    the latter case).

    In comparison, this data structure will make 'group_comparison_operands' have a worst-case
    runtime of O(n*log(n)): 'add_mapping()' and 'items()' are worst-case O(k*log(n) + v) and
    O(k*log(n)) respectively. This happens only in the rare case where the user keeps repeatedly
    making disjoint mappings before merging them in a way that persistently dodges the path
    compression optimization in '_lookup_root_id', which would end up constructing a single
    tree of height log_2(n). This makes root lookups no longer amoritized constant time when we
    finally call 'items()'.
    """

    def __init__(self) -> None:
        # Each key maps to a unique ID
        self._key_to_id: dict[TKey, int] = {}

        # Each id points to the parent id, forming a forest of upwards-pointing trees. If the
        # current id already is the root, it points to itself. We gradually flatten these trees
        # as we perform root lookups: eventually all nodes point directly to its root.
        self._id_to_parent_id: dict[int, int] = {}

        # Each root id in turn maps to the set of values.
        self._root_id_to_values: dict[int, set[TValue]] = {}

    def add_mapping(self, keys: set[TKey], values: set[TValue]) -> None:
        """Adds a 'Set[TKey] -> Set[TValue]' mapping. If there already exists a mapping
        containing one or more of the given keys, we merge the input mapping with the old one.

        Note that the given set of keys must be non-empty -- otherwise, nothing happens.
        """
        if not keys:
            return

        subtree_roots = [self._lookup_or_make_root_id(key) for key in keys]
        new_root = subtree_roots[0]

        root_values = self._root_id_to_values[new_root]
        root_values.update(values)
        for subtree_root in subtree_roots[1:]:
            if subtree_root == new_root or subtree_root not in self._root_id_to_values:
                continue
            self._id_to_parent_id[subtree_root] = new_root
            root_values.update(self._root_id_to_values.pop(subtree_root))