Skip to main content

relibc/
c_vec.rs

1//! Equivalent of Rust's `Vec<T>`, but using relibc's own allocator.
2
3use crate::{
4    io::{self, Write},
5    platform::{self, WriteByte, types::*},
6};
7use core::{
8    cmp, fmt,
9    iter::IntoIterator,
10    mem,
11    ops::{Deref, DerefMut},
12    ptr::{self, NonNull},
13    slice,
14};
15
16/// Error that occurs when an allocation fails
17#[derive(Debug, Default, Hash, PartialEq, Eq, Clone, Copy)]
18pub struct AllocError;
19
20/// A normal vector allocated in Rust needs to be dropped from Rust
21/// too, in order to avoid UB. This CVec is an abstraction that works
22/// using only C allocations functions and can therefore be dropped
23/// from C. Just like the Rust Vec, this does bounds checks to assure
24/// you never reach isize::MAX. Unless you need to drop something from
25/// C, prefer Rust's builtin Vec.
26pub struct CVec<T> {
27    ptr: NonNull<T>,
28    len: usize,
29    cap: usize,
30}
31impl<T> CVec<T> {
32    #[allow(clippy::new_without_default)]
33    pub fn new() -> Self {
34        Self {
35            ptr: NonNull::dangling(),
36            len: 0,
37            cap: 0,
38        }
39    }
40    fn check_bounds(i: usize) -> Result<usize, AllocError> {
41        if i > isize::MAX as usize {
42            Err(AllocError)
43        } else {
44            Ok(i)
45        }
46    }
47    fn check_mul(x: usize, y: usize) -> Result<usize, AllocError> {
48        x.checked_mul(y)
49            .ok_or(AllocError)
50            .and_then(Self::check_bounds)
51    }
52    pub fn with_capacity(cap: usize) -> Result<Self, AllocError> {
53        if cap == 0 {
54            return Ok(Self::new());
55        }
56        let size = Self::check_mul(cap, mem::size_of::<T>())?;
57        let ptr = NonNull::new(unsafe { platform::alloc(size).cast::<T>() }).ok_or(AllocError)?;
58        Ok(Self { ptr, len: 0, cap })
59    }
60    unsafe fn resize(&mut self, cap: usize) -> Result<(), AllocError> {
61        let size = Self::check_mul(cap, mem::size_of::<T>())?;
62        let ptr = if cap == 0 {
63            NonNull::dangling()
64        } else if self.cap > 0 {
65            NonNull::new(
66                unsafe { platform::realloc(self.ptr.as_ptr().cast::<c_void>(), size) }.cast::<T>(),
67            )
68            .ok_or(AllocError)?
69        } else {
70            NonNull::new((unsafe { platform::alloc(size) }).cast::<T>()).ok_or(AllocError)?
71        };
72        self.ptr = ptr;
73        self.cap = cap;
74        Ok(())
75    }
76    unsafe fn drop_range(&mut self, start: usize, end: usize) {
77        let mut start = unsafe { self.ptr.as_ptr().add(start) };
78        let end = unsafe { self.ptr.as_ptr().add(end) };
79        while start < end {
80            unsafe { ptr::drop_in_place(start) };
81            start = unsafe { start.add(1) };
82        }
83    }
84
85    // Push stuff
86
87    pub fn reserve(&mut self, required: usize) -> Result<(), AllocError> {
88        let required_len = self
89            .len
90            .checked_add(required)
91            .ok_or(AllocError)
92            .and_then(Self::check_bounds)?;
93        if required_len > self.cap {
94            let new_cap = cmp::min(required_len.next_power_of_two(), isize::MAX as usize);
95            unsafe {
96                self.resize(new_cap)?;
97            }
98        }
99        Ok(())
100    }
101    pub fn push(&mut self, elem: T) -> Result<(), AllocError> {
102        self.reserve(1)?;
103        unsafe {
104            ptr::write(self.ptr.as_ptr().add(self.len), elem);
105        }
106        self.len += 1; // no need to bounds check, as new len <= cap
107        Ok(())
108    }
109    pub fn extend_from_slice(&mut self, elems: &[T]) -> Result<(), AllocError>
110    where
111        T: Copy,
112    {
113        self.reserve(elems.len())?;
114        unsafe {
115            ptr::copy_nonoverlapping(elems.as_ptr(), self.ptr.as_ptr().add(self.len), elems.len());
116        }
117        self.len += elems.len(); // no need to bounds check, as new len <= cap
118        Ok(())
119    }
120    pub fn append(&mut self, other: &mut Self) -> Result<(), AllocError> {
121        let len = other.len;
122        other.len = 0; // move
123        self.reserve(len)?;
124        unsafe {
125            ptr::copy_nonoverlapping(other.as_ptr(), self.ptr.as_ptr().add(self.len), len);
126        }
127        self.len += other.len(); // no need to bounds check, as new len <= cap
128        Ok(())
129    }
130
131    // Pop stuff
132
133    pub fn truncate(&mut self, len: usize) {
134        if len < self.len {
135            unsafe {
136                let old_len = self.len;
137                self.drop_range(len, old_len);
138            }
139            self.len = len;
140        }
141    }
142    pub fn shrink_to_fit(&mut self) -> Result<(), AllocError> {
143        if self.len < self.cap {
144            unsafe {
145                let new_cap = self.len;
146                self.resize(new_cap)?;
147            }
148        }
149        Ok(())
150    }
151    pub fn pop(&mut self) -> Option<T> {
152        if self.is_empty() {
153            None
154        } else {
155            let elem = unsafe { ptr::read(self.as_ptr().add(self.len - 1)) };
156            self.len -= 1;
157            Some(elem)
158        }
159    }
160
161    // Misc stuff
162
163    pub fn capacity(&self) -> usize {
164        self.cap
165    }
166    pub fn as_ptr(&self) -> *const T {
167        self.ptr.as_ptr()
168    }
169    pub fn as_mut_ptr(&mut self) -> *mut T {
170        self.ptr.as_ptr()
171    }
172    /// Leaks the inner data. This is safe to drop from C!
173    pub fn leak(mut self) -> *mut T {
174        let ptr = self.as_mut_ptr();
175        mem::forget(self);
176        ptr
177    }
178}
179impl<T> Deref for CVec<T> {
180    type Target = [T];
181
182    fn deref(&self) -> &Self::Target {
183        unsafe { slice::from_raw_parts(self.ptr.as_ptr(), self.len) }
184    }
185}
186impl<T> DerefMut for CVec<T> {
187    fn deref_mut(&mut self) -> &mut Self::Target {
188        unsafe { slice::from_raw_parts_mut(self.ptr.as_ptr(), self.len) }
189    }
190}
191impl<T> Drop for CVec<T> {
192    fn drop(&mut self) {
193        unsafe {
194            let len = self.len;
195            self.drop_range(0, len);
196        }
197    }
198}
199impl<'a, T> IntoIterator for &'a CVec<T> {
200    type Item = <&'a [T] as IntoIterator>::Item;
201    type IntoIter = <&'a [T] as IntoIterator>::IntoIter;
202    fn into_iter(self) -> Self::IntoIter {
203        <&[T]>::into_iter(self)
204    }
205}
206impl<'a, T> IntoIterator for &'a mut CVec<T> {
207    type Item = <&'a mut [T] as IntoIterator>::Item;
208    type IntoIter = <&'a mut [T] as IntoIterator>::IntoIter;
209    fn into_iter(self) -> Self::IntoIter {
210        <&mut [T]>::into_iter(&mut *self)
211    }
212}
213
214impl Write for CVec<u8> {
215    fn write(&mut self, buf: &[u8]) -> io::Result<usize> {
216        self.extend_from_slice(buf).map_err(|err| {
217            io::Error::new(
218                io::ErrorKind::Other,
219                "AllocStringWriter::write failed to allocate",
220            )
221        })?;
222        Ok(buf.len())
223    }
224    fn flush(&mut self) -> io::Result<()> {
225        Ok(())
226    }
227}
228impl fmt::Write for CVec<u8> {
229    fn write_str(&mut self, s: &str) -> fmt::Result {
230        self.write(s.as_bytes()).map_err(|_| fmt::Error)?;
231        Ok(())
232    }
233}
234impl WriteByte for CVec<u8> {
235    fn write_u8(&mut self, byte: u8) -> fmt::Result {
236        self.write(&[byte]).map_err(|_| fmt::Error)?;
237        Ok(())
238    }
239}
240
241#[cfg(test)]
242mod tests {
243    use super::CVec;
244
245    #[test]
246    fn push_pop() {
247        let mut vec = CVec::new();
248        vec.push(1).unwrap();
249        vec.push(2).unwrap();
250        vec.push(3).unwrap();
251        assert_eq!(&vec[..], &[1, 2, 3]);
252        assert_eq!(vec.pop().unwrap(), 3);
253        assert_eq!(&vec[..], &[1, 2]);
254    }
255    #[test]
256    fn extend_from_slice() {
257        use crate::io::Write;
258
259        let mut vec = CVec::new();
260        vec.extend_from_slice(&[1, 2, 3]).unwrap();
261        vec.extend_from_slice(&[4, 5, 6]).unwrap();
262        assert_eq!(&vec[..], &[1, 2, 3, 4, 5, 6]);
263        assert_eq!(vec.write(&[7, 8, 9]).unwrap(), 3);
264        assert_eq!(&vec[..], &[1, 2, 3, 4, 5, 6, 7, 8, 9]);
265    }
266    #[test]
267    fn dropped() {
268        use alloc::rc::Rc;
269
270        let counter = Rc::new(());
271        let mut vec = CVec::with_capacity(3).unwrap();
272        vec.push(Rc::clone(&counter)).unwrap();
273        vec.push(Rc::clone(&counter)).unwrap();
274        vec.push(Rc::clone(&counter)).unwrap();
275        assert_eq!(Rc::strong_count(&counter), 4);
276
277        let popped = vec.pop().unwrap();
278        assert_eq!(Rc::strong_count(&counter), 4);
279        drop(popped);
280        assert_eq!(Rc::strong_count(&counter), 3);
281
282        vec.push(Rc::clone(&counter)).unwrap();
283        vec.push(Rc::clone(&counter)).unwrap();
284        vec.push(Rc::clone(&counter)).unwrap();
285
286        assert_eq!(vec.len(), 5);
287        assert_eq!(Rc::strong_count(&counter), 6);
288        vec.truncate(1);
289        assert_eq!(Rc::strong_count(&counter), 2);
290
291        drop(vec);
292        assert_eq!(Rc::strong_count(&counter), 1);
293    }
294}