rust_utils/
vec.rs

1use crate::sealed::Sealed;
2
3use alloc::vec::Vec;
4use core::ptr;
5
6pub trait VecExt<T>: Sealed {
7    fn insert_vec(&mut self, index: usize, other: Vec<T>);
8
9    fn remove_if<F>(&mut self, f: F)
10    where
11        F: FnMut(&mut T) -> bool;
12}
13
14impl<T> Sealed for Vec<T> {}
15
16impl<T> VecExt<T> for Vec<T> {
17    fn insert_vec(&mut self, index: usize, mut v: Vec<T>) {
18        self.reserve_exact(v.len());
19
20        unsafe {
21            let len = self.len();
22            let base = self.as_mut_ptr();
23            let additional = v.len();
24
25            // move existing elements
26            ptr::copy(base.add(index), base.add(index + additional), len - index);
27
28            // copy from `v`
29            ptr::copy_nonoverlapping(v.as_ptr(), base.add(index), additional);
30
31            // set length
32            v.set_len(0);
33            self.set_len(len + additional);
34        }
35    }
36
37    fn remove_if<F>(&mut self, mut f: F)
38    where
39        F: FnMut(&mut T) -> bool,
40    {
41        self.retain_mut(|x| !f(x))
42    }
43}
44
45#[cfg(test)]
46mod tests {
47    use super::*;
48
49    fn o<T: ToOwned + ?Sized>(x: &T) -> T::Owned {
50        x.to_owned()
51    }
52
53    #[test]
54    fn test_insert_vec() {
55        {
56            let mut v = vec![1, 2, 3, 4, 5];
57            v.insert_vec(2, vec![6, 7, 8]);
58            assert_eq!(v, vec![1, 2, 6, 7, 8, 3, 4, 5]);
59        }
60
61        // owned strings
62        {
63            let mut v = vec![o("a"), o("b"), o("c"), o("d"), o("e")];
64            v.insert_vec(2, vec![o("f"), o("g"), o("h")]);
65            assert_eq!(v, vec![o("a"), o("b"), o("f"), o("g"), o("h"), o("c"), o("d"), o("e")]);
66        }
67    }
68}