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}