1
use std::alloc::Layout;
2
use std::array;
3
use std::cell::Cell;
4
use std::cell::UnsafeCell;
5
use std::marker::PhantomData;
6
use std::mem::ManuallyDrop;
7
use std::ptr::NonNull;
8
use std::sync::Mutex;
9
use std::sync::MutexGuard;
10

            
11
use allocator_api2::alloc::AllocError;
12
use allocator_api2::alloc::Allocator;
13
use thread_local::ThreadLocal;
14

            
15
use crate::FreeList;
16
use crate::FreeListEntry;
17

            
18
/// The number of entries in every free list chunk.
19
const FREE_LIST_CHUNK_SIZE: usize = 1000;
20

            
21
/// This is a memory pool or also called fixed-size block allocator for a
22
/// concrete type `T`. It stores blocks of `N` to minimize the overhead of
23
/// individual memory allocations, which are typically in the range of one or
24
/// two words.
25
///
26
/// Behaves like `Allocator`, except that it only allocates for layouts of `T`.
27
/// Requires periodic calls to `remove_free_blocks` to prevent memory usage from
28
/// growing indefinitely.
29
///
30
/// This allocator minimizes contention by maintaining per-thread state for
31
/// the common allocation/deallocation paths and only takes a lock when a new
32
/// block needs to be allocated. This does mean that external synchronisation
33
/// is required to prevent concurrent allocations overlapping with `remove_free_blocks`.
34
///
35
/// A single thread-local freelist is maintained per thread:
36
/// - `free`: popped from during allocation and pushed to during deallocation.
37
pub struct BlockAllocator<T: Send, const N: usize> {
38
    /// Owns the block list; only locked when a new block must be allocated.
39
    /// This is the only shared state accessed during allocation.
40
    blocks: Mutex<BlockList<T, N>>,
41

            
42
    /// Per-thread state for allocation: current block and bump offset. This eliminates
43
    /// contention in the common path.
44
    alloc_state: ThreadLocal<ThreadLocalAllocState<T, N>>,
45
}
46

            
47
/// The block list and bump pointer, protected by the blocks mutex.
48
struct BlockList<T, const N: usize> {
49
    /// The block that is currently being bump-allocated from. We avoid Box here
50
    /// to allow multiple blocks to point to the same next block without
51
    /// violating Box's noalias.
52
    head_block: Option<NonNull<Block<T, N>>>,
53

            
54
    /// Shared chunks of free entries represented by list heads. Each head
55
    /// points to a null-terminated list of length `FREE_LIST_CHUNK_SIZE`
56
    /// (except possibly the final chunk).
57
    free_chunks: Vec<NonNull<Entry<T>>>,
58
}
59

            
60
impl<T, const N: usize> Drop for BlockList<T, N> {
61
101
    fn drop(&mut self) {
62
        // Drop the head block; Block's Drop impl recursively drops the list.
63
101
        if let Some(block_ptr) = self.head_block.take() {
64
101
            // Safety: we own all blocks in the list.
65
101
            unsafe { drop(Box::from_raw(block_ptr.as_ptr())) };
66
101
        }
67
101
    }
68
}
69

            
70
impl<T: Send, const N: usize> Default for BlockAllocator<T, N> {
71
    fn default() -> Self {
72
        Self::new()
73
    }
74
}
75

            
76
impl<T: Send, const N: usize> BlockAllocator<T, N> {
77
17741
    pub fn new() -> Self {
78
17741
        Self {
79
17741
            blocks: Mutex::new(BlockList {
80
17741
                head_block: None,
81
17741
                free_chunks: Vec::new(),
82
17741
            }),
83
17741
            alloc_state: ThreadLocal::new(),
84
17741
        }
85
17741
    }
86

            
87
    /// Allocates a slot for one object of type `T`.
88
20276989
    pub fn allocate_object(&self) -> Result<NonNull<T>, AllocError> {
89
20276989
        let state = self.alloc_state.get_or(ThreadLocalAllocState::new);
90

            
91
        // Fast path 1: try thread-local free list.
92
20276989
        if let Some(entry) = state.free.try_pop() {
93
13833618
            return Ok(entry.cast());
94
6443371
        }
95

            
96
        // Fast path 2: lock-free bump allocation from current thread's block.
97
6443371
        let block_ptr = state.current_block.get();
98
6443371
        let offset = state.bump_offset.get();
99
6443371
        if !block_ptr.is_null() && offset < N {
100
5976545
            state.bump_offset.set(offset + 1);
101
            return unsafe {
102
5976545
                let data_ptr = (*block_ptr).data.get() as *mut Entry<T>;
103
5976545
                let entry_ptr = data_ptr.add(offset);
104
5976545
                Ok(NonNull::new_unchecked(
105
5976545
                    std::ptr::addr_of_mut!((*entry_ptr).data) as *mut T
106
5976545
                ))
107
            };
108
466826
        }
109

            
110
        // Slow path: acquire the lock once and reuse it across refill /
111
        // new-block allocation.
112
466826
        let mut guard = self.blocks.lock().expect("Lock poisoned");
113

            
114
466826
        if self.refill_local_free_from_chunks(state, &mut guard)
115
450203
            && let Some(entry) = state.free.try_pop()
116
        {
117
450203
            return Ok(entry.cast());
118
16623
        }
119

            
120
16623
        self.allocate_new_block(state, guard)
121
20276989
    }
122

            
123
    /// Refills the calling thread's local freelist from one shared chunk.
124
466826
    fn refill_local_free_from_chunks(
125
466826
        &self,
126
466826
        state: &ThreadLocalAllocState<T, N>,
127
466826
        guard: &mut MutexGuard<'_, BlockList<T, N>>,
128
466826
    ) -> bool {
129
466826
        let Some(current) = guard.free_chunks.pop() else {
130
16623
            return false;
131
        };
132

            
133
450203
        debug_assert!(
134
450203
            state.free.is_empty(),
135
            "local freelist must be empty before chunk refill"
136
        );
137
450203
        unsafe {
138
450203
            state.free.set_head(current.as_ptr());
139
450203
        }
140
450203
        true
141
466826
    }
142

            
143
    /// Slow path: allocate a new block and update thread-local state.
144
    #[cold]
145
16623
    fn allocate_new_block(
146
16623
        &self,
147
16623
        state: &ThreadLocalAllocState<T, N>,
148
16623
        mut guard: MutexGuard<'_, BlockList<T, N>>,
149
16623
    ) -> Result<NonNull<T>, AllocError> {
150
        // Allocate a new block and link it to the existing list.
151
16623
        let mut new_block = Block::new();
152
16623
        new_block.next = guard.head_block;
153
16623
        let new_block_ptr = unsafe { NonNull::new_unchecked(Box::into_raw(Box::new(new_block))) };
154
16623
        guard.head_block = Some(new_block_ptr);
155

            
156
16623
        drop(guard);
157

            
158
        // Update thread-local state with new block.
159
16623
        state.current_block.set(new_block_ptr.as_ptr());
160
16623
        state.bump_offset.set(1);
161

            
162
        // Return slot 0 of the new block.
163
        unsafe {
164
16623
            let data_ptr = (*new_block_ptr.as_ptr()).data.get() as *mut Entry<T>;
165
16623
            Ok(NonNull::new_unchecked(
166
16623
                std::ptr::addr_of_mut!((*data_ptr).data) as *mut T
167
16623
            ))
168
        }
169
16623
    }
170

            
171
    /// Deallocates a previously-allocated pointer.
172
19805982
    pub fn deallocate_object(&self, ptr: NonNull<T>) {
173
19805982
        let state = self.alloc_state.get_or(ThreadLocalAllocState::new);
174
        // SAFETY: `ptr` was returned by `allocate_object`, so it points to a block entry that
175
        // stays valid until `remove_free_blocks` reclaims it.
176
19805982
        unsafe {
177
19805982
            state.free.push(ptr.cast());
178
19805982
        }
179
19805982
    }
180

            
181
    /// Removes empty blocks from the block list.
182
    ///
183
    /// Should be called periodically to prevent memory usage from growing
184
    /// indefinitely.
185
    ///
186
    /// **Important**: This method must not be called concurrently with any
187
    /// allocations or deallocations. It should be called at a synchronization
188
    /// point where no other threads are active on this allocator.
189
    ///
190
    /// This method scans every thread-local `free` list, marks those entries
191
    /// with a sentinel value, then removes blocks where all entries are marked
192
    /// as free.
193
    ///
194
    /// Returns `(removed_blocks, free_size)`: the number of blocks removed, and
195
    /// the size of the merged `free` list before cleanup.
196
5582400
    pub fn remove_free_blocks(&mut self) -> (usize, usize)
197
5582400
    where
198
5582400
        T: BlockAllocatorSafe,
199
    {
200
        // Mark all elements in all thread-local freelists with a special
201
        // value that none of the live entries can have.
202
5582400
        let nonexisting_value = NONEXISTING_VALUE as *mut Entry<T>;
203

            
204
        // In debug mode, collect all freelist entry pointers.
205
        #[cfg(debug_assertions)]
206
5582400
        let mut freelist_ptrs: std::collections::HashSet<*mut Entry<T>> = std::collections::HashSet::new();
207

            
208
        // Helper: walk a freelist and mark every entry with the sentinel.
209
        // We only update previous entries to ensure that iter keeps working.
210
        // Returns the number of entries in the freelist.
211
5582400
        let mark_freelist =
212
            |list: &FreeList<Entry<T>>,
213
             #[cfg(debug_assertions)] ptrs: &mut std::collections::HashSet<*mut Entry<T>>| unsafe {
214
2862680
                let mut previous: Option<NonNull<Entry<T>>> = None;
215
2862680
                let mut count: usize = 0;
216
188285216
                for current in list.iter() {
217
                    #[cfg(debug_assertions)]
218
188285216
                    ptrs.insert(current.as_ptr());
219

            
220
188285216
                    if let Some(previous) = previous {
221
187343642
                        *(*previous.as_ptr()).next = nonexisting_value;
222
187343642
                    }
223
188285216
                    previous = Some(current);
224
188285216
                    count += 1;
225
                }
226

            
227
2862680
                if let Some(previous) = previous {
228
941574
                    *(*previous.as_ptr()).next = nonexisting_value;
229
1921206
                }
230
2862680
                count
231
2862680
            };
232

            
233
5582400
        let mut free_size = 0;
234
5582400
        for state in self.alloc_state.iter_mut() {
235
2862680
            free_size += mark_freelist(
236
2862680
                &state.free,
237
2862680
                #[cfg(debug_assertions)]
238
2862680
                &mut freelist_ptrs,
239
2862680
            );
240
2862680
            state.free.clear();
241
2862680
            // Keep the thread-local bump pointer state; those entries will be
242
2862680
            // reclaimed when we walk the blocks next.
243
2862680
        }
244

            
245
5582400
        let mut guard = self.blocks.lock().expect("Lock poisoned");
246

            
247
        // Mark entries currently staged in shared free chunks.
248
5582400
        for &chunk_head in &guard.free_chunks {
249
2604366
            let mut current = Some(chunk_head);
250
1259744014
            while let Some(entry) = current {
251
1257139648
                #[cfg(debug_assertions)]
252
1257139648
                freelist_ptrs.insert(entry.as_ptr());
253
1257139648

            
254
1257139648
                let next = unsafe { Entry::get_next(entry.as_ptr()) };
255
1257139648
                unsafe {
256
1257139648
                    *(*entry.as_ptr()).next = nonexisting_value;
257
1257139648
                }
258
1257139648
                free_size += 1;
259
1257139648
                current = NonNull::new(next);
260
1257139648
            }
261
        }
262
5582400
        guard.free_chunks.clear();
263

            
264
        // Debug check: verify that no live entry has the sentinel value.
265
        #[cfg(debug_assertions)]
266
        {
267
5585500
            for block_ptr in Self::iter_blocks(&guard) {
268
2851486
                let data = unsafe { &*(*block_ptr.as_ptr()).data.get() };
269
2916747264
                for entry in data {
270
2916747264
                    let entry_ptr = entry as *const Entry<T> as *mut Entry<T>;
271
2916747264
                    if !freelist_ptrs.contains(&entry_ptr) {
272
                        // This entry is live — it must not look like the sentinel.
273
                        unsafe {
274
1471322400
                            debug_assert!(
275
1471322400
                                !std::ptr::eq(*entry.next, nonexisting_value),
276
                                "Live entry at {entry_ptr:?} has the sentinel value (null in first word). \
277
                                 This violates the BlockAllocatorSafe contract."
278
                            );
279
                        }
280
1445424864
                    }
281
                }
282
            }
283
        }
284

            
285
5582400
        let removed = if guard.head_block.is_some() {
286
            // Remove blocks that are now empty, i.e., all their entries have nonexisting_value.
287
2848029
            let mut prev_next_field: *mut Option<NonNull<Block<T, N>>> = &mut guard.head_block;
288
2848029
            let mut removed_blocks = 0;
289

            
290
5699515
            while let Some(current_ptr) = unsafe { *prev_next_field } {
291
2851486
                let all_free = unsafe {
292
2851486
                    let data = &*(*current_ptr.as_ptr()).data.get();
293
864620728
                    data.iter().all(|entry| std::ptr::eq(*entry.next, nonexisting_value))
294
                };
295

            
296
2851486
                if all_free {
297
4854
                    // Unlink and drop the current block.
298
4854
                    let next = unsafe { (*current_ptr.as_ptr()).next.take() };
299
4854
                    unsafe { *prev_next_field = next };
300
4854
                    unsafe { drop(Box::from_raw(current_ptr.as_ptr())) };
301
4854
                    removed_blocks += 1;
302
4854
                    // prev_next_field stays the same — it now points to the next block.
303
2846632
                } else {
304
2846632
                    // Keep this block; advance to next.
305
2846632
                    prev_next_field = unsafe { &mut (*current_ptr.as_ptr()).next };
306
2846632
                }
307
            }
308

            
309
2848029
            removed_blocks
310
        } else {
311
            // No blocks, nothing to remove.
312
2734371
            0
313
        };
314

            
315
        // Recreate shared free chunks from remaining blocks.
316
5582400
        let mut rebuilt_chunk_heads: Vec<NonNull<Entry<T>>> = Vec::new();
317
5582400
        let mut chunk_head: Option<NonNull<Entry<T>>> = None;
318
5582400
        let mut chunk_tail: Option<NonNull<Entry<T>>> = None;
319
5582400
        let mut chunk_len = 0usize;
320

            
321
5585500
        for block_ptr in Self::iter_blocks(&guard) {
322
            // Walk entries via raw pointers; deriving `*mut` from `&Entry<T>`
323
            // would violate Stacked Borrows when we write through that pointer.
324
2846632
            let data_ptr = unsafe { (*block_ptr.as_ptr()).data.get() as *mut Entry<T> };
325
2911776768
            for i in 0..N {
326
                unsafe {
327
2911776768
                    let entry_ptr = data_ptr.add(i);
328
2911776768
                    if std::ptr::eq(*(*entry_ptr).next, nonexisting_value) {
329
1440454368
                        *(*entry_ptr).next = std::ptr::null_mut();
330
1440454368
                        let entry_ptr = NonNull::new_unchecked(entry_ptr);
331

            
332
1440454368
                        if let Some(tail) = chunk_tail {
333
1437397883
                            *(*tail.as_ptr()).next = entry_ptr.as_ptr();
334
1437397883
                            chunk_tail = Some(entry_ptr);
335
1437397883
                            chunk_len += 1;
336
1437397883
                        } else {
337
3056485
                            chunk_head = Some(entry_ptr);
338
3056485
                            chunk_tail = Some(entry_ptr);
339
3056485
                            chunk_len = 1;
340
3056485
                        }
341

            
342
1440454368
                        if chunk_len == FREE_LIST_CHUNK_SIZE {
343
292869
                            rebuilt_chunk_heads.push(chunk_head.expect("chunk head set"));
344
292869
                            chunk_head = None;
345
292869
                            chunk_tail = None;
346
292869
                            chunk_len = 0;
347
1440161499
                        }
348
1471322400
                    }
349
                }
350
            }
351
        }
352

            
353
5582400
        if let Some(head) = chunk_head {
354
2763616
            rebuilt_chunk_heads.push(head);
355
2818884
        }
356

            
357
5582400
        guard.free_chunks.extend(rebuilt_chunk_heads);
358

            
359
5582400
        drop(guard);
360
5582400
        (removed, free_size)
361
5582400
    }
