mz_ore/
vec.rs

Help
1// Copyright Materialize, Inc. and contributors. All rights reserved.
2//
3// Licensed under the Apache License, Version 2.0 (the "License");
4// you may not use this file except in compliance with the License.
5// You may obtain a copy of the License in the LICENSE file at the
6// root of this repository, or online at
7//
8//     http://www.apache.org/licenses/LICENSE-2.0
9//
10// Unless required by applicable law or agreed to in writing, software
11// distributed under the License is distributed on an "AS IS" BASIS,
12// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13// See the License for the specific language governing permissions and
14// limitations under the License.
15
16//! Vector utilities.
17
18use std::mem::{align_of, size_of};
19
20#[cfg(feature = "smallvec")]
21use smallvec::SmallVec;
22
23/// Create a new vector that re-uses the same allocation as an old one.
24/// The element types must have the same size and alignment.
25pub fn repurpose_allocation<T1, T2>(mut v: Vec<T1>) -> Vec<T2> {
26    assert_eq!(size_of::<T1>(), size_of::<T2>(), "same size");
27    assert_eq!(align_of::<T1>(), align_of::<T2>(), "same alignment");
28
29    v.clear();
30    let cap = v.capacity();
31    let p = v.as_mut_ptr().cast();
32    std::mem::forget(v);
33    // This is safe because `T1` and `T2` have the same size and alignment,
34    // `p`'s allocation is no longer owned by `v` (since that has been forgotten),
35    // and `p` was previously allocated with capacity `cap`.
36    unsafe { Vec::from_raw_parts(p, 0, cap) }
37}
38
39/// A trait for objects that behave like vectors.
40pub trait Vector<T> {
41    /// Appends an element to the vector.
42    fn push(&mut self, value: T);
43
44    /// Copies and appends all elements in a slice to the vector.
45    fn extend_from_slice(&mut self, other: &[T])
46    where
47        T: Copy;
48}
49
50impl<T> Vector<T> for Vec<T> {
51    fn push(&mut self, value: T) {
52        Vec::push(self, value)
53    }
54
55    fn extend_from_slice(&mut self, other: &[T])
56    where
57        T: Copy,
58    {
59        Vec::extend_from_slice(self, other)
60    }
61}
62
63#[cfg(feature = "smallvec")]
64impl<A> Vector<A::Item> for SmallVec<A>
65where
66    A: smallvec::Array,
67{
68    fn push(&mut self, value: A::Item) {
69        SmallVec::push(self, value)
70    }
71
72    fn extend_from_slice(&mut self, other: &[A::Item])
73    where
74        A::Item: Copy,
75    {
76        SmallVec::extend_from_slice(self, other)
77    }
78}
79
80#[cfg(feature = "compact_bytes")]
81impl Vector<u8> for compact_bytes::CompactBytes {
82    fn push(&mut self, value: u8) {
83        self.push(value)
84    }
85
86    fn extend_from_slice(&mut self, other: &[u8]) {
87        self.extend_from_slice(other)
88    }
89}
90
91/// Extension methods for `std::vec::Vec`
92pub trait VecExt<T> {
93    /// Creates an iterator which uses a closure to determine if an element should be removed.
94    ///
95    /// If the closure returns true, then the element is removed and yielded.
96    /// If the closure returns false, the element will remain in the vector and will not be yielded
97    /// by the iterator.
98    ///
99    /// Using this method and consuming the iterator is equivalent to the following code:
100    ///
101    /// ```
102    /// # let some_predicate = |x: &mut i32| { *x == 2 || *x == 3 || *x == 6 };
103    /// # let mut vec = vec![1, 2, 3, 4, 5, 6];
104    /// let mut i = 0;
105    /// while i < vec.len() {
106    ///     if some_predicate(&mut vec[i]) {
107    ///         let val = vec.swap_remove(i);
108    ///         // your code here
109    ///     } else {
110    ///         i += 1;
111    ///     }
112    /// }
113    ///
114    /// # assert_eq!(vec, vec![1, 5, 4]);
115    /// ```
116    ///
117    /// But `drain_filter_swapping` is easier to use.
118    ///
119    /// Note that `drain_filter_swapping` also lets you mutate every element in the filter closure,
120    /// regardless of whether you choose to keep or remove it.
121    ///
122    /// # Note
123    ///
124    /// Because the elements are removed using [`Vec::swap_remove`] the order of elements yielded
125    /// by the iterator and the order of the remaining elements in the original vector is **not**
126    /// maintained.
127    ///
128    /// # Examples
129    ///
130    /// Splitting an array into evens and odds, reusing the original allocation. Notice how the
131    /// order is not preserved in either results:
132    ///
133    /// ```
134    /// use mz_ore::vec::VecExt;
135    ///
136    /// let mut numbers = vec![1, 2, 3, 4, 5, 6, 8, 9, 11, 13, 14, 15];
137    ///
138    /// let evens = numbers.drain_filter_swapping(|x| *x % 2 == 0).collect::<Vec<_>>();
139    /// let odds = numbers;
140    ///
141    /// assert_eq!(evens, vec![2, 4, 14, 6, 8]);
142    /// assert_eq!(odds, vec![1, 15, 3, 13, 5, 11, 9]);
143    /// ```
144    #[must_use = "The vector is modified only if the iterator is consumed"]
145    fn drain_filter_swapping<F>(&mut self, filter: F) -> DrainFilterSwapping<'_, T, F>
146    where
147        F: FnMut(&mut T) -> bool;
148
149    /// Returns whether the vector is sorted using the given comparator function.
150    fn is_sorted_by<F>(&self, compare: F) -> bool
151    where
152        F: FnMut(&T, &T) -> bool;
153}
154
155/// Extension methods for `Vec<T>` where `T: PartialOrd`
156pub trait PartialOrdVecExt<T> {
157    /// Returns whether the vector is sorted.
158    // Remove once https://github.com/rust-lang/rust/issues/53485 is stabilized
159    fn is_sorted(&self) -> bool;
160    /// Returns whether the vector is sorted with strict inequality.
161    fn is_strictly_sorted(&self) -> bool;
162}
163
164impl<T> VecExt<T> for Vec<T> {
165    fn drain_filter_swapping<F>(&mut self, filter: F) -> DrainFilterSwapping<'_, T, F>
166    where
167        F: FnMut(&mut T) -> bool,
168    {
169        DrainFilterSwapping {
170            vec: self,
171            idx: 0,
172            pred: filter,
173        }
174    }
175
176    // implementation is from Vec::is_sorted_by, but with `windows` instead of
177    // the unstable `array_windows`
178    fn is_sorted_by<F>(&self, mut compare: F) -> bool
179    where
180        F: FnMut(&T, &T) -> bool,
181    {
182        self.windows(2).all(|win| compare(&win[0], &win[1]))
183    }
184}
185
186impl<T> PartialOrdVecExt<T> for Vec<T>
187where
188    T: PartialOrd,
189{
190    fn is_sorted(&self) -> bool {
191        self.is_sorted_by(|a, b| a <= b)
192    }
193
194    fn is_strictly_sorted(&self) -> bool {
195        self.is_sorted_by(|a, b| a < b)
196    }
197}
198
199/// An iterator which uses a closure to determine if an element should be removed.
200///
201/// This struct is created by [`VecExt::drain_filter_swapping`].
202/// See its documentation for more.
203///
204/// Warning: The vector is modified only if the iterator is consumed!
205///
206/// # Example
207///
208/// ```
209/// use mz_ore::vec::VecExt;
210///
211/// let mut v = vec![0, 1, 2];
212/// let iter: mz_ore::vec::DrainFilterSwapping<_, _> = v.drain_filter_swapping(|x| *x % 2 == 0);
213/// ```
214#[derive(Debug)]
215pub struct DrainFilterSwapping<'a, T, F> {
216    vec: &'a mut Vec<T>,
217    /// The index of the item that will be inspected by the next call to `next`.
218    idx: usize,
219    /// The filter test predicate.
220    pred: F,
221}
222
223impl<'a, T, F> Iterator for DrainFilterSwapping<'a, T, F>
224where
225    F: FnMut(&mut T) -> bool,
226{
227    type Item = T;
228
229    fn next(&mut self) -> Option<Self::Item> {
230        loop {
231            let item = self.vec.get_mut(self.idx)?;
232            if (self.pred)(item) {
233                return Some(self.vec.swap_remove(self.idx));
234            } else {
235                self.idx += 1;
236            }
237        }
238    }
239
240    fn size_hint(&self) -> (usize, Option<usize>) {
241        (0, Some(self.vec.len() - self.idx))
242    }
243}
244
245/// Remove the elements from `v` at the positions indicated by `indexes`, and return the removed
246/// elements in a new vector.
247///
248/// `indexes` shouldn't have duplicates. (Might panic or behave incorrectly in case of
249/// duplicates.)
250pub fn swap_remove_multiple<T>(v: &mut Vec<T>, mut indexes: Vec<usize>) -> Vec<T> {
251    indexes.sort();
252    indexes.reverse();
253    let mut result = Vec::new();
254    for r in indexes {
255        result.push(v.swap_remove(r));
256    }
257    result
258}
259
260#[cfg(test)]
261mod test {
262    use super::*;
263
264    #[crate::test]
265    fn miri_test_repurpose() {
266        let v: Vec<usize> = vec![0, 1, 2];
267
268        let mut other: Vec<isize> = repurpose_allocation(v);
269
270        assert!(other.is_empty());
271        other.push(-1);
272        assert_eq!(other[0], -1);
273
274        struct Gus1 {
275            s: String,
276        }
277        impl Drop for Gus1 {
278            fn drop(&mut self) {
279                println!("happy {}", self.s);
280            }
281        }
282
283        struct Gus2 {
284            s: String,
285        }
286        impl Drop for Gus2 {
287            fn drop(&mut self) {
288                println!("happy {}", self.s);
289            }
290        }
291
292        // also exercise non-`Copy`, `Drop`-impling values as well
293        let v: Vec<Gus1> = vec![Gus1 {
294            s: "hmm".to_string(),
295        }];
296
297        let mut other: Vec<Gus2> = repurpose_allocation(v);
298
299        assert!(other.is_empty());
300        other.push(Gus2 {
301            s: "hmm2".to_string(),
302        });
303        assert_eq!(other[0].s, "hmm2");
304    }
305
306    #[crate::test]
307    #[should_panic(expected = "same size")]
308    fn miri_test_wrong_size() {
309        let v: Vec<usize> = vec![0, 1, 2];
310        let _: Vec<()> = repurpose_allocation(v);
311    }
312
313    #[crate::test]
314    #[should_panic(expected = "same alignment")]
315    fn miri_test_wrong_align() {
316        #[repr(align(8))]
317        #[derive(Default)]
318        struct Gus1 {
319            _i: [u8; 16],
320        }
321
322        #[repr(align(16))]
323        #[derive(Default)]
324        struct Gus2 {
325            _i: [u8; 8],
326        }
327
328        use std::mem::size_of;
329        assert_eq!(size_of::<Gus1>(), size_of::<Gus2>(), "same size in test");
330
331        // You need a value in here to have miri catch the problem, if we remove
332        // the alignment check
333        let v: Vec<Gus1> = vec![Default::default()];
334        let _: Vec<Gus2> = repurpose_allocation(v);
335    }
336
337    #[crate::test]
338    fn test_is_sorted_by() {
339        assert!(vec![0, 1, 2].is_sorted_by(|a, b| a < b));
340        assert!(vec![0, 1, 2].is_sorted_by(|a, b| a <= b));
341        assert!(!vec![0, 1, 2].is_sorted_by(|a, b| a > b));
342        assert!(!vec![0, 1, 2].is_sorted_by(|a, b| a >= b));
343        assert!(vec![0, 1, 2].is_sorted_by(|_a, _b| true));
344        assert!(!vec![0, 1, 2].is_sorted_by(|_a, _b| false));
345
346        assert!(!vec![0, 1, 1, 2].is_sorted_by(|a, b| a < b));
347        assert!(vec![0, 1, 1, 2].is_sorted_by(|a, b| a <= b));
348        assert!(!vec![0, 1, 1, 2].is_sorted_by(|a, b| a > b));
349        assert!(!vec![0, 1, 1, 2].is_sorted_by(|a, b| a >= b));
350        assert!(vec![0, 1, 1, 2].is_sorted_by(|_a, _b| true));
351        assert!(!vec![0, 1, 1, 2].is_sorted_by(|_a, _b| false));
352
353        assert!(!vec![2, 1, 0].is_sorted_by(|a, b| a < b));
354        assert!(!vec![2, 1, 0].is_sorted_by(|a, b| a <= b));
355        assert!(vec![2, 1, 0].is_sorted_by(|a, b| a > b));
356        assert!(vec![2, 1, 0].is_sorted_by(|a, b| a >= b));
357        assert!(vec![2, 1, 0].is_sorted_by(|_a, _b| true));
358        assert!(!vec![2, 1, 0].is_sorted_by(|_a, _b| false));
359
360        assert!(!vec![5, 1, 9, 42].is_sorted_by(|a, b| a < b));
361        assert!(!vec![5, 1, 9, 42].is_sorted_by(|a, b| a <= b));
362        assert!(!vec![5, 1, 9, 42].is_sorted_by(|a, b| a > b));
363        assert!(!vec![5, 1, 9, 42].is_sorted_by(|a, b| a >= b));
364        assert!(vec![5, 1, 9, 42].is_sorted_by(|_a, _b| true));
365        assert!(!vec![5, 1, 9, 42].is_sorted_by(|_a, _b| false));
366    }
367}
mz_ore/vec.rs

mz_ore/
vec.rs
