1 // SPDX-License-Identifier: GPL-2.0 1 // SPDX-License-Identifier: GPL-2.0
2 2
3 //! Extensions to [`Vec`] for fallible allocat 3 //! Extensions to [`Vec`] for fallible allocations.
4 4
5 use super::{AllocError, Flags}; 5 use super::{AllocError, Flags};
6 use alloc::vec::Vec; 6 use alloc::vec::Vec;
7 7
8 /// Extensions to [`Vec`]. 8 /// Extensions to [`Vec`].
9 pub trait VecExt<T>: Sized { 9 pub trait VecExt<T>: Sized {
10 /// Creates a new [`Vec`] instance with at 10 /// Creates a new [`Vec`] instance with at least the given capacity.
11 /// 11 ///
12 /// # Examples 12 /// # Examples
13 /// 13 ///
14 /// ``` 14 /// ```
15 /// let v = Vec::<u32>::with_capacity(20, 15 /// let v = Vec::<u32>::with_capacity(20, GFP_KERNEL)?;
16 /// 16 ///
17 /// assert!(v.capacity() >= 20); 17 /// assert!(v.capacity() >= 20);
18 /// # Ok::<(), Error>(()) 18 /// # Ok::<(), Error>(())
19 /// ``` 19 /// ```
20 fn with_capacity(capacity: usize, flags: F 20 fn with_capacity(capacity: usize, flags: Flags) -> Result<Self, AllocError>;
21 21
22 /// Appends an element to the back of the 22 /// Appends an element to the back of the [`Vec`] instance.
23 /// 23 ///
24 /// # Examples 24 /// # Examples
25 /// 25 ///
26 /// ``` 26 /// ```
27 /// let mut v = Vec::new(); 27 /// let mut v = Vec::new();
28 /// v.push(1, GFP_KERNEL)?; 28 /// v.push(1, GFP_KERNEL)?;
29 /// assert_eq!(&v, &[1]); 29 /// assert_eq!(&v, &[1]);
30 /// 30 ///
31 /// v.push(2, GFP_KERNEL)?; 31 /// v.push(2, GFP_KERNEL)?;
32 /// assert_eq!(&v, &[1, 2]); 32 /// assert_eq!(&v, &[1, 2]);
33 /// # Ok::<(), Error>(()) 33 /// # Ok::<(), Error>(())
34 /// ``` 34 /// ```
35 fn push(&mut self, v: T, flags: Flags) -> 35 fn push(&mut self, v: T, flags: Flags) -> Result<(), AllocError>;
36 36
37 /// Pushes clones of the elements of slice 37 /// Pushes clones of the elements of slice into the [`Vec`] instance.
38 /// 38 ///
39 /// # Examples 39 /// # Examples
40 /// 40 ///
41 /// ``` 41 /// ```
42 /// let mut v = Vec::new(); 42 /// let mut v = Vec::new();
43 /// v.push(1, GFP_KERNEL)?; 43 /// v.push(1, GFP_KERNEL)?;
44 /// 44 ///
45 /// v.extend_from_slice(&[20, 30, 40], GFP 45 /// v.extend_from_slice(&[20, 30, 40], GFP_KERNEL)?;
46 /// assert_eq!(&v, &[1, 20, 30, 40]); 46 /// assert_eq!(&v, &[1, 20, 30, 40]);
47 /// 47 ///
48 /// v.extend_from_slice(&[50, 60], GFP_KER 48 /// v.extend_from_slice(&[50, 60], GFP_KERNEL)?;
49 /// assert_eq!(&v, &[1, 20, 30, 40, 50, 60 49 /// assert_eq!(&v, &[1, 20, 30, 40, 50, 60]);
50 /// # Ok::<(), Error>(()) 50 /// # Ok::<(), Error>(())
51 /// ``` 51 /// ```
52 fn extend_from_slice(&mut self, other: &[T 52 fn extend_from_slice(&mut self, other: &[T], flags: Flags) -> Result<(), AllocError>
53 where 53 where
54 T: Clone; 54 T: Clone;
55 55
56 /// Ensures that the capacity exceeds the 56 /// Ensures that the capacity exceeds the length by at least `additional` elements.
57 /// 57 ///
58 /// # Examples 58 /// # Examples
59 /// 59 ///
60 /// ``` 60 /// ```
61 /// let mut v = Vec::new(); 61 /// let mut v = Vec::new();
62 /// v.push(1, GFP_KERNEL)?; 62 /// v.push(1, GFP_KERNEL)?;
63 /// 63 ///
64 /// v.reserve(10, GFP_KERNEL)?; 64 /// v.reserve(10, GFP_KERNEL)?;
65 /// let cap = v.capacity(); 65 /// let cap = v.capacity();
66 /// assert!(cap >= 10); 66 /// assert!(cap >= 10);
67 /// 67 ///
68 /// v.reserve(10, GFP_KERNEL)?; 68 /// v.reserve(10, GFP_KERNEL)?;
69 /// let new_cap = v.capacity(); 69 /// let new_cap = v.capacity();
70 /// assert_eq!(new_cap, cap); 70 /// assert_eq!(new_cap, cap);
71 /// 71 ///
72 /// # Ok::<(), Error>(()) 72 /// # Ok::<(), Error>(())
73 /// ``` 73 /// ```
74 fn reserve(&mut self, additional: usize, f 74 fn reserve(&mut self, additional: usize, flags: Flags) -> Result<(), AllocError>;
75 } 75 }
76 76
77 impl<T> VecExt<T> for Vec<T> { 77 impl<T> VecExt<T> for Vec<T> {
78 fn with_capacity(capacity: usize, flags: F 78 fn with_capacity(capacity: usize, flags: Flags) -> Result<Self, AllocError> {
79 let mut v = Vec::new(); 79 let mut v = Vec::new();
80 <Self as VecExt<_>>::reserve(&mut v, c 80 <Self as VecExt<_>>::reserve(&mut v, capacity, flags)?;
81 Ok(v) 81 Ok(v)
82 } 82 }
83 83
84 fn push(&mut self, v: T, flags: Flags) -> 84 fn push(&mut self, v: T, flags: Flags) -> Result<(), AllocError> {
85 <Self as VecExt<_>>::reserve(self, 1, 85 <Self as VecExt<_>>::reserve(self, 1, flags)?;
86 let s = self.spare_capacity_mut(); 86 let s = self.spare_capacity_mut();
87 s[0].write(v); 87 s[0].write(v);
88 88
89 // SAFETY: We just initialised the fir 89 // SAFETY: We just initialised the first spare entry, so it is safe to increase the length
90 // by 1. We also know that the new len 90 // by 1. We also know that the new length is <= capacity because of the previous call to
91 // `reserve` above. 91 // `reserve` above.
92 unsafe { self.set_len(self.len() + 1) 92 unsafe { self.set_len(self.len() + 1) };
93 Ok(()) 93 Ok(())
94 } 94 }
95 95
96 fn extend_from_slice(&mut self, other: &[T 96 fn extend_from_slice(&mut self, other: &[T], flags: Flags) -> Result<(), AllocError>
97 where 97 where
98 T: Clone, 98 T: Clone,
99 { 99 {
100 <Self as VecExt<_>>::reserve(self, oth 100 <Self as VecExt<_>>::reserve(self, other.len(), flags)?;
101 for (slot, item) in core::iter::zip(se 101 for (slot, item) in core::iter::zip(self.spare_capacity_mut(), other) {
102 slot.write(item.clone()); 102 slot.write(item.clone());
103 } 103 }
104 104
105 // SAFETY: We just initialised the `ot 105 // SAFETY: We just initialised the `other.len()` spare entries, so it is safe to increase
106 // the length by the same amount. We a 106 // the length by the same amount. We also know that the new length is <= capacity because
107 // of the previous call to `reserve` a 107 // of the previous call to `reserve` above.
108 unsafe { self.set_len(self.len() + oth 108 unsafe { self.set_len(self.len() + other.len()) };
109 Ok(()) 109 Ok(())
110 } 110 }
111 111
112 #[cfg(any(test, testlib))] 112 #[cfg(any(test, testlib))]
113 fn reserve(&mut self, additional: usize, _ 113 fn reserve(&mut self, additional: usize, _flags: Flags) -> Result<(), AllocError> {
114 Vec::reserve(self, additional); 114 Vec::reserve(self, additional);
115 Ok(()) 115 Ok(())
116 } 116 }
117 117
118 #[cfg(not(any(test, testlib)))] 118 #[cfg(not(any(test, testlib)))]
119 fn reserve(&mut self, additional: usize, f 119 fn reserve(&mut self, additional: usize, flags: Flags) -> Result<(), AllocError> {
120 let len = self.len(); 120 let len = self.len();
121 let cap = self.capacity(); 121 let cap = self.capacity();
122 122
123 if cap - len >= additional { 123 if cap - len >= additional {
124 return Ok(()); 124 return Ok(());
125 } 125 }
126 126
127 if core::mem::size_of::<T>() == 0 { 127 if core::mem::size_of::<T>() == 0 {
128 // The capacity is already `usize: 128 // The capacity is already `usize::MAX` for SZTs, we can't go higher.
129 return Err(AllocError); 129 return Err(AllocError);
130 } 130 }
131 131
132 // We know cap is <= `isize::MAX` beca 132 // We know cap is <= `isize::MAX` because `Layout::array` fails if the resulting byte size
133 // is greater than `isize::MAX`. So th 133 // is greater than `isize::MAX`. So the multiplication by two won't overflow.
134 let new_cap = core::cmp::max(cap * 2, 134 let new_cap = core::cmp::max(cap * 2, len.checked_add(additional).ok_or(AllocError)?);
135 let layout = core::alloc::Layout::arra 135 let layout = core::alloc::Layout::array::<T>(new_cap).map_err(|_| AllocError)?;
136 136
137 let (old_ptr, len, cap) = destructure( 137 let (old_ptr, len, cap) = destructure(self);
138 138
139 // We need to make sure that `ptr` is 139 // We need to make sure that `ptr` is either NULL or comes from a previous call to
140 // `krealloc_aligned`. A `Vec<T>`'s `p 140 // `krealloc_aligned`. A `Vec<T>`'s `ptr` value is not guaranteed to be NULL and might be
141 // dangling after being created with ` 141 // dangling after being created with `Vec::new`. Instead, we can rely on `Vec<T>`'s capacity
142 // to be zero if no memory has been al 142 // to be zero if no memory has been allocated yet.
143 let ptr = if cap == 0 { 143 let ptr = if cap == 0 {
144 core::ptr::null_mut() 144 core::ptr::null_mut()
145 } else { 145 } else {
146 old_ptr 146 old_ptr
147 }; 147 };
148 148
149 // SAFETY: `ptr` is valid because it's 149 // SAFETY: `ptr` is valid because it's either NULL or comes from a previous call to
150 // `krealloc_aligned`. We also verifie 150 // `krealloc_aligned`. We also verified that the type is not a ZST.
151 let new_ptr = unsafe { super::allocato 151 let new_ptr = unsafe { super::allocator::krealloc_aligned(ptr.cast(), layout, flags) };
152 if new_ptr.is_null() { 152 if new_ptr.is_null() {
153 // SAFETY: We are just rebuilding 153 // SAFETY: We are just rebuilding the existing `Vec` with no changes.
154 unsafe { rebuild(self, old_ptr, le 154 unsafe { rebuild(self, old_ptr, len, cap) };
155 Err(AllocError) 155 Err(AllocError)
156 } else { 156 } else {
157 // SAFETY: `ptr` has been realloca 157 // SAFETY: `ptr` has been reallocated with the layout for `new_cap` elements. New cap
158 // is greater than `cap`, so it co 158 // is greater than `cap`, so it continues to be >= `len`.
159 unsafe { rebuild(self, new_ptr.cas 159 unsafe { rebuild(self, new_ptr.cast::<T>(), len, new_cap) };
160 Ok(()) 160 Ok(())
161 } 161 }
162 } 162 }
163 } 163 }
164 164
165 #[cfg(not(any(test, testlib)))] 165 #[cfg(not(any(test, testlib)))]
166 fn destructure<T>(v: &mut Vec<T>) -> (*mut T, 166 fn destructure<T>(v: &mut Vec<T>) -> (*mut T, usize, usize) {
167 let mut tmp = Vec::new(); 167 let mut tmp = Vec::new();
168 core::mem::swap(&mut tmp, v); 168 core::mem::swap(&mut tmp, v);
169 let mut tmp = core::mem::ManuallyDrop::new 169 let mut tmp = core::mem::ManuallyDrop::new(tmp);
170 let len = tmp.len(); 170 let len = tmp.len();
171 let cap = tmp.capacity(); 171 let cap = tmp.capacity();
172 (tmp.as_mut_ptr(), len, cap) 172 (tmp.as_mut_ptr(), len, cap)
173 } 173 }
174 174
175 /// Rebuilds a `Vec` from a pointer, length, a 175 /// Rebuilds a `Vec` from a pointer, length, and capacity.
176 /// 176 ///
177 /// # Safety 177 /// # Safety
178 /// 178 ///
179 /// The same as [`Vec::from_raw_parts`]. 179 /// The same as [`Vec::from_raw_parts`].
180 #[cfg(not(any(test, testlib)))] 180 #[cfg(not(any(test, testlib)))]
181 unsafe fn rebuild<T>(v: &mut Vec<T>, ptr: *mut 181 unsafe fn rebuild<T>(v: &mut Vec<T>, ptr: *mut T, len: usize, cap: usize) {
182 // SAFETY: The safety requirements from th 182 // SAFETY: The safety requirements from this function satisfy those of `from_raw_parts`.
183 let mut tmp = unsafe { Vec::from_raw_parts 183 let mut tmp = unsafe { Vec::from_raw_parts(ptr, len, cap) };
184 core::mem::swap(&mut tmp, v); 184 core::mem::swap(&mut tmp, v);
185 } 185 }
Linux® is a registered trademark of Linus Torvalds in the United States and other countries.
TOMOYO® is a registered trademark of NTT DATA CORPORATION.