fso_graph/
afforest.rs

1use std::{mem::ManuallyDrop, sync::atomic::Ordering};
2
3use rayon::prelude::*;
4
5use crate::prelude::*;
6
7/// A union find data structure based on [1].
8///
9/// Note, that this data structure requires calling `compress`
10/// before calling `find` in order to return the correct set id.
11///
12/// [1]  Michael Sutton, Tal Ben-Nun, Amnon Barak:
13///      "Optimizing Parallel Graph Connectivity Computation via Subgraph Sampling",
14///       Symposium on Parallel and Distributed Processing, IPDPS 2018
15pub struct Afforest<NI: Idx>(Box<[Atomic<NI>]>);
16
17unsafe impl<NI: Idx> Send for Afforest<NI> {}
18unsafe impl<NI: Idx> Sync for Afforest<NI> {}
19
20impl<NI: Idx> UnionFind<NI> for Afforest<NI> {
21    // Corresponds to the `link` method described in [1].
22    fn union(&self, u: NI, v: NI) {
23        let mut p1 = self.find(u);
24        let mut p2 = self.find(v);
25
26        while p1 != p2 {
27            let high = NI::max(p1, p2);
28            let low = p1 + p2 - high;
29            let p_high = self.find(high);
30
31            if p_high == low
32                || (p_high == high && self.update_parent(self.find(high), high, low).is_ok())
33            {
34                break;
35            }
36            p1 = self.parent(self.parent(high));
37            p2 = self.parent(low);
38        }
39    }
40
41    fn find(&self, u: NI) -> NI {
42        self.parent(u)
43    }
44
45    fn len(&self) -> usize {
46        self.0.len()
47    }
48
49    // Corresponds to the `compress` method described in [1].
50    fn compress(&self) {
51        (0..self.len()).into_par_iter().map(NI::new).for_each(|n| {
52            while self.parent(n) != self.parent(self.parent(n)) {
53                self.0[n.index()].store(self.parent(self.parent(n)), Ordering::SeqCst)
54            }
55        });
56    }
57}
58
59impl<NI: Idx> Afforest<NI> {
60    /// Creates a new disjoint-set struct of `size` elements.
61    ///
62    /// # Examples
63    ///
64    /// ```
65    /// use fso_graph::prelude::*;
66    ///
67    /// let af = Afforest::new(3);
68    /// af.union(0, 1);
69    /// af.compress();
70    ///
71    /// let set0 = af.find(0);
72    /// let set1 = af.find(1);
73    /// assert_eq!(set0, set1);
74    /// ```
75    pub fn new(size: usize) -> Self {
76        let mut v = Vec::with_capacity(size);
77
78        (0..size)
79            .into_par_iter()
80            .map(|i| Atomic::new(NI::new(i)))
81            .collect_into_vec(&mut v);
82
83        Self(v.into_boxed_slice())
84    }
85
86    fn parent(&self, i: NI) -> NI {
87        self.0[i.index()].load(Ordering::SeqCst)
88    }
89
90    fn update_parent(&self, id: NI, current: NI, new: NI) -> Result<NI, NI> {
91        self.0[id.index()].compare_exchange_weak(current, new, Ordering::SeqCst, Ordering::Relaxed)
92    }
93}
94
95impl<NI: Idx> Components<NI> for Afforest<NI> {
96    fn component(&self, node: NI) -> NI {
97        self.find(node)
98    }
99
100    fn to_vec(self) -> Vec<NI> {
101        let mut components = ManuallyDrop::new(self.0.into_vec());
102        let (ptr, len, cap) = (
103            components.as_mut_ptr(),
104            components.len(),
105            components.capacity(),
106        );
107
108        // SAFETY: NI and NI::Atomic have the same memory layout
109        unsafe {
110            let ptr = ptr as *mut Vec<NI>;
111            let ptr = ptr as *mut _;
112            Vec::from_raw_parts(ptr, len, cap)
113        }
114    }
115}
116
117#[cfg(test)]
118mod test {
119    use crate::prelude::*;
120
121    #[test]
122    fn test_union() {
123        let af = Afforest::new(10);
124
125        af.union(9, 7);
126        af.union(7, 4);
127        af.union(4, 2);
128        af.union(2, 0);
129
130        af.compress();
131
132        assert_eq!(af.find(9), 0);
133    }
134}
fso_graph/afforest.rs

fso_graph/
afforest.rs
