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 pool = MemoryPool::<usize>::with_capacity(2);
121        let mut a = vec![];
122        for i in 0..100 {
123            let p = pool.allocate(i);
124            a.push(p);
125            for (i, &p) in a.iter().enumerate() {
126                assert_eq!(unsafe { *p.as_ref() }, i);
127            }
128        }
129    }
130
131    #[test]
132    fn test_drop() {
133        #[derive(Debug)]
134        struct CheckDrop<T>(T);
135        thread_local! {
136            static CNT: RefCell<usize> = const { RefCell::new(0) };
137        }
138        impl<T> Drop for CheckDrop<T> {
139            fn drop(&mut self) {
140                CNT.with(|cnt| *cnt.borrow_mut() += 1);
141            }
142        }
143        const Q: usize = 100;
144        let mut cnt = 0usize;
145        let mut rng = Xorshift::default();
146        for _ in 0..10 {
147            let mut pool = MemoryPool::new();
148            let mut a = vec![];
149            for _ in 0..Q {
150                if a.is_empty() || rng.gen_bool(0.8) {
151                    let k = rng.rand(!0);
152                    a.push(pool.allocate(CheckDrop(k)));
153                } else {
154                    let i = rng.rand(a.len() as _) as usize;
155                    let p = a.swap_remove(i);
156                    pool.deallocate(p);
157                    cnt += 1;
158                }
159                assert_eq!(cnt, CNT.with(|cnt| *cnt.borrow()));
160            }
161            cnt += a.len();
162        }
163        assert_eq!(cnt, CNT.with(|cnt| *cnt.borrow()));
164    }
165}
competitive/data_structure/allocator.rs

competitive/data_structure/
allocator.rs
