sequential_id_alloc/
lib.rs

1//! A simple sequential ID allocator that guarantees sequential allocation.
2//!
3//! This crate provides a macro to generate sequential ID allocators that differ from
4//! traditional bitmap/slab allocators by not immediately reusing freed IDs. Instead,
5//! freed IDs are only reused when the allocation pointer wraps around to them again.
6//! This ensures IDs are allocated sequentially and predictably.
7//!
8//! # Example
9//!
10//! ```
11//! use sequential_id_alloc::sequential_id_alloc;
12//!
13//! // Create an allocator for u8 IDs (0-255)
14//! sequential_id_alloc!(MyIdAllocator, u8, 256, u8);
15//!
16//! let mut allocator = MyIdAllocator::default();
17//!
18//! // Allocate IDs sequentially
19//! assert_eq!(allocator.alloc(), Some(0u8));
20//! assert_eq!(allocator.alloc(), Some(1u8));
21//! assert_eq!(allocator.alloc(), Some(2u8));
22//!
23//! // Free an ID - it won't be reused immediately
24//! allocator.dealloc(1u8);
25//! assert_eq!(allocator.alloc(), Some(3u8)); // Gets 3, not 1
26//!
27//! // Check if an ID is allocated
28//! assert!(allocator.contains(0u8));
29//! assert!(!allocator.contains(1u8));
30//!
31//! // Get allocation statistics
32//! assert_eq!(allocator.size(), 3); // Currently 3 IDs allocated
33//! assert!(!allocator.is_full());   // Not all IDs are allocated
34//! ```
35
36pub use bitvec::BitArr;
37
38/// A simple sequential id allocator.
39///
40/// The allocator will allocate new ids sequentially from 0 to u8::MAX. Unlike
41/// other bitmap / slab allocators, freed ids will not be reused right away,
42/// only when the next id pointer wraps around.
43#[macro_export]
44macro_rules! sequential_id_alloc {
45    ($ty:ident, $output_ty:ident, $max:expr, $arr_ty:ident) => {
46        #[derive(Debug)]
47pub struct $ty<T = $output_ty, A = $crate::BitArr![for $max - 1, in $arr_ty]> {
48            next_ptr: usize,
49            bits: A,
50            _output_type: std::marker::PhantomData<T>,
51        }
52
53        impl<T> Default for $ty<T> {
54            fn default() -> Self {
55                Self {
56                    next_ptr: Default::default(),
57                    bits: Default::default(),
58                    _output_type: Default::default(),
59                }
60            }
61        }
62
63        impl<T> $ty<T>
64        where
65            T: Into<usize>,
66            T: TryFrom<usize>,
67        {
68            pub const fn max() -> usize {
69                $max
70            }
71
72            pub fn alloc(&mut self) -> Option<T> {
73                let index = self
74                    .bits
75                    .into_iter()
76                    .enumerate()
77                    .cycle()
78                    .skip(self.next_ptr)
79                    .take($max)
80                    .filter_map(|(i, b)| if !b { Some(i) } else { None })
81                    .next()?;
82
83                self.bits.set(index, true);
84                self.next_ptr = index + 1;
85
86                T::try_from(index).ok()
87            }
88
89            pub fn dealloc(&mut self, id: T) {
90                let id = id.into();
91                self.bits.set(id, false);
92            }
93
94            pub fn contains(&self, id: T) -> bool {
95                let id = id.into();
96                self.bits[id]
97            }
98
99            pub fn is_full(&self) -> bool {
100                self.bits.count_zeros() == 0
101            }
102
103            pub fn size(&self) -> usize {
104                self.bits.count_ones()
105            }
106        }
107    };
108}
109
110#[cfg(test)]
111mod tests {
112    use super::*;
113    use proptest::prelude::*;
114
115    sequential_id_alloc!(SequentialIdAllocU8, u8, 256, u8);
116
117    #[test]
118    fn test_sequential_alloc_1() {
119        let mut ids = SequentialIdAllocU8::default();
120
121        assert!(!ids.contains(0u8));
122        assert_eq!(ids.alloc(), Some(0u8));
123        assert!(ids.contains(0u8));
124        assert_eq!(ids.size(), 1usize);
125
126        ids.dealloc(0);
127        assert_eq!(ids.alloc(), Some(1));
128        assert_eq!(ids.size(), 1);
129    }
130
131    #[test]
132    fn test_sequential_alloc_2() {
133        let mut ids = SequentialIdAllocU8::default();
134
135        assert_eq!(ids.alloc(), Some(0u8));
136        assert_eq!(ids.size(), 1);
137
138        assert_eq!(ids.alloc(), Some(1));
139        assert_eq!(ids.size(), 2);
140
141        ids.dealloc(0);
142        assert_eq!(ids.alloc(), Some(2));
143        assert_eq!(ids.size(), 2);
144    }
145
146    #[test]
147    fn test_wrap() {
148        let mut ids = SequentialIdAllocU8::<u8>::default();
149        for _ in 0..SequentialIdAllocU8::<u8>::max() {
150            assert!(ids.alloc().is_some());
151        }
152
153        assert!(ids.is_full());
154        assert!(ids.alloc().is_none());
155
156        ids.dealloc(5);
157        assert!(!ids.is_full());
158        assert_eq!(ids.alloc(), Some(5));
159    }
160
161    #[test]
162    fn test_arbitrary_inputs() {
163        let alloc = std::cell::RefCell::new(SequentialIdAllocU8::default());
164        proptest!(ProptestConfig::with_cases(1000), |(n in 0..u8::MAX, allocate in proptest::bool::weighted(0.8))| {
165            let mut alloc = alloc.borrow_mut();
166            let size = alloc.size();
167            let is_full = alloc.is_full();
168            if allocate {
169                let n = alloc.alloc();
170                match (is_full, n) {
171                    (true, None) => {
172                        assert_eq!(alloc.size(), size);
173                    }
174                    (false, Some(n)) => {
175                        assert_eq!(alloc.size(), size + 1);
176                        assert!(alloc.contains(n));
177                    }
178                    _ => unreachable!(),
179                }
180            } else if alloc.contains(n) {
181                alloc.dealloc(n);
182                assert_eq!(alloc.size(), size - 1);
183                assert!(!alloc.contains(n));
184                assert!(!alloc.is_full());
185            }
186        });
187    }
188}