cranelift_egraph/
unionfind.rs

1//! Simple union-find data structure.
2
3use crate::{trace, Id};
4use cranelift_entity::SecondaryMap;
5use std::hash::{Hash, Hasher};
6
7/// A union-find data structure. The data structure can allocate
8/// `Id`s, indicating eclasses, and can merge eclasses together.
9#[derive(Clone, Debug)]
10pub struct UnionFind {
11    parent: SecondaryMap<Id, Id>,
12}
13
14impl UnionFind {
15    /// Create a new `UnionFind`.
16    pub fn new() -> Self {
17        UnionFind {
18            parent: SecondaryMap::new(),
19        }
20    }
21
22    /// Create a new `UnionFind` with the given capacity.
23    pub fn with_capacity(cap: usize) -> Self {
24        UnionFind {
25            parent: SecondaryMap::with_capacity(cap),
26        }
27    }
28
29    /// Add an `Id` to the `UnionFind`, with its own equivalence class
30    /// initially. All `Id`s must be added before being queried or
31    /// unioned.
32    pub fn add(&mut self, id: Id) {
33        self.parent[id] = id;
34    }
35
36    /// Find the canonical `Id` of a given `Id`.
37    pub fn find(&self, mut node: Id) -> Id {
38        while node != self.parent[node] {
39            node = self.parent[node];
40        }
41        node
42    }
43
44    /// Find the canonical `Id` of a given `Id`, updating the data
45    /// structure in the process so that future queries for this `Id`
46    /// (and others in its chain up to the root of the equivalence
47    /// class) will be faster.
48    pub fn find_and_update(&mut self, mut node: Id) -> Id {
49        // "Path splitting" mutating find (Tarjan and Van Leeuwen).
50        let orig = node;
51        while node != self.parent[node] {
52            let next = self.parent[self.parent[node]];
53            self.parent[node] = next;
54            node = next;
55        }
56        trace!("find_and_update: {} -> {}", orig, node);
57        node
58    }
59
60    /// Merge the equivalence classes of the two `Id`s.
61    pub fn union(&mut self, a: Id, b: Id) {
62        let a = self.find_and_update(a);
63        let b = self.find_and_update(b);
64        let (a, b) = (std::cmp::min(a, b), std::cmp::max(a, b));
65        if a != b {
66            // Always canonicalize toward lower IDs.
67            self.parent[b] = a;
68            trace!("union: {}, {}", a, b);
69        }
70    }
71
72    /// Determine if two `Id`s are equivalent, after
73    /// canonicalizing. Update union-find data structure during our
74    /// canonicalization to make future lookups faster.
75    pub fn equiv_id_mut(&mut self, a: Id, b: Id) -> bool {
76        self.find_and_update(a) == self.find_and_update(b)
77    }
78
79    /// Hash an `Id` after canonicalizing it. Update union-find data
80    /// structure to make future lookups/hashing faster.
81    pub fn hash_id_mut<H: Hasher>(&mut self, hash: &mut H, id: Id) {
82        let id = self.find_and_update(id);
83        id.hash(hash);
84    }
85}
cranelift_egraph/unionfind.rs

cranelift_egraph/
unionfind.rs
