toolbox_rs/
partition_id.rs

1use bincode::Decode;
2use bincode::Encode;
3use core::cmp::max;
4use std::{
5    fmt::Display,
6    hash::Hash,
7    ops::{BitAnd, BitOr},
8};
9
10/// represents the hiearchical partition id scheme. The root id has ID 1 and
11/// children are shifted to the left by one and plus 0/1. The parent child
12/// relationship can thus be queried in constant time.
13#[derive(Clone, Copy, Debug, Eq, Hash, Ord, PartialEq, PartialOrd, Encode, Decode)]
14pub struct PartitionID(pub u32);
15
16impl PartitionID {
17    /// Returns the root id
18    pub fn root() -> PartitionID {
19        PartitionID(1)
20    }
21
22    /// Returns the parent of a given ID.
23    /// Note that the parent of the root id is always 1
24    pub fn parent(&self) -> PartitionID {
25        let new_id = max(1, self.0 >> 1);
26        PartitionID::new(new_id)
27    }
28
29    pub fn parent_at_level(&self, level: u32) -> PartitionID {
30        let parent = self.0 & (0xffff_ffff ^ ((1 << level) - 1));
31        PartitionID::new(parent)
32    }
33
34    /// Returns a left-right ordered tuple of children for a given ID
35    pub fn children(&self) -> (PartitionID, PartitionID) {
36        let temp = self.0 << 1;
37        (PartitionID(temp), PartitionID(temp + 1))
38    }
39
40    /// Returns the left child of a ID
41    pub fn left_child(&self) -> PartitionID {
42        let temp = self.0 << 1;
43        PartitionID(temp)
44    }
45
46    /// Returns the right child of a ID
47    pub fn right_child(&self) -> PartitionID {
48        let temp = self.0 << 1;
49        PartitionID(temp + 1)
50    }
51
52    /// Transform ID to its left-most descendant k levels down
53    pub fn make_leftmost_descendant(&mut self, k: usize) {
54        self.0 <<= k;
55    }
56
57    /// Transform ID to its right-most descendant k levels down
58    pub fn make_rightmost_descendant(&mut self, k: usize) {
59        self.make_leftmost_descendant(k);
60        self.0 += (1 << k) - 1;
61    }
62
63    /// Transform the ID into its left child
64    pub fn make_left_child(&mut self) {
65        self.make_leftmost_descendant(1);
66    }
67
68    /// Transform the ID into its right child
69    pub fn make_right_child(&mut self) {
70        self.make_rightmost_descendant(1);
71    }
72
73    /// Returns a new PartitionID from an u32
74    pub fn new(id: u32) -> Self {
75        // the id scheme is designed in a way that the number of leading zeros is always odd
76        debug_assert!(id != 0);
77        PartitionID(id)
78    }
79
80    /// The level in this scheme is defined by the the number of leading zeroes.
81    pub fn level(&self) -> u8 {
82        // magic number 31 := 32 - 1, as 1 is the root's ID
83        (31 - self.0.leading_zeros()).try_into().unwrap()
84    }
85
86    /// Returns whether the ID id a left child
87    pub fn is_left_child(&self) -> bool {
88        self.0 % 2 == 0
89    }
90
91    /// Returns whether the ID id a right child
92    pub fn is_right_child(&self) -> bool {
93        self.0 % 2 == 1
94    }
95
96    // Returns the lowest common ancestor of this and the other ID
97    pub fn lowest_common_ancestor(&self, other: &PartitionID) -> PartitionID {
98        let mut left = *self;
99        let mut right = *other;
100
101        let left_level = left.level();
102        let right_level = right.level();
103
104        if left_level > right_level {
105            left.0 >>= left_level - right_level;
106        }
107        if right_level > left_level {
108            right.0 >>= right_level - left_level;
109        }
110
111        while left != right {
112            left = left.parent();
113            right = right.parent();
114        }
115        left
116    }
117
118    pub fn extract_bit(&self, index: usize) -> bool {
119        let mask = 1 << index;
120        mask & self.0 > 0
121    }
122}
123
124impl Display for PartitionID {
125    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
126        write!(f, "{}", self.0)
127    }
128}
129
130impl From<PartitionID> for usize {
131    fn from(s: PartitionID) -> usize {
132        s.0.try_into().unwrap()
133    }
134}
135
136impl core::fmt::Binary for PartitionID {
137    fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
138        let val = self.0;
139        core::fmt::Binary::fmt(&val, f) // delegate to u32's implementation
140    }
141}
142
143impl BitAnd for PartitionID {
144    type Output = Self;
145
146    fn bitand(self, rhs: Self) -> Self::Output {
147        Self(self.0 & rhs.0)
148    }
149}
150
151impl BitOr for PartitionID {
152    type Output = Self;
153
154    fn bitor(self, rhs: Self) -> Self::Output {
155        Self(self.0 | rhs.0)
156    }
157}
158
159#[cfg(test)]
160mod tests {
161
162    use crate::partition_id::PartitionID;
163
164    #[test]
165    fn parent_id() {
166        let id = PartitionID::new(4);
167        assert_eq!(id.parent(), PartitionID::new(2));
168    }
169
170    #[test]
171    fn new_id() {
172        let id = PartitionID::new(1);
173        assert_eq!(id.parent(), PartitionID::root());
174    }
175
176    #[test]
177    fn children_ids() {
178        let id = PartitionID::new(0b0101_0101_0101_0101u32);
179        assert_eq!(id.level(), 14);
180        let (child0, child1) = id.children();
181        assert_eq!(child0, PartitionID::new(0b1010_1010_1010_1010u32));
182        assert_eq!(child1, PartitionID::new(0b1010_1010_1010_1011u32));
183    }
184
185    #[test]
186    fn level() {
187        let root = PartitionID::root();
188        assert_eq!(root.level(), 0);
189        let (child0, child1) = root.children();
190
191        assert_eq!(child0.level(), 1);
192        assert_eq!(child1.level(), 1);
193    }
194
195    #[test]
196    fn root_parent() {
197        let root = PartitionID::root();
198        let roots_parent = root.parent();
199        assert_eq!(root, roots_parent);
200    }
201
202    #[test]
203    fn left_right_childs() {
204        let id = PartitionID(12345);
205        let (left_child, right_child) = id.children();
206        assert_eq!(left_child, id.left_child());
207        assert_eq!(right_child, id.right_child());
208    }
209
210    #[test]
211    fn is_left_right_child() {
212        let id = PartitionID(12345);
213        let (left_child, right_child) = id.children();
214        assert_eq!(left_child, id.left_child());
215        assert_eq!(right_child, id.right_child());
216        assert!(left_child.is_left_child());
217        assert!(right_child.is_right_child());
218    }
219
220    #[test]
221    fn make_left_child() {
222        let mut id = PartitionID(12345);
223        let (left_child, _) = id.children();
224        id.make_left_child();
225        assert_eq!(left_child, id);
226    }
227
228    #[test]
229    fn make_right_child() {
230        let mut id = PartitionID(12345);
231        let (_, right_child) = id.children();
232        id.make_right_child();
233        assert_eq!(right_child, id);
234    }
235
236    #[test]
237    fn into_usize() {
238        let id = PartitionID(12345);
239        let id_usize = usize::from(id);
240        assert_eq!(12345, id_usize);
241    }
242
243    #[test]
244    fn make_leftmost_descendant() {
245        let id = PartitionID(1);
246        let mut current = id;
247        for i in 1..30 {
248            let mut id = id;
249            id.make_leftmost_descendant(i);
250            assert_eq!(current.left_child(), id);
251            current = current.left_child();
252        }
253    }
254
255    #[test]
256    fn make_rightmost_descendant() {
257        let id = PartitionID(1);
258        let mut current = id;
259        for i in 1..30 {
260            let mut id = id;
261            id.make_rightmost_descendant(i);
262            assert_eq!(current.right_child(), id);
263            current = current.right_child();
264        }
265    }
266
267    #[test]
268    fn display() {
269        for i in 0..100 {
270            let id = PartitionID(i);
271            let string = format!("{id}");
272            let recast_id = PartitionID(string.parse::<u32>().unwrap());
273            assert_eq!(id, recast_id);
274        }
275    }
276
277    #[test]
278    fn partial_eq() {
279        for i in 0..100 {
280            let id = PartitionID(i);
281            let string = format!("{id}");
282            let recast_id = PartitionID(string.parse::<u32>().unwrap());
283            assert_eq!(id, recast_id);
284        }
285    }
286
287    #[test]
288    fn parent_at_level() {
289        let id = PartitionID::new(0xffff_ffff);
290        let levels = [0, 3, 9, 15, 20];
291        let results = [
292            PartitionID::new(0b11111111111111111111111111111111),
293            PartitionID::new(0b11111111111111111111111111111000),
294            PartitionID::new(0b11111111111111111111111000000000),
295            PartitionID::new(0b11111111111111111000000000000000),
296            PartitionID::new(0b11111111111100000000000000000000),
297        ];
298        levels
299            .iter()
300            .zip(results.iter())
301            .for_each(|(level, expected)| {
302                assert_eq!(id.parent_at_level(*level), *expected);
303            });
304    }
305
306    #[test]
307    fn binary_trait() {
308        let id = PartitionID::new(0xffff_ffff);
309        let levels = [0, 3, 9, 15, 20];
310        let results = [
311            "0b11111111111111111111111111111111",
312            "0b11111111111111111111111111111000",
313            "0b11111111111111111111111000000000",
314            "0b11111111111111111000000000000000",
315            "0b11111111111100000000000000000000",
316        ];
317        levels
318            .iter()
319            .zip(results.iter())
320            .for_each(|(level, expected)| {
321                assert_eq!(format!("{:#032b}", id.parent_at_level(*level)), *expected);
322            });
323    }
324
325    #[test]
326    fn lowest_common_ancestor() {
327        let a = PartitionID(0b1000);
328        let b = PartitionID(0b1001);
329        assert_eq!(a.lowest_common_ancestor(&b), b.lowest_common_ancestor(&a));
330
331        let expected = PartitionID(0b100);
332        assert_eq!(a.lowest_common_ancestor(&b), expected);
333
334        let a = PartitionID(0b1001);
335        let b = PartitionID(0b1111);
336        assert_eq!(a.lowest_common_ancestor(&b), b.lowest_common_ancestor(&a));
337
338        assert_eq!(a.lowest_common_ancestor(&b), PartitionID::root());
339    }
340
341    #[test]
342    fn bitand() {
343        let a = PartitionID(0b1000);
344        let b = PartitionID(0b1001);
345        assert_eq!(PartitionID(0b1000), a & b);
346    }
347
348    #[test]
349    fn bitor() {
350        let a = PartitionID(0b1000);
351        let b = PartitionID(0b1001);
352        assert_eq!(PartitionID(0b1001), a | b);
353    }
354
355    #[test]
356    fn extract_bit() {
357        let a = PartitionID(0b1001);
358        assert!(a.extract_bit(0));
359        assert!(!a.extract_bit(1));
360        assert!(!a.extract_bit(2));
361        assert!(a.extract_bit(3));
362        assert!(!a.extract_bit(4));
363
364        let a = PartitionID(0b100000000100000001000);
365        // [0, 3, 7, 11, 15]
366        assert!(!a.extract_bit(0));
367        assert!(a.extract_bit(3));
368        assert!(!a.extract_bit(7));
369        assert!(a.extract_bit(11));
370        assert!(!a.extract_bit(15));
371    }
372
373    #[test]
374    fn make_left_child_parent() {
375        let mut id = PartitionID(12345);
376        let original = id;
377        id.make_left_child();
378        assert_eq!(original, id.parent());
379        assert!(id.is_left_child());
380        assert!(!id.is_right_child());
381    }
382
383    #[test]
384    fn make_right_child_parent() {
385        let mut id = PartitionID(12345);
386        let original = id;
387        id.make_right_child();
388        assert_eq!(original, id.parent());
389        assert!(!id.is_left_child());
390        assert!(id.is_right_child());
391    }
392
393    #[test]
394    fn lowest_common_ancestor_comprehensive() {
395        // Test case 1: Same node - should return itself
396        let node = PartitionID(0b1000);
397        assert_eq!(node.lowest_common_ancestor(&node), node);
398
399        // Test case 2: Parent-child relationship
400        let parent = PartitionID(0b100);
401        let left_child = PartitionID(0b1000);
402        let right_child = PartitionID(0b1001);
403        assert_eq!(left_child.lowest_common_ancestor(&parent), parent);
404        assert_eq!(right_child.lowest_common_ancestor(&parent), parent);
405        assert_eq!(parent.lowest_common_ancestor(&left_child), parent);
406        assert_eq!(parent.lowest_common_ancestor(&right_child), parent);
407
408        // Test case 3: Siblings
409        assert_eq!(left_child.lowest_common_ancestor(&right_child), parent);
410        assert_eq!(right_child.lowest_common_ancestor(&left_child), parent);
411
412        // Test case 4: Nodes at very different levels
413        let deep_node = PartitionID(0b1000_0000); // Level 6
414        let shallow_node = PartitionID(0b10); // Level 1
415        let expected_ancestor = PartitionID(0b10); // Their LCA should be at level 1
416        assert_eq!(
417            deep_node.lowest_common_ancestor(&shallow_node),
418            expected_ancestor
419        );
420        assert_eq!(
421            shallow_node.lowest_common_ancestor(&deep_node),
422            expected_ancestor
423        );
424
425        // Test case 5: Root cases
426        let root = PartitionID::root();
427        let any_node = PartitionID(0b1111);
428        assert_eq!(root.lowest_common_ancestor(&any_node), root);
429        assert_eq!(any_node.lowest_common_ancestor(&root), root);
430        assert_eq!(root.lowest_common_ancestor(&root), root);
431
432        // Test case 6: Complex tree relationship
433        // Create a more complex scenario:
434        //           1
435        //         /   \
436        //       2       3
437        //      / \     / \
438        //     4   5   6   7
439        //    /
440        //   8
441        let node_1 = PartitionID::root(); // 1
442        let node_2 = PartitionID(0b10); // 2
443        let node_3 = PartitionID(0b11); // 3
444        let node_4 = PartitionID(0b100); // 4
445        let node_5 = PartitionID(0b101); // 5
446        let node_6 = PartitionID(0b110); // 6
447        let node_7 = PartitionID(0b111); // 7
448        let node_8 = PartitionID(0b1000); // 8
449
450        // Test various relationships
451        assert_eq!(node_4.lowest_common_ancestor(&node_5), node_2); // Siblings under 2
452        assert_eq!(node_6.lowest_common_ancestor(&node_7), node_3); // Siblings under 3
453        assert_eq!(node_4.lowest_common_ancestor(&node_6), node_1); // Cousins
454        assert_eq!(node_8.lowest_common_ancestor(&node_5), node_2); // Uncle relationship
455        assert_eq!(node_8.lowest_common_ancestor(&node_7), node_1); // Different subtrees
456    }
457}
toolbox_rs/partition_id.rs

toolbox_rs/
partition_id.rs
