1
use std::cmp::Ordering;
2
use std::fmt;
3
use std::marker::PhantomData;
4
use std::slice::Iter;
5

            
6
use itertools::Itertools;
7

            
8
use crate::vec_difference::Difference;
9

            
10
#[macro_export]
11
macro_rules! vecset {
12
    () => {
13
        $crate::VecSet::new()
14
    };
15
    ($elem:expr; $n:expr) => {{
16
        let mut __set = $crate::VecSet::new();
17
        let __count: usize = $n;
18
        if __count > 0 {
19
            __set.insert($elem);
20
        }
21
        __set
22
    }};
23
    ($($x:expr),+ $(,)?) => {{
24
        let mut __set = $crate::VecSet::new();
25
        $( let _ = __set.insert($x); )*
26
        __set
27
    }};
28
}
29

            
30
///
31
/// A set that is internally represented by a sorted vector. Mostly useful for
32
/// a compact representation of sets that are not changed often.
33
///
34
#[derive(Clone, PartialEq, Eq, PartialOrd, Ord, Hash)]
35
pub struct VecSet<T> {
36
    /// The internal storage with the invariant that the array is sorted.
37
    sorted_array: Vec<T>,
38
}
39

            
40
impl<T: Ord> VecSet<T> {
41
    /// Creates a new, empty VecSet.
42
7967
    pub fn new() -> Self {
43
7967
        Self {
44
7967
            sorted_array: Vec::new(),
45
7967
        }
46
7967
    }
47

            
48
    /// Creates a VecSet from the given vector without assuming anything of the given vector.
49
10646
    pub fn from_vec(mut vec: Vec<T>) -> Self {
50
10646
        vec.sort_unstable();
51
10646
        vec.dedup();
52
10646
        Self { sorted_array: vec }
53
10646
    }
54

            
55
    /// Returns the capacity of the set.
56
    pub fn capacity(&self) -> usize {
57
        self.sorted_array.capacity()
58
    }
59

            
60
    /// Returns true iff the set contains the given element.
61
    pub fn contains(&self, element: &T) -> bool {
62
        self.sorted_array.binary_search(element).is_ok()
63
    }
64

            
65
    /// Clears the set, removing all elements.
66
    pub fn clear(&mut self) {
67
        self.sorted_array.clear();
68
    }
69

            
70
    /// Retains only the elements specified by the predicate.
71
7760
    pub fn retain<F>(&mut self, mut f: F)
72
7760
    where
73
7760
        F: FnMut(&T) -> bool,
74
    {
75
        // Removing elements does not change the order.
76
18687
        self.sorted_array.retain(|e| f(e));
77
7760
    }
78

            
79
    /// Returns true iff this set is a subset of the other set.
80
70963
    pub fn is_subset(&self, other: &VecSet<T>) -> bool {
81
70963
        let mut self_iter = self.sorted_array.iter();
82
70963
        let mut other_iter = other.sorted_array.iter();
83

            
84
        // Traverse both sets in order, checking that all elements of self are in other.
85
70963
        let mut self_next = self_iter.next();
86
70963
        let mut other_next = other_iter.next();
87

            
88
195271
        while let Some(self_val) = self_next {
89
192891
            match other_next {
90
180660
                Some(other_val) => {
91
180660
                    match self_val.cmp(other_val) {
92
49817
                        Ordering::Equal => {
93
49817
                            self_next = self_iter.next();
94
49817
                            other_next = other_iter.next();
95
49817
                        }
96
74491
                        Ordering::Greater => other_next = other_iter.next(),
97
56352
                        Ordering::Less => return false, // self_val is smaller, not in other
98
                    }
99
                }
100
12231
                None => return false, // other is exhausted
101
            }
102
        }
103

            
104
2380
        true
105
70963
    }
106

            
107
    /// Returns a new set only containing the given element.
108
15119
    pub fn singleton(element: T) -> Self {
109
15119
        Self {
110
15119
            sorted_array: vec![element],
111
15119
        }
112
15119
    }
113

            
114
    /// Returns true iff the set is empty.
115
3324
    pub fn is_empty(&self) -> bool {
116
3324
        self.sorted_array.is_empty()
117
3324
    }
118

            
119
    /// Returns the difference of this set and the other set.
120
1247
    pub fn difference<'a>(&'a self, other: &'a VecSet<T>) -> impl Iterator<Item = &'a T> {
121
1247
        Difference::<'a, T, _> {
122
1247
            self_iter: self.sorted_array.iter(),
123
1247
            other_iter: other.sorted_array.iter(),
124
1247
            other_next: None,
125
1247
            marker: PhantomData,
126
1247
        }
127
1247
    }
