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
#[macro_export]
9
macro_rules! vecbag {
10
    () => {
11
        $crate::VecBag::new()
12
    };
13
    ($elem:expr; $n:expr) => {
14
        $crate::VecBag::from_vec(vec![$elem; $n])
15
    };
16
    ($($x:expr),+ $(,)?) => {{
17
        let mut __bag = $crate::VecBag::new();
18
        $( let _ = __bag.insert($x); )*
19
        __bag
20
    }};
21
}
22

            
23
///
24
/// A bag (multiset) that is internally represented by a sorted vector.
25
/// Mostly useful for a compact representation of bags that are not changed often.
26
///
27
#[derive(Clone, PartialEq, Eq, PartialOrd, Ord, Hash)]
28
pub struct VecBag<T: Ord> {
29
    /// The internal storage with the invariant that the array is sorted.
30
    sorted_array: Vec<T>,
31
}
32

            
33
impl<T: Ord> VecBag<T> {
34
    /// Creates a new, empty VecBag.
35
3969
    pub fn new() -> Self {
36
3969
        Self {
37
3969
            sorted_array: Vec::new(),
38
3969
        }
39
3969
    }
40

            
41
    /// Creates a VecBag from the given vector without assuming anything of the given vector.
42
16004
    pub fn from_vec(mut vec: Vec<T>) -> Self {
43
16004
        vec.sort_unstable();
44
16004
        Self { sorted_array: vec }
45
16004
    }
46

            
47
    /// Returns the capacity of the bag.
48
    pub fn capacity(&self) -> usize {
49
        self.sorted_array.capacity()
50
    }
51

            
52
    /// Returns true iff the bag contains at least one occurrence of the given element.
53
    pub fn contains(&self, element: &T) -> bool {
54
        self.sorted_array.binary_search(element).is_ok()
55
    }
56

            
57
    /// Clears the bag, removing all elements.
58
    pub fn clear(&mut self) {
59
        self.sorted_array.clear();
60
    }
61

            
62
    /// Retains only the elements specified by the predicate.
63
    pub fn retain<F>(&mut self, mut f: F)
64
    where
65
        F: FnMut(&T) -> bool,
66
    {
67
        // Removing elements does not change the order.
68
        self.sorted_array.retain(|e| f(e));
69
    }
70

            
71
    /// Returns true iff this bag is a subbag of the other bag.
72
    pub fn is_subset(&self, other: &VecBag<T>) -> bool {
73
        let mut self_iter = self.sorted_array.iter();
74
        let mut other_iter = other.sorted_array.iter();
75

            
76
        // Traverse both bags in order, checking multiplicity as well.
77
        let mut self_next = self_iter.next();
78
        let mut other_next = other_iter.next();
79

            
80
        while let Some(self_val) = self_next {
81
            match other_next {
82
                Some(other_val) => match self_val.cmp(other_val) {
83
                    Ordering::Equal => {
84
                        self_next = self_iter.next();
85
                        other_next = other_iter.next();
86
                    }
87
                    Ordering::Greater => other_next = other_iter.next(),
88
                    Ordering::Less => return false,
89
                },
90
                None => return false,
91
            }
92
        }
93

            
94
        true
95
    }
96

            
97
    /// Returns a new bag only containing the given element.
98
1700
    pub fn singleton(element: T) -> Self {
99
1700
        Self {
100
1700
            sorted_array: vec![element],
101
1700
        }
102
1700
    }
103

            
104
    /// Returns true iff the bag is empty.
105
22249
    pub fn is_empty(&self) -> bool {
106
22249
        self.sorted_array.is_empty()
107
22249
    }
108

            
109
    /// Returns the difference of this bag and the other bag.
110
100
    pub fn difference<'a>(&'a self, other: &'a VecBag<T>) -> impl Iterator<Item = &'a T> {
111
100
        Difference::<'a, T, _> {
112
100
            self_iter: self.sorted_array.iter(),
113
100
            other_iter: other.sorted_array.iter(),
114
100
            other_next: None,
115
100
            marker: PhantomData,
116
100
        }
117
100
    }
118

            
119
    /// Inserts one occurrence of the given element into the bag.
120
3555
    pub fn insert(&mut self, element: T) -> bool {
121
        // Finds the location where to insert the element to keep the array sorted.
122
3555
        let position = match self.sorted_array.binary_search(&element) {
123
3555
            Ok(position) | Err(position) => position,
124
        };
125

            
126
3555
        self.sorted_array.insert(position, element);
127
3555
        true
128
3555
    }
129

            
130
    /// Extends this bag with the elements from the given iterator, returns true iff at least one element was inserted.
