oom_safe/
vec_ext.rs

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

oom_safe/
vec_ext.rs
