metrale_grammar/fsm/
union_find.rs

1// SPDX-License-Identifier: MIT OR Apache-2.0
2//
3// UnionFindSet — a disjoint-set forest keyed by `i32` state ids.
4// Port of `cpp/support/union_find_set.h`, specialized to `i32` (the FSM
5// only ever unions state ids). Uses path compression + union-by-size.
6
7use ahash::AHashMap;
8
9/// A disjoint-set (union-find) structure over `i32` elements.
10///
11/// Elements must be `add`ed before they can be `find`/`union`-ed; this
12/// matches the C++ semantics where membership is explicit.
13#[derive(Debug, Default)]
14pub struct UnionFindSet {
15    /// Maps element -> (parent, subtree size).
16    parent_size: AHashMap<i32, (i32, usize)>,
17}
18
19impl UnionFindSet {
20    /// Create an empty union-find set.
21    pub fn new() -> Self {
22        Self {
23            parent_size: AHashMap::new(),
24        }
25    }
26
27    /// Add `element` as its own singleton set. Returns `false` if it
28    /// was already present.
29    pub fn add(&mut self, element: i32) -> bool {
30        if self.parent_size.contains_key(&element) {
31            return false;
32        }
33        self.parent_size.insert(element, (element, 1));
34        true
35    }
36
37    /// Remove every element.
38    pub fn clear(&mut self) {
39        self.parent_size.clear();
40    }
41
42    /// True if `element` has been added.
43    pub fn contains(&self, element: i32) -> bool {
44        self.parent_size.contains_key(&element)
45    }
46
47    /// Find the representative of `element`'s set, compressing the path.
48    /// Panics if `element` was never added (matches C++ `XGRAMMAR_CHECK`).
49    pub fn find(&mut self, element: i32) -> i32 {
50        let parent = self
51            .parent_size
52            .get(&element)
53            .expect("Element not found in union-find set.")
54            .0;
55        if parent != element {
56            let root = self.find(parent);
57            self.parent_size.get_mut(&element).unwrap().0 = root;
58            root
59        } else {
60            element
61        }
62    }
63
64    /// Merge the sets containing `a` and `b`. Both must have been added.
65    pub fn union(&mut self, a: i32, b: i32) {
66        let mut root_a = self.find(a);
67        let mut root_b = self.find(b);
68        if root_a == root_b {
69            return;
70        }
71        let size_a = self.parent_size[&root_a].1;
72        let size_b = self.parent_size[&root_b].1;
73        // Keep root_a as the larger set (union by size).
74        if size_a < size_b {
75            std::mem::swap(&mut root_a, &mut root_b);
76        }
77        self.parent_size.get_mut(&root_b).unwrap().0 = root_a;
78        let merged = self.parent_size[&root_a].1 + self.parent_size[&root_b].1;
79        self.parent_size.get_mut(&root_a).unwrap().1 = merged;
80    }
81
82    /// Collect all sets, each sorted ascending, ordered by their first
83    /// element — deterministic, matching the C++ `GetAllSets`.
84    pub fn all_sets(&mut self) -> Vec<Vec<i32>> {
85        let elements: Vec<i32> = self.parent_size.keys().copied().collect();
86        let mut root_to_idx: AHashMap<i32, usize> = AHashMap::new();
87        let mut result: Vec<Vec<i32>> = Vec::new();
88        for value in elements {
89            let root = self.find(value);
90            let idx = *root_to_idx.entry(root).or_insert_with(|| {
91                result.push(Vec::new());
92                result.len() - 1
93            });
94            result[idx].push(value);
95        }
96        for vec in &mut result {
97            vec.sort_unstable();
98        }
99        result.sort_by(|v1, v2| v1[0].cmp(&v2[0]));
100        result
101    }
102}
103
104#[cfg(test)]
105mod tests {
106    use super::*;
107
108    #[test]
109    fn add_is_idempotent() {
110        let mut uf = UnionFindSet::new();
111        assert!(uf.add(1));
112        assert!(!uf.add(1));
113        assert!(uf.contains(1));
114        assert!(!uf.contains(2));
115    }
116
117    #[test]
118    fn union_merges_sets() {
119        let mut uf = UnionFindSet::new();
120        for i in 0..4 {
121            uf.add(i);
122        }
123        uf.union(0, 1);
124        uf.union(2, 3);
125        assert_eq!(uf.find(0), uf.find(1));
126        assert_eq!(uf.find(2), uf.find(3));
127        assert_ne!(uf.find(0), uf.find(2));
128        uf.union(1, 2);
129        assert_eq!(uf.find(0), uf.find(3));
130    }
131
132    #[test]
133    fn all_sets_deterministic() {
134        let mut uf = UnionFindSet::new();
135        for i in 0..6 {
136            uf.add(i);
137        }
138        uf.union(4, 1);
139        uf.union(5, 2);
140        let sets = uf.all_sets();
141        // Sets: {1,4}, {2,5}, {0}, {3} — ordered by first element.
142        assert_eq!(sets, vec![vec![0], vec![1, 4], vec![2, 5], vec![3]]);
143    }
144
145    #[test]
146    fn clear_empties() {
147        let mut uf = UnionFindSet::new();
148        uf.add(7);
149        uf.clear();
150        assert!(!uf.contains(7));
151    }
152
153    #[test]
154    #[should_panic(expected = "not found")]
155    fn find_missing_panics() {
156        let mut uf = UnionFindSet::new();
157        uf.find(99);
158    }
159}
metrale_grammar/fsm/union_find.rs

metrale_grammar/fsm/
union_find.rs
