competitive/data_structure/
allocator.rs

1use std::{
2    marker::PhantomData,
3    mem::{replace, size_of, take},
4    ptr::{self, NonNull, read, write},
5};
6
7pub trait Allocator<T> {
8    fn allocate(&mut self, value: T) -> NonNull<T>;
9    fn deallocate(&mut self, ptr: NonNull<T>) -> T;
10}
11
12#[derive(Debug)]
13pub struct MemoryPool<T> {
14    pool: Vec<T>,
15    chunks: Vec<Vec<T>>,
16    unused: Vec<NonNull<T>>,
17}
18
19impl<T> Default for MemoryPool<T> {
20    fn default() -> Self {
21        Self::with_capacity(CAP / 1usize.max(size_of::<T>()))
22    }
23}
24
25const CAP: usize = 1024;
26
27impl<T> MemoryPool<T> {
28    pub fn new() -> Self {
29        Default::default()
30    }
31    pub fn with_capacity(capacity: usize) -> Self {
32        let pool = Vec::with_capacity(capacity.max(1));
33        Self {
34            pool,
35            chunks: Vec::new(),
36            unused: Vec::new(),
37        }
38    }
39}
40
41impl<T> Drop for MemoryPool<T> {
42    fn drop(&mut self) {
43        self.chunks.push(take(&mut self.pool));
44        let mut removed = vec![vec![]; self.chunks.len()];
45        for p in self.unused.iter() {
46            let p = p.as_ptr();
47            for (chunk, removed) in self.chunks.iter().zip(&mut removed).rev() {
48                let ptr = chunk.as_ptr() as *mut _;
49                let len = chunk.len();
50                if ptr <= p && p < ptr.wrapping_add(len) {
51                    removed.push(p);
52                }
53            }
54        }
55        for (chunk, removed) in self.chunks.iter_mut().zip(&mut removed) {
56            removed.sort_unstable();
57            for &p in removed.iter() {
58                unsafe {
59                    let len = chunk.len();
60                    let base_ptr = chunk.as_mut_ptr();
61                    ptr::copy(base_ptr.add(len - 1), p, 1);
62                    chunk.set_len(len - 1);
63                }
64            }
65        }
66    }
67}
68
69impl<T> Allocator<T> for MemoryPool<T> {
70    fn allocate(&mut self, value: T) -> NonNull<T> {
71        if let Some(mut ptr) = self.unused.pop() {
72            unsafe { write(ptr.as_mut(), value) };
73            ptr
74        } else {
75            let len = self.pool.len();
76            if len >= self.pool.capacity() {
77                let new_capacity = self.pool.capacity() * 2;
78                let new_pool = Vec::with_capacity(new_capacity);
79                self.chunks.push(replace(&mut self.pool, new_pool));
80            }
81            let len = self.pool.len();
82            debug_assert!(len < self.pool.capacity());
83            self.pool.push(value);
84            unsafe { NonNull::new_unchecked(self.pool.as_mut_ptr().add(len)) }
85        }
86    }
87
88    fn deallocate(&mut self, ptr: NonNull<T>) -> T {
89        self.unused.push(ptr);
90        unsafe { read(ptr.as_ptr()) }
91    }
92}
93
94#[derive(Debug)]
95pub struct BoxAllocator<T>(PhantomData<fn() -> T>);
96
97impl<T> Default for BoxAllocator<T> {
98    fn default() -> Self {
99        Self(PhantomData)
100    }
101}
102
103impl<T> Allocator<T> for BoxAllocator<T> {
104    fn allocate(&mut self, value: T) -> NonNull<T> {
105        unsafe { NonNull::new_unchecked(Box::leak(Box::new(value))) }
106    }
107    fn deallocate(&mut self, ptr: NonNull<T>) -> T {
108        unsafe { *Box::from_raw(ptr.as_ptr()) }
109    }
110}
111
112#[cfg(test)]
113mod tests {
114    use super::*;
115    use crate::tools::Xorshift;
116    use std::cell::RefCell;
117
118    #[test]
119    fn test_alloc() {
120        let mut rng = Xorshift::default();
121        for _ in 0..100 {
122            let mut pool = MemoryPool::with_capacity(rng.random(0..=32));
123            let mut live: Vec<(NonNull<u64>, u64)> = Vec::new();
124            for _ in 0..1000 {
125                if !live.is_empty() && rng.random(0..3) == 0 {
126                    let i = rng.random(0..live.len());
127                    let (ptr, value) = live.swap_remove(i);
128                    assert_eq!(pool.deallocate(ptr), value);
129                } else {
130                    let value = rng.rand64();
131                    live.push((pool.allocate(value), value));
132                }
133                for &(ptr, value) in &live {
134                    assert_eq!(unsafe { *ptr.as_ref() }, value);
135                }
136            }
137        }
138    }
139
140    #[test]
141    fn test_drop() {
142        #[derive(Debug)]
143        struct CheckDrop<T>(T);
144        thread_local! {
145            static CNT: RefCell<usize> = const { RefCell::new(0) };
146        }
147        impl<T> Drop for CheckDrop<T> {
148            fn drop(&mut self) {
149                CNT.with(|cnt| *cnt.borrow_mut() += 1);
150            }
151        }
152        const Q: usize = 100;
153        let mut cnt = 0usize;
154        let mut rng = Xorshift::default();
155        for _ in 0..10 {
156            let mut pool = MemoryPool::new();
157            let mut a = vec![];
158            for _ in 0..Q {
159                if a.is_empty() || rng.gen_bool(0.8) {
160                    let k = rng.rand(!0);
161                    a.push(pool.allocate(CheckDrop(k)));
162                } else {
163                    let i = rng.rand(a.len() as _) as usize;
164                    let p = a.swap_remove(i);
165                    pool.deallocate(p);
166                    cnt += 1;
167                }
168                assert_eq!(cnt, CNT.with(|cnt| *cnt.borrow()));
169            }
170            cnt += a.len();
171        }
172        assert_eq!(cnt, CNT.with(|cnt| *cnt.borrow()));
173    }
174}
competitive/data_structure/allocator.rs

competitive/data_structure/
allocator.rs
