alloc_safe/
vec.rs

1//! Vec extensions.
2
3use crate::alloc::AllocError;
4use crate::sealed::Sealed;
5use std::alloc::{Allocator, Global};
6
7/// Extension for `Vec<T, A>`
8pub trait VecAllocExt<T, A: Allocator>: Sized + Sealed {
9    fn try_with_capacity_in(capacity: usize, alloc: A) -> Result<Self, AllocError>;
10    fn try_push(&mut self, value: T) -> Result<(), AllocError>;
11    fn try_resize(&mut self, new_len: usize, value: T) -> Result<(), AllocError>
12    where
13        T: Copy;
14    fn try_resize_with<F>(&mut self, new_len: usize, f: F) -> Result<(), AllocError>
15    where
16        F: FnMut() -> T;
17    fn try_copy_from_slice(&mut self, other: &[T]) -> Result<(), AllocError>
18    where
19        T: Copy;
20}
21
22/// Extension for `Vec<T>`
23pub trait VecExt<T>: VecAllocExt<T, Global> {
24    fn try_with_capacity(capacity: usize) -> Result<Self, AllocError>;
25}
26
27impl<T, A: Allocator> Sealed for Vec<T, A> {}
28
29impl<T, A: Allocator> VecAllocExt<T, A> for Vec<T, A> {
30    #[inline]
31    fn try_with_capacity_in(capacity: usize, alloc: A) -> Result<Self, AllocError> {
32        let mut vec = Vec::new_in(alloc);
33        vec.try_reserve(capacity)?;
34        Ok(vec)
35    }
36
37    #[inline]
38    fn try_push(&mut self, value: T) -> Result<(), AllocError> {
39        if self.len() == self.capacity() {
40            self.try_reserve(1)?;
41        }
42        self.push(value);
43        Ok(())
44    }
45
46    #[inline]
47    fn try_resize(&mut self, new_len: usize, value: T) -> Result<(), AllocError>
48    where
49        T: Copy,
50    {
51        if new_len > self.len() {
52            self.try_reserve(new_len - self.len())?;
53        }
54        self.resize(new_len, value);
55        Ok(())
56    }
57
58    #[inline]
59    fn try_resize_with<F>(&mut self, new_len: usize, f: F) -> Result<(), AllocError>
60    where
61        F: FnMut() -> T,
62    {
63        if new_len > self.len() {
64            self.try_reserve(new_len - self.len())?;
65        }
66        self.resize_with(new_len, f);
67        Ok(())
68    }
69
70    fn try_copy_from_slice(&mut self, other: &[T]) -> Result<(), AllocError>
71    where
72        T: Copy,
73    {
74        self.try_reserve(other.len())?;
75        self.extend_from_slice(other);
76        Ok(())
77    }
78}
79
80impl<T> VecExt<T> for Vec<T> {
81    #[inline]
82    fn try_with_capacity(capacity: usize) -> Result<Self, AllocError> {
83        let mut vec = Vec::new();
84        vec.try_reserve(capacity)?;
85        Ok(vec)
86    }
87}
88
89#[cfg(test)]
90mod tests {
91    use super::*;
92    use std::alloc::{AllocError, Layout};
93    use std::ptr::NonNull;
94
95    struct Alloc {
96        inner: Global,
97        _mark: u64,
98    }
99
100    impl Alloc {
101        pub const fn new() -> Self {
102            Alloc {
103                inner: Global,
104                _mark: u64::MAX,
105            }
106        }
107    }
108
109    unsafe impl Allocator for Alloc {
110        fn allocate(&self, layout: Layout) -> Result<NonNull<[u8]>, AllocError> {
111            self.inner.allocate(layout)
112        }
113
114        unsafe fn deallocate(&self, ptr: NonNull<u8>, layout: Layout) {
115            self.inner.deallocate(ptr, layout)
116        }
117    }
118
119    #[test]
120    fn test_vec() {
121        let mut v = Vec::try_with_capacity(1).unwrap();
122        v.try_push(1).unwrap();
123        v.try_push(2).unwrap();
124        v.try_push(3).unwrap();
125        assert_eq!(v, [1, 2, 3]);
126    }
127
128    #[test]
129    fn test_vec_with_alloc() {
130        let mut v = Vec::try_with_capacity_in(1, Alloc::new()).unwrap();
131        v.try_push(1).unwrap();
132        v.try_push(2).unwrap();
133        v.try_push(3).unwrap();
134        assert_eq!(v, [1, 2, 3]);
135    }
136
137    #[test]
138    fn test_vec_resize() {
139        let mut v = Vec::new_in(Alloc::new());
140        v.try_resize(3, 1).unwrap();
141        assert_eq!(v, [1, 1, 1]);
142    }
143
144    #[test]
145    fn test_vec_resize_with() {
146        let mut v = Vec::new_in(Alloc::new());
147        let mut i = 0;
148        v.try_resize_with(3, || {
149            i += 1;
150            i
151        })
152        .unwrap();
153        assert_eq!(v, [1, 2, 3]);
154    }
155
156    #[test]
157    fn test_vec_extend_from_slice() {
158        let mut v = Vec::new_in(Alloc::new());
159        v.try_copy_from_slice(&[1, 2, 3]).unwrap();
160        assert_eq!(v, [1, 2, 3]);
161    }
162}
alloc_safe/vec.rs

alloc_safe/
vec.rs
