rustc_mir_transform/coverage/counters/
union_find.rs

Help
1use std::cmp::Ordering;
2use std::mem;
3
4use rustc_index::{Idx, IndexVec};
5
6#[cfg(test)]
7mod tests;
8
9/// Simple implementation of a union-find data structure, i.e. a disjoint-set
10/// forest.
11#[derive(Debug)]
12pub(crate) struct UnionFind<Key: Idx> {
13    table: IndexVec<Key, UnionFindEntry<Key>>,
14}
15
16#[derive(Debug)]
17struct UnionFindEntry<Key> {
18    /// Transitively points towards the "root" of the set containing this key.
19    ///
20    /// Invariant: A root key is its own parent.
21    parent: Key,
22    /// When merging two "root" keys, their ranks determine which key becomes
23    /// the new root, to prevent the parent tree from becoming unnecessarily
24    /// tall. See [`UnionFind::unify`] for details.
25    rank: u32,
26}
27
28impl<Key: Idx> UnionFind<Key> {
29    /// Creates a new disjoint-set forest containing the keys `0..num_keys`.
30    /// Initially, every key is part of its own one-element set.
31    pub(crate) fn new(num_keys: usize) -> Self {
32        // Initially, every key is the root of its own set, so its parent is itself.
33        Self { table: IndexVec::from_fn_n(|key| UnionFindEntry { parent: key, rank: 0 }, num_keys) }
34    }
35
36    /// Returns the "root" key of the disjoint-set containing the given key.
37    /// If two keys have the same root, they belong to the same set.
38    ///
39    /// Also updates internal data structures to make subsequent `find`
40    /// operations faster.
41    pub(crate) fn find(&mut self, key: Key) -> Key {
42        // Loop until we find a key that is its own parent.
43        let mut curr = key;
44        while let parent = self.table[curr].parent
45            && curr != parent
46        {
47            // Perform "path compression" by peeking one layer ahead, and
48            // setting the current key's parent to that value.
49            // (This works even when `parent` is the root of its set, because
50            // of the invariant that a root is its own parent.)
51            let parent_parent = self.table[parent].parent;
52            self.table[curr].parent = parent_parent;
53
54            // Advance by one step and continue.
55            curr = parent;
56        }
57        curr
58    }
59
60    /// Merges the set containing `a` and the set containing `b` into one set.
61    ///
62    /// Returns the common root of both keys, after the merge.
63    pub(crate) fn unify(&mut self, a: Key, b: Key) -> Key {
64        let mut a = self.find(a);
65        let mut b = self.find(b);
66
67        // If both keys have the same root, they're already in the same set,
68        // so there's nothing more to do.
69        if a == b {
70            return a;
71        };
72
73        // Ensure that `a` has strictly greater rank, swapping if necessary.
74        // If both keys have the same rank, increment the rank of `a` so that
75        // future unifications will also prefer `a`, leading to flatter trees.
76        match Ord::cmp(&self.table[a].rank, &self.table[b].rank) {
77            Ordering::Less => mem::swap(&mut a, &mut b),
78            Ordering::Equal => self.table[a].rank += 1,
79            Ordering::Greater => {}
80        }
81
82        debug_assert!(self.table[a].rank > self.table[b].rank);
83        debug_assert_eq!(self.table[b].parent, b);
84
85        // Make `a` the parent of `b`.
86        self.table[b].parent = a;
87
88        a
89    }
90
91    /// Takes a "snapshot" of the current state of this disjoint-set forest, in
92    /// the form of a vector that directly maps each key to its current root.
93    pub(crate) fn snapshot(&mut self) -> IndexVec<Key, Key> {
94        self.table.indices().map(|key| self.find(key)).collect()
95    }
96}
rustc_mir_transform/coverage/counters/union_find.rs

rustc_mir_transform/coverage/counters/
union_find.rs