128

            
129
    /// Inserts the given element into the set, returns true iff the element was
130
    /// inserted.
131
28622
    pub fn insert(&mut self, element: T) -> bool {
132
        // Finds the location where to insert the element to keep the array sorted.
133
28622
        if let Err(position) = self.sorted_array.binary_search(&element) {
134
27389
            self.sorted_array.insert(position, element);
135
27389
            return true;
136
1233
        }
137

            
138
1233
        false
139
28622
    }
140

            
141
    /// Extends this set with the elements from the given iterator, returns true iff at least one element was inserted.
142
563
    pub fn extend<'a, I: IntoIterator<Item = &'a T>>(&mut self, iter: I) -> bool
143
563
    where
144
563
        T: Clone + 'a,
145
    {
146
563
        let mut inserted = false;
147
1494
        for element in iter {
148
1494
            if self.insert(element.clone()) {
149
1322
                inserted = true;
150
1322
            }
151
        }
152

            
153
563
        inserted
154
563
    }
155

            
156
    /// Returns an iterator over the elements in the set, they are yielded in sorted order.
157
9474
    pub fn iter(&self) -> impl Iterator<Item = &T> {
158
9474
        self.sorted_array.iter()
159
9474
    }
160

            
161
    /// Returns the number of elements in the set.
162
4
    pub fn len(&self) -> usize {
163
4
        self.sorted_array.len()
164
4
    }
165

            
166
    /// Consumes the set and returns a vector with the elements in sorted order.
167
4476
    pub fn to_vec(self) -> Vec<T> {
168
4476
        self.sorted_array
169
4476
    }
170
}
171

            
172
impl<T: Ord> FromIterator<T> for VecSet<T> {
173
4946
    fn from_iter<I: IntoIterator<Item = T>>(iter: I) -> Self {
174
4946
        Self::from_vec(iter.into_iter().collect())
175
4946
    }
176
}
177

            
178
impl<T: Ord> Default for VecSet<T> {
179
    fn default() -> Self {
180
        Self::new()
181
    }
182
}
183

            
184
impl<'a, T> IntoIterator for &'a VecSet<T> {
185
    type Item = &'a T;
186
    type IntoIter = Iter<'a, T>;
187

            
188
384
    fn into_iter(self) -> Self::IntoIter {
189
384
        self.sorted_array.iter()
190
384
    }
191
}
192

            
193
impl<T: fmt::Debug> fmt::Debug for VecSet<T> {
194
202
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
195
202
        write!(f, "{{{:?}}}", self.sorted_array.iter().format(", "))
196
202
    }
197
}
198

            
199
#[cfg(test)]
200
mod tests {
201
    use itertools::Itertools;
202
    use rand::RngExt;
203

            
204
    use merc_utilities::random_test;
205

            
206
    use crate::VecSet;
207

            
208
    #[test]
209
1
    fn test_random_vecset_difference() {
210
100
        random_test(100, |rng| {
211
100
            let size = rng.random_range(0..20);
212
100
            let size2 = rng.random_range(0..20);
213
901
            let vec1: Vec<u32> = (0..size).map(|_| rng.random_range(0..100)).collect();
214
855
            let vec2: Vec<u32> = (0..size2).map(|_| rng.random_range(0..100)).collect();
215

            
216
100
            let set1 = VecSet::from_vec(vec1.clone());
217
100
            let set2 = VecSet::from_vec(vec2.clone());
218

            
219
100
            println!("left: {:?}", set1);
220
100
            println!("right: {:?}", set2);
221

            
222
100
            let difference: Vec<u32> = set1.difference(&set2).cloned().collect();
223
100
            let expected_difference: Vec<u32> = vec1
224
100
                .into_iter()
225
901
                .filter(|x| !vec2.contains(x))
226
100
                .sorted()
227
100
                .dedup()
228
100
                .collect();
229

            
230
100
            assert_eq!(difference, expected_difference);
231
100
        })
232
1
    }
233
}