Last active
February 22, 2026 19:47
-
-
Save oconnor663/6ea41052c8229e75840522d5167f1593 to your computer and use it in GitHub Desktop.
ThreeRena
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| use std::alloc::{Layout, alloc, dealloc}; | |
| use std::mem::{self, MaybeUninit}; | |
| use std::ops::{Deref, DerefMut}; | |
| use std::slice; | |
| pub struct FixedVec<'a, T> { | |
| buf: &'a mut [MaybeUninit<T>], | |
| len: usize, | |
| } | |
| impl<'a, T> FixedVec<'a, T> { | |
| pub fn new(buf: &'a mut [MaybeUninit<T>]) -> Self { | |
| Self { buf, len: 0 } | |
| } | |
| pub fn push(&mut self, value: T) { | |
| assert!(self.len < self.buf.len()); | |
| self.buf[self.len].write(value); | |
| self.len += 1; | |
| } | |
| pub fn pop(&mut self) -> Option<T> { | |
| if self.len == 0 { | |
| return None; | |
| } | |
| self.len -= 1; | |
| Some(unsafe { self.buf[self.len].assume_init_read() }) | |
| } | |
| } | |
| impl<'a, T> Deref for FixedVec<'a, T> { | |
| type Target = [T]; | |
| fn deref(&self) -> &[T] { | |
| unsafe { std::slice::from_raw_parts(self.buf.as_ptr() as *const T, self.len) } | |
| } | |
| } | |
| impl<'a, T> DerefMut for FixedVec<'a, T> { | |
| fn deref_mut(&mut self) -> &mut [T] { | |
| unsafe { std::slice::from_raw_parts_mut(self.buf.as_mut_ptr() as *mut T, self.len) } | |
| } | |
| } | |
| impl<'a, T> Drop for FixedVec<'a, T> { | |
| fn drop(&mut self) { | |
| for i in 0..self.len { | |
| // SAFETY: If one of these panics, the others will be leaked, but that's ok. | |
| unsafe { self.buf[i].assume_init_drop() }; | |
| } | |
| } | |
| } | |
| unsafe impl<'a, T> Send for FixedVec<'a, T> where T: Send {} | |
| unsafe impl<'a, T> Sync for FixedVec<'a, T> where T: Sync {} | |
| pub struct ThreeRena<A, B, C> { | |
| base_ptr: *mut u8, | |
| base_layout: Layout, | |
| a_ptr: *mut MaybeUninit<A>, | |
| b_ptr: *mut MaybeUninit<B>, | |
| c_ptr: *mut MaybeUninit<C>, | |
| capacity: usize, | |
| } | |
| impl<A, B, C> ThreeRena<A, B, C> { | |
| pub fn new(capacity: usize) -> Self { | |
| let mut size = 0; | |
| // It would be optimal to sort A, B, and C in alignment order, so that we would never need | |
| // to pad in between them. But let's just keep them in order for simplicity. | |
| let a_offset = 0; | |
| size += capacity * mem::size_of::<A>(); | |
| let b_align = mem::align_of::<B>(); | |
| if size % b_align != 0 { | |
| size += b_align - (size % b_align); | |
| } | |
| let b_offset = size; | |
| size += capacity * mem::size_of::<B>(); | |
| let c_align = mem::align_of::<C>(); | |
| if size % c_align != 0 { | |
| size += c_align - (size % c_align); | |
| } | |
| let c_offset = size; | |
| size += capacity * mem::size_of::<C>(); | |
| let a_align = mem::align_of::<A>(); | |
| let max_alignment = [a_align, b_align, c_align].into_iter().max().unwrap(); | |
| let base_layout = Layout::from_size_align(size, max_alignment).unwrap(); | |
| unsafe { | |
| let base_ptr = alloc(base_layout); | |
| Self { | |
| base_ptr, | |
| base_layout, | |
| a_ptr: base_ptr.add(a_offset) as *mut _, | |
| b_ptr: base_ptr.add(b_offset) as *mut _, | |
| c_ptr: base_ptr.add(c_offset) as *mut _, | |
| capacity, | |
| } | |
| } | |
| } | |
| // Note that the `FixedVec`s free all their contents (but not their backing memory) when they | |
| // drop. You can call `.vecs()` multiple times on the same `ThreeRena`, but you get empty | |
| // `FixedVec`s each time. | |
| pub fn vecs<'a>(&'a mut self) -> (FixedVec<'a, A>, FixedVec<'a, B>, FixedVec<'a, C>) { | |
| unsafe { | |
| ( | |
| FixedVec::new(slice::from_raw_parts_mut(self.a_ptr, self.capacity)), | |
| FixedVec::new(slice::from_raw_parts_mut(self.b_ptr, self.capacity)), | |
| FixedVec::new(slice::from_raw_parts_mut(self.c_ptr, self.capacity)), | |
| ) | |
| } | |
| } | |
| } | |
| impl<A, B, C> Drop for ThreeRena<A, B, C> { | |
| fn drop(&mut self) { | |
| unsafe { | |
| dealloc(self.base_ptr, self.base_layout); | |
| } | |
| } | |
| } | |
| unsafe impl<A, B, C> Send for ThreeRena<A, B, C> | |
| where | |
| A: Send, | |
| B: Send, | |
| C: Send, | |
| { | |
| } | |
| unsafe impl<A, B, C> Sync for ThreeRena<A, B, C> | |
| where | |
| A: Sync, | |
| B: Sync, | |
| C: Sync, | |
| { | |
| } | |
| #[test] | |
| fn test_threerena() { | |
| let mut arena = ThreeRena::<bool, i32, String>::new(3); | |
| let (mut bools, mut ints, mut strings) = arena.vecs(); | |
| bools.push(true); | |
| ints.push(1); | |
| strings.push("hello".to_string()); | |
| bools.push(false); | |
| ints.push(2); | |
| strings.push("world".to_string()); | |
| bools.push(true); | |
| ints.push(3); | |
| strings.push("!".to_string()); | |
| assert_eq!(bools[..], [true, false, true]); | |
| assert_eq!(ints[..], [1, 2, 3]); | |
| assert_eq!(strings[..], ["hello", "world", "!"]); | |
| assert_eq!(bools.pop(), Some(true)); | |
| assert_eq!(ints.pop(), Some(3)); | |
| assert_eq!(strings.pop().as_deref(), Some("!")); | |
| assert_eq!(bools[..], [true, false]); | |
| assert_eq!(ints[..], [1, 2]); | |
| assert_eq!(strings[..], ["hello", "world"]); | |
| // Dropping the `FixedVec`s frees their contents, and dropping the `ThreeRena` frees backing | |
| // allocation. If you comment out any of the `Drop` impls above, Miri will detect a memory leak | |
| // at this point. | |
| } |
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment