Skip to main content

gc_alloc/
vec.rs

1use std::{ffi::c_void, ptr::NonNull};
2
3use crate::{GcToken, gc};
4
5pub fn from_fn<T>(_token: &impl GcToken, len: usize, mut f: impl FnMut(usize) -> T) -> GcVec<T> {
6    let vec = VecInner::<T>::new(len);
7    for i in 0..len {
8        unsafe { vec.as_ptr().add(i).write(f(i)) };
9    }
10    vec.set_len(len);
11
12    register_finalizer(vec.as_ptr());
13
14    GcVec(unsafe { std::slice::from_raw_parts_mut(vec.as_ptr(), len).into() })
15}
16
17pub fn repeat<T: Clone>(token: &impl GcToken, val: T, len: usize) -> GcVec<T> {
18    from_fn(token, len, |_| val.clone())
19}
20
21pub fn from_iter<T, I: IntoIterator<Item = T>>(_token: &impl GcToken, iter: I) -> GcVec<T> {
22    let iter = iter.into_iter();
23    let (lower, _) = iter.size_hint();
24
25    let mut cap = lower.max(1);
26    let mut vec = VecInner::<T>::new(cap);
27    let mut len = 0;
28    for item in iter {
29        if len == cap {
30            cap = cap.checked_mul(2).expect("Capacity overflow");
31            let new_vec = VecInner::<T>::new(cap);
32            unsafe { std::ptr::copy_nonoverlapping(vec.as_ptr(), new_vec.as_ptr(), len) };
33            vec = new_vec;
34        }
35        unsafe { vec.as_ptr().add(len).write(item) };
36        len += 1;
37    }
38    vec.set_len(len);
39
40    register_finalizer(vec.as_ptr());
41
42    GcVec(unsafe { std::slice::from_raw_parts_mut(vec.as_ptr(), len).into() })
43}
44
45fn register_finalizer<T>(ptr: *mut T) {
46    if std::mem::needs_drop::<T>() {
47        extern "C" fn finalizer<T>(obj: *mut c_void, _: *mut c_void) {
48            let vec = VecInner(unsafe { NonNull::new_unchecked(obj as *mut T) });
49            let len = vec.num_to_drop();
50            for i in 0..len {
51                unsafe { std::ptr::drop_in_place(vec.as_ptr().add(i)) };
52            }
53        }
54
55        unsafe {
56            gc::GC_register_finalizer(
57                ptr as *mut std::ffi::c_void,
58                Some(finalizer::<T>),
59                std::ptr::null_mut(),
60                std::ptr::null_mut(),
61                std::ptr::null_mut(),
62            );
63        }
64    }
65}
66
67pub struct GcVec<T>(NonNull<[T]>);
68
69impl<T> GcVec<T> {
70    pub fn as_ptr(&self) -> *mut [T] {
71        self.0.as_ptr()
72    }
73
74    pub fn as_ref<'gc>(&self, _token: &'gc impl GcToken) -> &'gc [T] {
75        unsafe { &*self.as_ptr() }
76    }
77
78    #[allow(clippy::mut_from_ref)]
79    pub fn as_mut<'gc>(&mut self, _token: &'gc impl GcToken) -> &'gc mut [T] {
80        unsafe { &mut *self.as_ptr() }
81    }
82
83    /// # Safety
84    /// The returned reference cannot be used in a thread that is not registered with the GC.
85    pub unsafe fn as_ref_unconstrained(&self) -> &'static mut [T] {
86        unsafe { &mut *self.as_ptr() }
87    }
88}
89
90struct VecInner<T>(NonNull<T>);
91
92impl<T> VecInner<T> {
93    fn new(cap: usize) -> Self {
94        if std::mem::needs_drop::<T>() {
95            let padding = usize::max(std::mem::size_of::<usize>(), std::mem::size_of::<T>());
96            let alignment = usize::max(std::mem::align_of::<usize>(), std::mem::align_of::<T>());
97
98            let ptr = unsafe {
99                gc::GC_memalign(
100                    std::mem::size_of::<T>().strict_mul(cap).strict_add(padding),
101                    alignment,
102                ) as *mut T
103            };
104            let ptr = unsafe {
105                NonNull::new(ptr)
106                    .expect("Allocation failed")
107                    .byte_add(padding)
108            };
109            VecInner(ptr)
110        } else {
111            let ptr = unsafe {
112                gc::GC_memalign(
113                    std::mem::size_of::<T>().strict_mul(cap),
114                    std::mem::align_of::<T>(),
115                ) as *mut T
116            };
117            let ptr = NonNull::new(ptr).expect("Allocation failed");
118            VecInner(ptr)
119        }
120    }
121
122    fn as_ptr(&self) -> *mut T {
123        self.0.as_ptr()
124    }
125
126    fn num_to_drop(&self) -> usize {
127        if std::mem::needs_drop::<T>() {
128            let p_len =
129                unsafe { self.0.as_ptr().byte_sub(std::mem::size_of::<usize>()) as *const usize };
130            unsafe { p_len.read() }
131        } else {
132            0
133        }
134    }
135
136    fn set_len(&self, len: usize) {
137        if std::mem::needs_drop::<T>() {
138            let p_len =
139                unsafe { self.0.as_ptr().byte_sub(std::mem::size_of::<usize>()) as *mut usize };
140            unsafe { p_len.write(len) };
141        }
142    }
143}