1
  2
  3
  4
  5
  6
  7
  8
  9
 10
 11
 12
 13
 14
 15
 16
 17
 18
 19
 20
 21
 22
 23
 24
 25
 26
 27
 28
 29
 30
 31
 32
 33
 34
 35
 36
 37
 38
 39
 40
 41
 42
 43
 44
 45
 46
 47
 48
 49
 50
 51
 52
 53
 54
 55
 56
 57
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145

use core::alloc::GlobalAlloc;
pub use core::alloc::Layout;
use core::sync::atomic::{fence, Ordering};

/// An action that an allocator can take, either right before, or right after it
/// happens.
#[derive(Clone, Copy, Eq, PartialEq, Debug)]
pub enum AllocAction {
    /// alloc was called
    Alloc,
    /// alloc returned a pointer
    AllocResult { ptr: *mut u8 },
    /// alloc_zeroed was called
    AllocZeroed,
    /// alloc_zeroed returned a pointer
    AllocZeroedResult { ptr: *mut u8 },
    /// dealloc was called on a pointer
    Dealloc { ptr: *mut u8 },
    /// dealloc finished execution
    DeallocResult,
    /// realloc was called on a pointer
    Realloc { ptr: *mut u8, new_size: usize },
    /// realloc returned a pointer
    ReallocResult { ptr: *mut u8, new_size: usize },
}

/// Before or after an allocation call is executed.
#[derive(Clone, Copy, Eq, PartialEq, Debug)]
pub enum AllocRel {
    Before,
    After,
}

impl AllocAction {
    /// Whether the action is before or after the action itself.
    #[inline]
    pub fn relation(&self) -> AllocRel {
        use AllocAction::*;
        match self {
            Alloc | AllocZeroed => AllocRel::Before,
            Dealloc { ptr: _ } => AllocRel::Before,
            Realloc {
                ptr: _,
                new_size: _,
            } => AllocRel::Before,
            _ => AllocRel::After,
        }
    }
}

/// An allocator that watches the calls to its API, sends them to a struct, and
/// then makes the calls.
///
/// To use this struct, use the `#[global_allocator]` compiler directive and
/// construct the allocator with your own custom monitor, or one of the builtins.
/// Note that the new method of `interloc::StatsMonitor` is a `const fn`.
pub struct InterAlloc<'a, T, F>
where
    T: GlobalAlloc,
    F: AllocMonitor,
{
    // Taken directly from https://github.com/neoeinstein/stats_alloc, or stats_alloc
    // on crates.io - all credit to the original writer of this struct, the user
    // neoeinstein on GitHub.
    /// Inner allocator
    pub inner: T,
    /// Monitor on the calls to the allocator
    pub monitor: &'a F,
}

impl<'a, T, F> InterAlloc<'a, T, F>
where
    T: GlobalAlloc,
    F: AllocMonitor,
{
    /// A new instance of the allocator.
    #[inline]
    pub fn new(base: T, monitor: &'a F) -> Self {
        Self {
            inner: base,
            monitor: monitor,
        }
    }

    /// Call the monitor function.
    #[inline]
    fn monitor_(&self, layout: Layout, act: AllocAction) {
        self.monitor.monitor(layout, act);
    }
}

unsafe impl<'a, T, F> GlobalAlloc for InterAlloc<'a, T, F>
where
    T: GlobalAlloc,
    F: AllocMonitor,
{
    #[inline]
    unsafe fn alloc(&self, layout: Layout) -> *mut u8 {
        self.monitor_(layout, AllocAction::Alloc);
        fence(Ordering::SeqCst);
        let ptr = self.inner.alloc(layout);
        fence(Ordering::SeqCst);
        self.monitor_(layout, AllocAction::AllocResult { ptr });
        ptr
    }

    #[inline]
    unsafe fn dealloc(&self, ptr: *mut u8, layout: Layout) {
        self.monitor_(layout, AllocAction::Dealloc { ptr });
        fence(Ordering::SeqCst);
        self.inner.dealloc(ptr, layout);
        fence(Ordering::SeqCst);
        self.monitor_(layout, AllocAction::DeallocResult);
    }

    #[inline]
    unsafe fn alloc_zeroed(&self, layout: Layout) -> *mut u8 {
        self.monitor_(layout, AllocAction::AllocZeroed);
        fence(Ordering::SeqCst);
        let ptr = self.inner.alloc_zeroed(layout);
        fence(Ordering::SeqCst);
        self.monitor_(layout, AllocAction::AllocZeroedResult { ptr });
        ptr
    }

    #[inline]
    unsafe fn realloc(&self, ptr: *mut u8, layout: Layout, new_size: usize) -> *mut u8 {
        self.monitor_(layout, AllocAction::Realloc { ptr, new_size });
        fence(Ordering::SeqCst);
        let ptr = self.inner.realloc(ptr, layout, new_size);
        fence(Ordering::SeqCst);
        self.monitor_(layout, AllocAction::Realloc { ptr, new_size });
        ptr
    }
}

/// When attached to an `InterAlloc` instance, this struct's `monitor` method
/// is called before and after calls to the inner allocator. The ordering of
/// these method calls is enforced by `std::sync::atomic::fence`
/// with `std::sync::atomic::Ordering::SeqCst`.
pub trait AllocMonitor {
    /// The api to the monitor. This method is called right before and right after
    /// allocations happen.
    fn monitor(&self, layout: Layout, act: AllocAction);
}