rustc_mir_transform/coverage/counters/
union_find.rs

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
use std::cmp::Ordering;
use std::mem;

use rustc_index::{Idx, IndexVec};

#[cfg(test)]
mod tests;

/// Simple implementation of a union-find data structure, i.e. a disjoint-set
/// forest.
#[derive(Debug)]
pub(crate) struct UnionFind<Key: Idx> {
    table: IndexVec<Key, UnionFindEntry<Key>>,
}

#[derive(Debug)]
struct UnionFindEntry<Key> {
    /// Transitively points towards the "root" of the set containing this key.
    ///
    /// Invariant: A root key is its own parent.
    parent: Key,
    /// When merging two "root" keys, their ranks determine which key becomes
    /// the new root, to prevent the parent tree from becoming unnecessarily
    /// tall. See [`UnionFind::unify`] for details.
    rank: u32,
}

impl<Key: Idx> UnionFind<Key> {
    /// Creates a new disjoint-set forest containing the keys `0..num_keys`.
    /// Initially, every key is part of its own one-element set.
    pub(crate) fn new(num_keys: usize) -> Self {
        // Initially, every key is the root of its own set, so its parent is itself.
        Self { table: IndexVec::from_fn_n(|key| UnionFindEntry { parent: key, rank: 0 }, num_keys) }
    }

    /// Returns the "root" key of the disjoint-set containing the given key.
    /// If two keys have the same root, they belong to the same set.
    ///
    /// Also updates internal data structures to make subsequent `find`
    /// operations faster.
    pub(crate) fn find(&mut self, key: Key) -> Key {
        // Loop until we find a key that is its own parent.
        let mut curr = key;
        while let parent = self.table[curr].parent
            && curr != parent
        {
            // Perform "path compression" by peeking one layer ahead, and
            // setting the current key's parent to that value.
            // (This works even when `parent` is the root of its set, because
            // of the invariant that a root is its own parent.)
            let parent_parent = self.table[parent].parent;
            self.table[curr].parent = parent_parent;

            // Advance by one step and continue.
            curr = parent;
        }
        curr
    }

    /// Merges the set containing `a` and the set containing `b` into one set.
    ///
    /// Returns the common root of both keys, after the merge.
    pub(crate) fn unify(&mut self, a: Key, b: Key) -> Key {
        let mut a = self.find(a);
        let mut b = self.find(b);

        // If both keys have the same root, they're already in the same set,
        // so there's nothing more to do.
        if a == b {
            return a;
        };

        // Ensure that `a` has strictly greater rank, swapping if necessary.
        // If both keys have the same rank, increment the rank of `a` so that
        // future unifications will also prefer `a`, leading to flatter trees.
        match Ord::cmp(&self.table[a].rank, &self.table[b].rank) {
            Ordering::Less => mem::swap(&mut a, &mut b),
            Ordering::Equal => self.table[a].rank += 1,
            Ordering::Greater => {}
        }

        debug_assert!(self.table[a].rank > self.table[b].rank);
        debug_assert_eq!(self.table[b].parent, b);

        // Make `a` the parent of `b`.
        self.table[b].parent = a;

        a
    }

    /// Creates a snapshot of this disjoint-set forest that can no longer be
    /// mutated, but can be queried without mutation.
    pub(crate) fn freeze(&mut self) -> FrozenUnionFind<Key> {
        // Just resolve each key to its actual root.
        let roots = self.table.indices().map(|key| self.find(key)).collect();
        FrozenUnionFind { roots }
    }
}

/// Snapshot of a disjoint-set forest that can no longer be mutated, but can be
/// queried in O(1) time without mutation.
///
/// This is really just a wrapper around a direct mapping from keys to roots,
/// but with a [`Self::find`] method that resembles [`UnionFind::find`].
#[derive(Debug)]
pub(crate) struct FrozenUnionFind<Key: Idx> {
    roots: IndexVec<Key, Key>,
}

impl<Key: Idx> FrozenUnionFind<Key> {
    /// Returns the "root" key of the disjoint-set containing the given key.
    /// If two keys have the same root, they belong to the same set.
    pub(crate) fn find(&self, key: Key) -> Key {
        self.roots[key]
    }
}