Skip to main content

logicaffeine_base/
union_find.rs

1//! Union-Find (disjoint set) with path compression and union by rank.
2//!
3//! Shared by the kernel's congruence closure (`logicaffeine_kernel::cc`)
4//! and the compiler's equality-saturation e-graph — one equivalence engine
5//! underneath both the proof system and the optimizer.
6
7/// Union-Find over `usize` element ids.
8///
9/// Supports near-constant amortized time for `find` (with path compression)
10/// and `union` (by rank).
11pub struct UnionFind {
12    /// Parent pointer for each element (element is its own parent if root).
13    parent: Vec<usize>,
14    /// Rank (approximate tree depth) for union by rank optimization.
15    rank: Vec<usize>,
16}
17
18impl UnionFind {
19    pub fn new() -> Self {
20        UnionFind {
21            parent: Vec::new(),
22            rank: Vec::new(),
23        }
24    }
25
26    /// Add a new element, returns its ID.
27    pub fn make_set(&mut self) -> usize {
28        let id = self.parent.len();
29        self.parent.push(id);
30        self.rank.push(0);
31        id
32    }
33
34    /// Number of elements ever created (not the number of classes).
35    pub fn len(&self) -> usize {
36        self.parent.len()
37    }
38
39    pub fn is_empty(&self) -> bool {
40        self.parent.is_empty()
41    }
42
43    /// Find representative with path compression.
44    pub fn find(&mut self, x: usize) -> usize {
45        if self.parent[x] != x {
46            self.parent[x] = self.find(self.parent[x]);
47        }
48        self.parent[x]
49    }
50
51    /// Union by rank, returns true if a merge occurred.
52    pub fn union(&mut self, x: usize, y: usize) -> bool {
53        let rx = self.find(x);
54        let ry = self.find(y);
55        if rx == ry {
56            return false;
57        }
58
59        if self.rank[rx] < self.rank[ry] {
60            self.parent[rx] = ry;
61        } else if self.rank[rx] > self.rank[ry] {
62            self.parent[ry] = rx;
63        } else {
64            self.parent[ry] = rx;
65            self.rank[rx] += 1;
66        }
67        true
68    }
69}
70
71impl Default for UnionFind {
72    fn default() -> Self {
73        Self::new()
74    }
75}