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
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165

//! This module contains some composable building blocks to build allocator chains.

use std::cell::RefCell;

use super::{Allocator, Error, Block, BlockOwner};

/// This allocator always fails.
/// It will panic if you try to deallocate with it.
pub struct NullAllocator;

unsafe impl Allocator for NullAllocator {
    unsafe fn allocate_raw(&self, _size: usize, _align: usize) -> Result<Block, Error> {
        Err(Error::OutOfMemory)
    }

    unsafe fn reallocate_raw<'a>(&'a self, block: Block<'a>, _new_size: usize) -> Result<Block<'a>, (Error, Block<'a>)> {
        Err((Error::OutOfMemory, block))
    }

    unsafe fn deallocate_raw(&self, _block: Block) {
        panic!("Attempted to deallocate using null allocator.")
    }
}

impl BlockOwner for NullAllocator {
    fn owns_block(&self, _block: &Block) -> bool {
        false
    }
}

/// This allocator has a main and a fallback allocator.
/// It will always attempt to allocate first with the main allocator,
/// and second with the fallback.
pub struct Fallback<M: BlockOwner, F: BlockOwner> {
    main: M,
    fallback: F,
}

impl<M: BlockOwner, F: BlockOwner> Fallback<M, F> {
    /// Create a new `Fallback`
    pub fn new(main: M, fallback: F) -> Self {
        Fallback {
            main: main,
            fallback: fallback,
        }
    }
}

unsafe impl<M: BlockOwner, F: BlockOwner> Allocator for Fallback<M, F> {
    unsafe fn allocate_raw(&self, size: usize, align: usize) -> Result<Block, Error> {
        match self.main.allocate_raw(size, align) {
            Ok(block) => Ok(block),
            Err(_) => self.fallback.allocate_raw(size, align),
        }
    }

    unsafe fn reallocate_raw<'a>(&'a self, block: Block<'a>, new_size: usize) -> Result<Block<'a>, (Error, Block<'a>)> {
        if self.main.owns_block(&block) {
            self.main.reallocate_raw(block, new_size)
        } else if self.fallback.owns_block(&block) {
            self.fallback.reallocate_raw(block, new_size)
        } else {
            Err((Error::AllocatorSpecific("Neither fallback nor main owns this block.".into()), block))
        }
    }

    unsafe fn deallocate_raw(&self, block: Block) {
        if self.main.owns_block(&block) {
            self.main.deallocate_raw(block);
        } else if self.fallback.owns_block(&block) {
            self.fallback.deallocate_raw(block);
        }
    }
}

impl<M: BlockOwner, F: BlockOwner> BlockOwner for Fallback<M, F> {
    fn owns_block(&self, block: &Block) -> bool {
        self.main.owns_block(block) || self.fallback.owns_block(block)
    }
}

/// Something that logs an allocator's activity.
/// In practice, this may be an output stream,
/// a data collector, or seomthing else entirely.
pub trait ProxyLogger {
    /// Called after a successful allocation.
    fn allocate_success(&mut self, block: &Block);
    /// Called after a failed allocation.
    fn allocate_fail(&mut self, err: &Error, size: usize, align: usize);

    /// Called when deallocating a block.
    fn deallocate(&mut self, block: &Block);

    /// Called after a successful reallocation.
    fn reallocate_success(&mut self, old_block: &Block, new_block: &Block);
    /// Called after a failed reallocation.
    fn reallocate_fail(&mut self, err: &Error, block: &Block, req_size: usize);
}

/// This wraps an allocator and a logger, logging all allocations
/// and deallocations.
pub struct Proxy<A, L> {
    alloc: A,
    logger: RefCell<L>,
}

impl<A: Allocator, L: ProxyLogger> Proxy<A, L> {
    /// Create a new proxy allocator.
    pub fn new(alloc: A, logger: L) -> Self {
        Proxy {
            alloc: alloc,
            logger: RefCell::new(logger),
        }
    }
}

unsafe impl<A: Allocator, L: ProxyLogger> Allocator for Proxy<A, L> {
    unsafe fn allocate_raw(&self, size: usize, align: usize) -> Result<Block, Error> {
        let mut logger = self.logger.borrow_mut();
        match self.alloc.allocate_raw(size, align) {
            Ok(block) => {
                logger.allocate_success(&block);
                Ok(block)
            }
            Err(err) => {
                logger.allocate_fail(&err, size, align);
                Err(err)
            }
        }
    }

    unsafe fn reallocate_raw<'a>(&'a self, block: Block<'a>, new_size: usize) -> Result<Block<'a>, (Error, Block<'a>)> {
        let mut logger = self.logger.borrow_mut();
        let old_copy = Block::new(block.ptr(), block.size(), block.align());

        match self.alloc.reallocate_raw(block, new_size) {
            Ok(new_block) => {
                logger.reallocate_success(&old_copy, &new_block);
                Ok(new_block)
            }
            Err((err, old)) => {
                logger.reallocate_fail(&err, &old, new_size);
                Err((err, old))
            }
        }
    }

    unsafe fn deallocate_raw(&self, block: Block) {
        let mut logger = self.logger.borrow_mut();
        logger.deallocate(&block);
        self.alloc.deallocate_raw(block);
    }
}

#[cfg(test)]
mod tests {
    use super::super::*;

    #[test]
    #[should_panic]
    fn null_allocate() {
        let alloc = NullAllocator;
        alloc.allocate(1i32).unwrap();
    }
}