Created
August 27, 2024 23:53
-
-
Save rksm/d559d47ed2147c1fbfeefde8a3ef65bd to your computer and use it in GitHub Desktop.
generic SequentialIdAlloc
This file contains 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 bitvec::prelude::*; | |
| /// A simple sequential id allocator. | |
| /// | |
| /// The allocator will allocate new ids sequentially from 0 to u8::MAX. Unlike | |
| /// other bitmap / slab allocators, freed ids will not be reused right away, | |
| /// only when the next id pointer wraps around. | |
| pub struct SequentialIdAlloc<const MAX: usize, T = u8, U = BitArr![for 255, in u8]> { | |
| next_ptr: usize, | |
| bits: U, | |
| _output_type: std::marker::PhantomData<T>, | |
| } | |
| pub type SequentialIdAllocU8<const MAX: usize> = SequentialIdAlloc<256, u8>; | |
| impl<const MAX: usize, T> Default for SequentialIdAlloc<MAX, T> { | |
| fn default() -> Self { | |
| Self { | |
| next_ptr: Default::default(), | |
| bits: Default::default(), | |
| _output_type: Default::default(), | |
| } | |
| } | |
| } | |
| impl<const MAX: usize, T> SequentialIdAlloc<MAX, T> | |
| where | |
| T: Into<usize>, | |
| T: TryFrom<usize>, | |
| <T as TryFrom<usize>>::Error: std::fmt::Debug, | |
| { | |
| pub fn alloc(&mut self) -> Option<T> { | |
| let index = self | |
| .bits | |
| .iter() | |
| .enumerate() | |
| .cycle() | |
| .skip(self.next_ptr) | |
| .take(MAX) | |
| .filter_map(|(i, b)| if !b { Some(i) } else { None }) | |
| .next()?; | |
| self.bits.set(index, true); | |
| self.next_ptr = index + 1; | |
| Some(T::try_from(index).unwrap()) | |
| } | |
| pub fn dealloc(&mut self, id: T) { | |
| let id = id.into(); | |
| self.bits.set(id, false); | |
| } | |
| pub fn contains(&self, id: T) -> bool { | |
| let id = id.into(); | |
| self.bits[id] | |
| } | |
| pub fn is_full(&self) -> bool { | |
| self.bits.count_zeros() == 0 | |
| } | |
| pub fn size(&self) -> usize { | |
| self.bits.count_ones() | |
| } | |
| pub const fn max() -> usize { | |
| MAX | |
| } | |
| } | |
| #[cfg(test)] | |
| mod tests { | |
| use super::*; | |
| use proptest::prelude::*; | |
| #[test] | |
| fn test_sequential_alloc_1() { | |
| let mut ids = SequentialIdAllocU8::default(); | |
| assert!(!ids.contains(0u8)); | |
| assert_eq!(ids.alloc(), Some(0)); | |
| assert!(ids.contains(0)); | |
| assert_eq!(ids.size(), 1); | |
| ids.dealloc(0); | |
| assert_eq!(ids.alloc(), Some(1)); | |
| assert_eq!(ids.size(), 1); | |
| } | |
| #[test] | |
| fn test_sequential_alloc_2() { | |
| let mut ids = SequentialIdAllocU8::default(); | |
| assert_eq!(ids.alloc(), Some(0)); | |
| assert_eq!(ids.size(), 1); | |
| assert_eq!(ids.alloc(), Some(1)); | |
| assert_eq!(ids.size(), 2); | |
| ids.dealloc(0); | |
| assert_eq!(ids.alloc(), Some(2)); | |
| assert_eq!(ids.size(), 2); | |
| } | |
| #[test] | |
| fn test_wrap() { | |
| let mut ids = SequentialIdAllocU8::default(); | |
| for _ in 0..SequentialIdAllocU8::max() { | |
| assert!(ids.alloc().is_some()); | |
| } | |
| assert!(ids.is_full()); | |
| assert!(ids.alloc().is_none()); | |
| ids.dealloc(5); | |
| assert!(!ids.is_full()); | |
| assert_eq!(ids.alloc(), Some(5)); | |
| } | |
| #[test] | |
| fn test_arbitrary_inputs() { | |
| let alloc = std::cell::RefCell::new(SequentialIdAllocU8::default()); | |
| proptest!(ProptestConfig::with_cases(1000), |(n in 0..u8::MAX, allocate in proptest::bool::weighted(0.8))| { | |
| let mut alloc = alloc.borrow_mut(); | |
| let size = alloc.size(); | |
| let is_full = alloc.is_full(); | |
| if allocate { | |
| let n = alloc.alloc(); | |
| match (is_full, n) { | |
| (true, None) => { | |
| assert_eq!(alloc.size(), size); | |
| } | |
| (false, Some(n)) => { | |
| assert_eq!(alloc.size(), size + 1); | |
| assert!(alloc.contains(n)); | |
| } | |
| _ => unreachable!(), | |
| } | |
| } else if alloc.contains(n) { | |
| alloc.dealloc(n); | |
| assert_eq!(alloc.size(), size - 1); | |
| assert!(!alloc.contains(n)); | |
| assert!(!alloc.is_full()); | |
| } | |
| }); | |
| } | |
| } |
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment