alloc_test/alloc/
allocator.rs

1use std::alloc::{GlobalAlloc, Layout};
2
3#[derive(Debug, Default)]
4pub struct TracingAllocator<H: 'static, A>(A, H)
5where
6    A: GlobalAlloc;
7
8impl<H, A> TracingAllocator<H, A>
9where
10    A: GlobalAlloc,
11{
12    pub const fn new(hooks: H, allocator: A) -> Self {
13        TracingAllocator(allocator, hooks)
14    }
15}
16
17// pub const fn default_tracing_allocator() -> TracingAllocator<MemoryTracingHooks, System> {
18//     TracingAllocator(System, MemoryTracingHooks)
19// }
20
21unsafe impl<H, A> GlobalAlloc for TracingAllocator<H, A>
22where
23    A: GlobalAlloc,
24    H: AllocHooks,
25{
26    unsafe fn alloc(&self, layout: Layout) -> *mut u8 {
27        let size = layout.size();
28        let align = layout.align();
29        let pointer = self.0.alloc(layout);
30        self.1.on_alloc(pointer, size, align);
31        pointer
32    }
33
34    unsafe fn dealloc(&self, pointer: *mut u8, layout: Layout) {
35        let size = layout.size();
36        let align = layout.align();
37        self.0.dealloc(pointer, layout);
38        self.1.on_dealloc(pointer, size, align);
39    }
40
41    unsafe fn alloc_zeroed(&self, layout: Layout) -> *mut u8 {
42        let size = layout.size();
43        let align = layout.align();
44        let pointer = self.0.alloc_zeroed(layout);
45        self.1.on_alloc_zeroed(pointer, size, align);
46        pointer
47    }
48
49    unsafe fn realloc(&self, old_pointer: *mut u8, layout: Layout, new_size: usize) -> *mut u8 {
50        let old_size = layout.size();
51        let align = layout.align();
52        let new_pointer = self.0.realloc(old_pointer, layout, new_size);
53        self.1
54            .on_realloc(old_pointer, new_pointer, old_size, new_size, align);
55        new_pointer
56    }
57}
58
59/// Trait for implementing allocation hooks in a tracing allocator.
60///
61/// # Safety
62///
63/// Implementors must ensure that hook methods do not allocate memory
64/// (to avoid infinite recursion) and are safe to call from any thread
65/// at any time during allocation operations.
66pub unsafe trait AllocHooks {
67    fn on_alloc(&self, pointer: *mut u8, size: usize, align: usize);
68    fn on_dealloc(&self, pointer: *mut u8, size: usize, align: usize);
69    fn on_alloc_zeroed(&self, pointer: *mut u8, size: usize, align: usize);
70    fn on_realloc(
71        &self,
72        old_pointer: *mut u8,
73        new_pointer: *mut u8,
74        old_size: usize,
75        new_size: usize,
76        align: usize,
77    );
78}