362

            
363
    /// Returns an iterator over the blocks.
364
    ///
365
    /// The caller must pass the already-acquired guard to avoid a deadlock.
366
11164800
    fn iter_blocks<'a>(guard: &'a MutexGuard<'_, BlockList<T, N>>) -> BlockIter<'a, T, N> {
367
11164800
        BlockIter {
368
11164800
            current: guard.head_block,
369
11164800
            _marker: PhantomData,
370
11164800
        }
371
11164800
    }
372
}
373

            
374
/// Per-thread state for allocation: current block and bump offset.
375
/// This eliminates contention in the hot path.
376
struct ThreadLocalAllocState<T, const N: usize> {
377
    /// The block currently being bump-allocated from.
378
    current_block: Cell<*mut Block<T, N>>,
379

            
380
    /// Bump offset within `current_block`. N or greater means the block is full.
381
    bump_offset: Cell<usize>,
382

            
383
    /// Thread-local free list (only popped from during allocation).
384
    free: FreeList<Entry<T>>,
385
}
386

            
387
impl<T, const N: usize> ThreadLocalAllocState<T, N> {
388
8680
    fn new() -> Self {
389
8680
        Self {
390
8680
            current_block: Cell::new(std::ptr::null_mut()),
391
8680
            bump_offset: Cell::new(N),
392
8680
            free: FreeList::new(),
393
8680
        }
394
8680
    }
395
}
396

            
397
unsafe impl<T: Send, const N: usize> Send for ThreadLocalAllocState<T, N> {}
398

            
399
/// Implementing this trait for a type `T` asserts that the special sentinel
400
/// never occurs as a valid entry.
401
///
402
/// # Safety
403
///
404
/// Marker trait asserting two properties of `T`, both required because
405
/// [`BlockAllocator::remove_free_blocks`] reads the first
406
/// `size_of::<*mut _>()` bytes of every *live* entry through the freelist
407
/// union:
408
///
409
/// - The sentinel value—a pointer value where all bytes are set to `0xFF`
410
///   (i.e., `usize::MAX` cast to `*mut Entry<T>`)—can never appear as the
411
///   first `size_of::<*mut _>()` bytes of any valid value of `T`.
412
/// - The first `size_of::<*mut _>()` bytes of any valid value of `T` are
413
///   always fully initialized (no padding or uninitialized bytes in the first
414
///   pointer-sized word), since reading uninitialized bytes is undefined
415
///   behaviour.
416
pub unsafe trait BlockAllocatorSafe {}
417

            
418
/// Sentinel value used to identify entries that are not on the freelist.
419
const NONEXISTING_VALUE: usize = usize::MAX;
420

            
421
/// The [BlockAllocator] is thread-safe.
422
unsafe impl<T: Send, const N: usize> Send for BlockAllocator<T, N> {}
423
unsafe impl<T: Send, const N: usize> Sync for BlockAllocator<T, N> {}
424

            
425
/// `AllocBlock` implements the [`Allocator`] trait using the underlying [`BlockAllocator`].
426
pub struct AllocBlock<T: Send, const N: usize> {
427
    block_allocator: BlockAllocator<T, N>,
428
}
429

            
430
impl<T: Send, const N: usize> Default for AllocBlock<T, N> {
431
    fn default() -> Self {
432
        Self::new()
433
    }
434
}
435

            
436
impl<T: Send, const N: usize> AllocBlock<T, N> {
437
    /// Creates a new `AllocBlock`.
438
17640
    pub fn new() -> Self {
439
17640
        Self {
440
17640
            block_allocator: BlockAllocator::new(),
441
17640
        }
442
17640
    }
443

            
444
    /// Removes free blocks from the underlying block allocator, see [`BlockAllocator::remove_free_blocks`].
445
5582300
    pub fn remove_free_blocks(&mut self) -> usize
446
5582300
    where
447
5582300
        T: BlockAllocatorSafe,
448
    {
449
5582300
        self.block_allocator.remove_free_blocks().0
450
5582300
    }
451
}
452

            
453
unsafe impl<T: Send, const N: usize> Allocator for AllocBlock<T, N> {
454
20126689
    fn allocate(&self, layout: Layout) -> Result<NonNull<[u8]>, AllocError> {
455
        // The blocks only fit objects with exactly T's layout; returning one for a larger
456
        // request would hand the caller a too-small allocation.
457
20126689
        if layout != Layout::new::<T>() {
458
            return Err(AllocError);
459
20126689
        }
460

            
461
20126689
        let ptr = self.block_allocator.allocate_object()?;
462

            
463
        // Convert NonNull<T> to NonNull<[u8]> with the correct size
464
20126689
        let byte_ptr = ptr.cast::<u8>();
465
20126689
        let slice_ptr = NonNull::slice_from_raw_parts(byte_ptr, std::mem::size_of::<T>());
466

            
467
20126689
        Ok(slice_ptr)
468
20126689
    }
469

            
470
19755820
    unsafe fn deallocate(&self, ptr: NonNull<u8>, layout: Layout) {
471
19755820
        debug_assert_eq!(
472
            layout,
473
19755820
            Layout::new::<T>(),
474
            "The requested layout should match the type T"
475
        );
476
19755820
        self.block_allocator.deallocate_object(ptr.cast::<T>());
477
19755820
    }
478
}
479

            
480
union Entry<T> {
481
    /// Stores the actual element.
482
    data: ManuallyDrop<T>,
483

            
484
    /// If the element is free, this points to the next entry in the freelist, or null if this is the last entry.
485
    next: ManuallyDrop<*mut Entry<T>>,
486
}
487

            
488
// Safety: `Entry<T>` stores a single intrusive next-pointer in `next` used only
489
// while the slot is on the freelist.
490
unsafe impl<T> FreeListEntry for Entry<T> {
491
1459708685
    unsafe fn get_next(ptr: *mut Self) -> *mut Self {
492
        // Safety: caller ensures `ptr` is a valid freelist node.
493
1459708685
        unsafe { *(*ptr).next }
494
1459708685
    }
495

            
496
19805982
    unsafe fn set_next(ptr: *mut Self, next: *mut Self) {
497
        // Safety: caller ensures `ptr` is a valid freelist node.
498
19805982
        unsafe {
499
19805982
            *(*ptr).next = next;
500
19805982
        }
501
19805982
    }
502
}
503

            
504
/// An iterator over the blocks in the block allocator.
505
struct BlockIter<'a, T, const N: usize> {
506
    current: Option<NonNull<Block<T, N>>>,
