pliron 0.17.0

Programming Languages Intermediate RepresentatiON
Documentation 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142

// SPDX-License-Identifier: Apache-2.0
// Copyright (c) The pliron contributors

//! A union-find (disjoint-set) data structure.

use alloc::{vec, vec::Vec};
use core::hash::Hash;

use crate::utils::table::HMap;

/// A union-find (disjoint-set) structure over values of type `T`.
///
/// Elements need not be registered up front: any element passed to a method
/// is added, as a singleton set, on first sight. Uses path compression and
/// union-by-size; member lists are maintained so that enumerating a set
/// ([Self::set_members]) is O(set size).
#[derive(Clone, Debug)]
pub struct UnionFind<T: Eq + Hash + Clone> {
    parent: HMap<T, T>,
    /// Members of each set, keyed by the set's representative.
    members: HMap<T, Vec<T>>,
}

impl<T: Eq + Hash + Clone> Default for UnionFind<T> {
    fn default() -> Self {
        Self {
            parent: HMap::default(),
            members: HMap::default(),
        }
    }
}

impl<T: Eq + Hash + Clone> UnionFind<T> {
    /// Get the representative of `v`'s set, adding a singleton set if `v` is new.
    pub fn find(&mut self, v: T) -> T {
        match self.parent.get(&v) {
            None => {
                self.parent.insert(v.clone(), v.clone());
                self.members.insert(v.clone(), vec![v.clone()]);
                v
            }
            Some(p) if *p == v => v,
            Some(p) => {
                let root = self.find(p.clone());
                // Path compression.
                self.parent.insert(v, root.clone());
                root
            }
        }
    }

    /// Merge the sets of `a` and `b`.
    pub fn union(&mut self, a: T, b: T) {
        let ra = self.find(a);
        let rb = self.find(b);
        if ra == rb {
            return;
        }
        // Union by size.
        let (big, small) = if self.members[&ra].len() >= self.members[&rb].len() {
            (ra, rb)
        } else {
            (rb, ra)
        };
        self.parent.insert(small.clone(), big.clone());
        let small_members = self
            .members
            .remove(&small)
            .expect("set root must have a members list");
        self.members
            .get_mut(&big)
            .expect("set root must have a members list")
            .extend(small_members);
    }

    /// Are `a` and `b` in the same set?
    pub fn in_same_set(&mut self, a: T, b: T) -> bool {
        self.find(a) == self.find(b)
    }

    /// All values in `v`'s set (including `v` itself), in no particular order.
    pub fn set_members(&mut self, v: T) -> &[T] {
        let root = self.find(v);
        &self.members[&root]
    }
}

#[cfg(test)]
mod tests {
    use super::UnionFind;
    use alloc::{vec, vec::Vec};

    fn sorted_members(uf: &mut UnionFind<u32>, v: u32) -> Vec<u32> {
        let mut members = uf.set_members(v).to_vec();
        members.sort();
        members
    }

    #[test]
    fn singletons_on_first_sight() {
        let mut uf = UnionFind::<u32>::default();
        assert_eq!(uf.find(7), 7);
        assert!(!uf.in_same_set(1, 2));
        assert_eq!(sorted_members(&mut uf, 7), vec![7]);
    }

    #[test]
    fn union_merges_sets() {
        let mut uf = UnionFind::<u32>::default();
        uf.union(1, 2);
        uf.union(3, 4);
        assert!(uf.in_same_set(1, 2));
        assert!(uf.in_same_set(3, 4));
        assert!(!uf.in_same_set(2, 3));

        uf.union(2, 3);
        assert!(uf.in_same_set(1, 4));
        assert_eq!(sorted_members(&mut uf, 4), vec![1, 2, 3, 4]);
    }

    #[test]
    fn union_is_idempotent() {
        let mut uf = UnionFind::<u32>::default();
        uf.union(1, 2);
        uf.union(2, 1);
        uf.union(1, 2);
        assert_eq!(sorted_members(&mut uf, 1), vec![1, 2]);
    }

    #[test]
    fn transitive_chain() {
        let mut uf = UnionFind::<u32>::default();
        for i in 0..10 {
            uf.union(i, i + 1);
        }
        assert!(uf.in_same_set(0, 10));
        assert_eq!(uf.set_members(5).len(), 11);
        // All elements resolve to the same representative.
        let root = uf.find(0);
        assert!((0..=10).all(|i| uf.find(i) == root));
    }
}