mz_ore/
vec.rs

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    /// Returns whether the vector is sorted using the given comparator function.
94    fn is_sorted_by<F>(&self, compare: F) -> bool
95    where
96        F: FnMut(&T, &T) -> bool;
97}
98
99/// Extension methods for `Vec<T>` where `T: PartialOrd`
100pub trait PartialOrdVecExt<T> {
101    /// Returns whether the vector is sorted.
102    // Remove once https://github.com/rust-lang/rust/issues/53485 is stabilized
103    fn is_sorted(&self) -> bool;
104    /// Returns whether the vector is sorted with strict inequality.
105    fn is_strictly_sorted(&self) -> bool;
106}
107
108impl<T> VecExt<T> for Vec<T> {
109    // implementation is from Vec::is_sorted_by, but with `windows` instead of
110    // the unstable `array_windows`
111    fn is_sorted_by<F>(&self, mut compare: F) -> bool
112    where
113        F: FnMut(&T, &T) -> bool,
114    {
115        self.windows(2).all(|win| compare(&win[0], &win[1]))
116    }
117}
118
119impl<T> PartialOrdVecExt<T> for Vec<T>
120where
121    T: PartialOrd,
122{
123    fn is_sorted(&self) -> bool {
124        self.is_sorted_by(|a, b| a <= b)
125    }
126
127    fn is_strictly_sorted(&self) -> bool {
128        self.is_sorted_by(|a, b| a < b)
129    }
130}
131
132/// Remove the elements from `v` at the positions indicated by `indexes`, and return the removed
133/// elements in a new vector.
134///
135/// `indexes` shouldn't have duplicates. (Might panic or behave incorrectly in case of
136/// duplicates.)
137pub fn swap_remove_multiple<T>(v: &mut Vec<T>, mut indexes: Vec<usize>) -> Vec<T> {
138    indexes.sort();
139    indexes.reverse();
140    let mut result = Vec::new();
141    for r in indexes {
142        result.push(v.swap_remove(r));
143    }
144    result
145}
146
147#[cfg(test)]
148mod test {
149    use super::*;
150
151    #[crate::test]
152    fn miri_test_repurpose() {
153        let v: Vec<usize> = vec![0, 1, 2];
154
155        let mut other: Vec<isize> = repurpose_allocation(v);
156
157        assert!(other.is_empty());
158        other.push(-1);
159        assert_eq!(other[0], -1);
160
161        struct Gus1 {
162            s: String,
163        }
164        impl Drop for Gus1 {
165            fn drop(&mut self) {
166                println!("happy {}", self.s);
167            }
168        }
169
170        struct Gus2 {
171            s: String,
172        }
173        impl Drop for Gus2 {
174            fn drop(&mut self) {
175                println!("happy {}", self.s);
176            }
177        }
178
179        // also exercise non-`Copy`, `Drop`-impling values as well
180        let v: Vec<Gus1> = vec![Gus1 {
181            s: "hmm".to_string(),
182        }];
183
184        let mut other: Vec<Gus2> = repurpose_allocation(v);
185
186        assert!(other.is_empty());
187        other.push(Gus2 {
188            s: "hmm2".to_string(),
189        });
190        assert_eq!(other[0].s, "hmm2");
191    }
192
193    #[crate::test]
194    #[should_panic(expected = "same size")]
195    fn miri_test_wrong_size() {
196        let v: Vec<usize> = vec![0, 1, 2];
197        let _: Vec<()> = repurpose_allocation(v);
198    }
199
200    #[crate::test]
201    #[should_panic(expected = "same alignment")]
202    fn miri_test_wrong_align() {
203        #[repr(align(8))]
204        #[derive(Default)]
205        struct Gus1 {
206            _i: [u8; 16],
207        }
208
209        #[repr(align(16))]
210        #[derive(Default)]
211        struct Gus2 {
212            _i: [u8; 8],
213        }
214
215        use std::mem::size_of;
216        assert_eq!(size_of::<Gus1>(), size_of::<Gus2>(), "same size in test");
217
218        // You need a value in here to have miri catch the problem, if we remove
219        // the alignment check
220        let v: Vec<Gus1> = vec![Default::default()];
221        let _: Vec<Gus2> = repurpose_allocation(v);
222    }
223
224    #[crate::test]
225    fn test_is_sorted_by() {
226        assert!(vec![0, 1, 2].is_sorted_by(|a, b| a < b));
227        assert!(vec![0, 1, 2].is_sorted_by(|a, b| a <= b));
228        assert!(!vec![0, 1, 2].is_sorted_by(|a, b| a > b));
229        assert!(!vec![0, 1, 2].is_sorted_by(|a, b| a >= b));
230        assert!(vec![0, 1, 2].is_sorted_by(|_a, _b| true));
231        assert!(!vec![0, 1, 2].is_sorted_by(|_a, _b| false));
232
233        assert!(!vec![0, 1, 1, 2].is_sorted_by(|a, b| a < b));
234        assert!(vec![0, 1, 1, 2].is_sorted_by(|a, b| a <= b));
235        assert!(!vec![0, 1, 1, 2].is_sorted_by(|a, b| a > b));
236        assert!(!vec![0, 1, 1, 2].is_sorted_by(|a, b| a >= b));
237        assert!(vec![0, 1, 1, 2].is_sorted_by(|_a, _b| true));
238        assert!(!vec![0, 1, 1, 2].is_sorted_by(|_a, _b| false));
239
240        assert!(!vec![2, 1, 0].is_sorted_by(|a, b| a < b));
241        assert!(!vec![2, 1, 0].is_sorted_by(|a, b| a <= b));
242        assert!(vec![2, 1, 0].is_sorted_by(|a, b| a > b));
243        assert!(vec![2, 1, 0].is_sorted_by(|a, b| a >= b));
244        assert!(vec![2, 1, 0].is_sorted_by(|_a, _b| true));
245        assert!(!vec![2, 1, 0].is_sorted_by(|_a, _b| false));
246
247        assert!(!vec![5, 1, 9, 42].is_sorted_by(|a, b| a < b));
248        assert!(!vec![5, 1, 9, 42].is_sorted_by(|a, b| a <= b));
249        assert!(!vec![5, 1, 9, 42].is_sorted_by(|a, b| a > b));
250        assert!(!vec![5, 1, 9, 42].is_sorted_by(|a, b| a >= b));
251        assert!(vec![5, 1, 9, 42].is_sorted_by(|_a, _b| true));
252        assert!(!vec![5, 1, 9, 42].is_sorted_by(|_a, _b| false));
253    }
254}
mz_ore/vec.rs

mz_ore/
vec.rs