131
    pub fn extend<'a, I: IntoIterator<Item = &'a T>>(&mut self, iter: I) -> bool
132
    where
133
        T: Clone + 'a,
134
    {
135
        let mut inserted = false;
136
        for element in iter {
137
            self.insert(element.clone());
138
            inserted = true;
139
        }
140

            
141
        inserted
142
    }
143

            
144
    /// Returns an iterator over the elements in the bag, they are yielded in sorted order.
145
22036
    pub fn iter(&self) -> impl Iterator<Item = &T> {
146
22036
        self.sorted_array.iter()
147
22036
    }
148

            
149
    /// Returns the number of elements in the bag.
150
    pub fn len(&self) -> usize {
151
        self.sorted_array.len()
152
    }
153

            
154
    /// Consumes the bag and returns a vector with the elements in sorted order.
155
    pub fn to_vec(self) -> Vec<T> {
156
        self.sorted_array
157
    }
158
}
159

            
160
impl<T: Ord> FromIterator<T> for VecBag<T> {
161
    fn from_iter<I: IntoIterator<Item = T>>(iter: I) -> Self {
162
        Self::from_vec(iter.into_iter().collect())
163
    }
164
}
165

            
166
/// A lazy iterator that yields the difference of two bags. The elements are yielded in sorted order.
167
struct Difference<'a, T, I> {
168
    self_iter: I,
169
    other_iter: I,
170
    other_next: Option<&'a T>,
171
    marker: PhantomData<&'a T>,
172
}
173

            
174
impl<'a, T, I> Iterator for Difference<'a, T, I>
175
where
176
    I: Iterator<Item = &'a T>,
177
    T: Ord + PartialEq,
178
{
179
    type Item = &'a T;
180

            
181
692
    fn next(&mut self) -> Option<Self::Item> {
182
692
        if self.other_next.is_none() {
183
269
            self.other_next = self.other_iter.next();
184
423
        }
185

            
186
1032
        for self_val in self.self_iter.by_ref() {
187
            loop {
188
1493
                match self.other_next {
189
1324
                    Some(other_val) => match self_val.cmp(other_val) {
190
                        Ordering::Equal => {
191
440
                            self.other_next = self.other_iter.next();
192
440
                            break;
193
                        }
194
461
                        Ordering::Greater => self.other_next = self.other_iter.next(),
195
423
                        Ordering::Less => return Some(self_val),
196
                    },
197
169
                    None => return Some(self_val),
198
                }
199
            }
200
        }
201

            
202
100
        None
203
692
    }
204
}
205

            
206
impl<T: Ord> Default for VecBag<T> {
207
    fn default() -> Self {
208
        Self::new()
209
    }
210
}
211

            
212
impl<'a, T: Ord> IntoIterator for &'a VecBag<T> {
213
    type Item = &'a T;
214
    type IntoIter = Iter<'a, T>;
215

            
216
    fn into_iter(self) -> Self::IntoIter {
217
        self.sorted_array.iter()
218
    }
219
}
220

            
221
impl<T: Ord + fmt::Debug> fmt::Debug for VecBag<T> {
222
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
223
        write!(f, "{{{:?}}}", self.sorted_array.iter().format(", "))
224
    }
225
}
226

            
227
#[cfg(test)]
228
mod tests {
229
    use rand::RngExt;
230

            
231
    use merc_utilities::random_test;
232

            
233
    use crate::VecBag;
234

            
235
    #[test]
236
1
    fn test_random_vecbag_difference() {
237
100
        random_test(100, |rng| {
238
100
            let size = rng.random_range(0..20);
239
100
            let size2 = rng.random_range(0..20);
240
1032
            let vec1: Vec<u32> = (0..size).map(|_| rng.random_range(0..10)).collect();
241
1047
            let vec2: Vec<u32> = (0..size2).map(|_| rng.random_range(0..10)).collect();
242

            
243
100
            let bag1 = VecBag::from_vec(vec1.clone());
244
100
            let bag2 = VecBag::from_vec(vec2.clone());
245

            
246
100
            let difference: Vec<u32> = bag1.difference(&bag2).cloned().collect();
247

            
248
100
            let mut expected_difference = vec1;
249
100
            expected_difference.sort_unstable();
250

            
251
100
            let mut sorted_vec2 = vec2;
252
100
            sorted_vec2.sort_unstable();
253

            
254
1047
            for value in sorted_vec2 {
255
1047
                if let Ok(index) = expected_difference.binary_search(&value) {
256
440
                    expected_difference.remove(index);
257
607
                }
258
            }
259

            
260
100
            assert_eq!(difference, expected_difference);
261
100
        })
262
1
    }
263
}