507
    _marker: PhantomData<&'a BlockList<T, N>>,
508
}
509

            
510
impl<'a, T, const N: usize> Iterator for BlockIter<'a, T, N> {
511
    type Item = NonNull<Block<T, N>>;
512

            
513
16862918
    fn next(&mut self) -> Option<Self::Item> {
514
16862918
        let current_ptr = self.current?;
515

            
516
        // Move to the next block for the next iteration.
517
5698118
        self.current = unsafe { (*current_ptr.as_ptr()).next };
518

            
519
5698118
        Some(current_ptr)
520
16862918
    }
521
}
522

            
523
/// We maintain a list of blocks that store N elements each.
524
struct Block<T, const N: usize> {
525
    /// Wrapped in UnsafeCell to avoid Stacked Borrows invalidation of
526
    /// outstanding pointers during subsequent allocations from the same block.
527
    data: UnsafeCell<[Entry<T>; N]>,
528

            
529
    /// Pointer to the next block (raw to avoid Box noalias).
530
    next: Option<NonNull<Block<T, N>>>,
531
}
532

            
533
impl<T, const N: usize> Block<T, N> {
534
16623
    fn new() -> Self {
535
        Self {
536
16623
            data: UnsafeCell::new(array::from_fn(|_i| Entry {
537
13834656
                next: ManuallyDrop::new(std::ptr::null_mut()),
538
13834656
            })),
539
16623
            next: None,
540
        }
541
16623
    }
542
}
543

            
544
impl<T, const N: usize> Drop for Block<T, N> {
545
8067
    fn drop(&mut self) {
546
        // Iteratively drop the list to avoid stack overflow on long lists.
547
8067
        let mut current = self.next.take();
548
11179
        while let Some(block_ptr) = current {
549
3112
            let mut block = unsafe { Box::from_raw(block_ptr.as_ptr()) };
550
3112
            current = block.next.take();
551
3112
        }
552
8067
    }
553
}
554

            
555
#[cfg(test)]
556
mod tests {
557
    use std::ptr::NonNull;
558
    use std::sync::Arc;
559

            
560
    use rand::RngExt;
561

            
562
    use merc_utilities::random_test;
563

            
564
    use super::BlockAllocator;
565
    use super::BlockAllocatorSafe;
566

            
567
    // In practice u64 is used only in tests; real clients must audit their types.
568
    unsafe impl BlockAllocatorSafe for usize {}
569

            
570
    #[test]
571
    #[cfg_attr(miri, ignore)]
572
1
    fn test_block_allocator() {
573
100
        random_test(100, |rng| {
574
100
            let mut allocator: BlockAllocator<usize, 32> = BlockAllocator::new();
575

            
576
            // Allocate 1000 elements and keep track of ptr to value mapping.
577
100
            let mut allocated: Vec<(NonNull<usize>, usize)> = Vec::new();
578
100000
            for _ in 0..1000 {
579
100000
                let ptr = allocator.allocate_object().unwrap();
580
100000
                let value: usize = rng.random_range(0..=usize::MAX - 1);
581
100000
                unsafe {
582
100000
                    ptr.as_ptr().write(value);
583
100000
                }
584
100000
                allocated.push((ptr, value));
585
100000
            }
586

            
587
            // Deallocate a random subset and keep track of which entries remain live.
588
100
            let mut remaining = Vec::new();
589
100000
            for (ptr, value) in allocated {
590
100000
                if rng.random_bool(0.5) {
591
49862
                    allocator.deallocate_object(ptr);
592
50138
                } else {
593
50138
                    remaining.push((ptr, value));
594
50138
                }
595
            }
596

            
597
            // All remaining elements must still hold their original values.
598
50138
            for (ptr, expected) in &remaining {
599
                unsafe {
600
50138
                    assert_eq!(*ptr.as_ref(), *expected);
601
                }
602
            }
603

            
604
100
            let (removed, free_size) = allocator.remove_free_blocks();
605
100
            println!("{removed} removed, {free_size} free");
606

            
607
50000
            for _ in 0..500 {
608
50000
                let ptr = allocator.allocate_object().unwrap();
609
50000
                let value: usize = rng.random_range(0..=usize::MAX - 1);
610
50000
                unsafe {
611
50000
                    ptr.as_ptr().write(value);
612
50000
                }
613
50000
                remaining.push((ptr, value));
614
50000
            }
615

            
616
            // All remaining elements must have the correct values.
617
100138
            for (ptr, expected) in &remaining {
618
                unsafe {
619
100138
                    assert_eq!(*ptr.as_ref(), *expected);
620
                }
621
            }
622
100
        })
623
1
    }
624

            
625
    #[test]
626
    #[cfg_attr(miri, ignore)]
627
1
    fn test_block_allocator_parallel_freelist() {
628
1
        let block_allocator = Arc::new(BlockAllocator::<u32, 32>::new());
629

            
630
1
        let threads: Vec<_> = (0..=2)
631
3
            .map(|_| {
632
3
                let block_allocator = block_allocator.clone();
633

            
634
3
                std::thread::spawn(move || {
635
                    // Not sure if this could actually trigger the ABA problem,
636
                    // but this is only used to detect data races.
637
3
                    let mut ptrs = Vec::new();
638
300
                    for _ in 0..100 {
639
300
                        let ptr = block_allocator.allocate_object().unwrap();
640
300
                        unsafe {
641
300
                            ptr.as_ptr().write(42);
642
300
                        }
643
300
                        ptrs.push(ptr);
644
300
                    }
645

            
646
300
                    for ptr in ptrs {
647
                        unsafe {
648
300
                            assert_eq!(*ptr.as_ref(), 42);
649
                        }
650
300
                        block_allocator.deallocate_object(ptr);
651
                    }
652
3
                })
653
3
            })
654
1
            .collect();
655

            
656
3
        for thread in threads {
657
3
            thread.join().unwrap();
658
3
        }
659
1
    }
660
}