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
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
367
368
369
370
371
372
373
374
375
376
377
378
379
380
381
382
383
384
385
386
387
388
389
390
391
392
393
394
395
396
397
398
399
400
401
402
403
404
405
406
407
408
409
410
411
412
413
414
415
416
417
418
419
420
421
422
//! Custom memory allocators and utilities for using them.
//!
//! # Examples
//! ```rust
//! #![feature(placement_in_syntax)]
//!
//! use std::io;
//! use allocators::{Allocator, Scoped, BlockOwner, Proxy};
//!
//! #[derive(Debug)]
//! struct Bomb(u8);
//!
//! impl Drop for Bomb {
//!     fn drop(&mut self) {
//!         println!("Boom! {}", self.0);
//!     }
//! }
//! // new scoped allocator with 4 kilobytes of memory.
//! let alloc = Scoped::new(4 * 1024).unwrap();
//!
//! alloc.scope(|inner| {
//!     let mut bombs = Vec::new();
//!     // allocate makes the value on the stack first.
//!     for i in 0..100 { bombs.push(inner.allocate(Bomb(i)).unwrap())}
//!     // there's also in-place allocation!
//!     let bomb_101 = in inner.make_place().unwrap() { Bomb(101) };
//!     // watch the bombs go off!
//! });
//!
//!
//! // You can make allocators backed by other allocators.
//! {
//!     let secondary_alloc = Scoped::new_from(&alloc, 1024).unwrap();
//!     let mut val = secondary_alloc.allocate(0i32).unwrap();
//!     *val = 1;
//! }
//!
//! // Let's wrap our allocator in a proxy to log what it's doing.
//! let proxied = Proxy::new(alloc, io::stdout());
//! let logged_allocation = proxied.allocate([0u8; 32]).unwrap();
//! ```

#![feature(
    alloc,
    coerce_unsized,
    core_intrinsics,
    heap_api,
    placement_new_protocol,
    placement_in_syntax,
    raw,
    unsize,
)]

use std::any::Any;
use std::borrow::{Borrow, BorrowMut};
use std::error::Error;
use std::fmt;
use std::marker::{PhantomData, Unsize};
use std::mem;
use std::ops::Place as StdPlace;
use std::ops::{CoerceUnsized, Deref, DerefMut, InPlace, Placer};

use alloc::heap;

extern crate alloc;

pub mod composable;
pub mod scoped;

pub use composable::*;
pub use scoped::Scoped;

/// A custom memory allocator.
pub unsafe trait Allocator {
    /// Attempts to allocate the value supplied to it.
    ///
    /// # Examples
    /// ```rust
    /// use allocators::{Allocator, Allocated};
    /// fn alloc_array<A: Allocator>(allocator: &A) -> Allocated<[u8; 1000], A> {
    ///     allocator.allocate([0; 1000]).ok().unwrap()
    /// }
    /// ```
    #[inline(always)]
    fn allocate<T>(&self, val: T) -> Result<Allocated<T, Self>, (AllocatorError, T)>
        where Self: Sized
    {
        match self.make_place() {
            Ok(place) => {
                Ok(in place { val })
            }
            Err(err) => {
                Err((err, val))
            }
        }
    }

    /// Attempts to create a place to allocate into.
    /// For the general purpose, calling `allocate` on the allocator is enough.
    /// However, when you know the value you are allocating is too large
    /// to be constructed on the stack, you should use in-place allocation.
    /// 
    /// # Examples
    /// ```rust
    /// #![feature(placement_in_syntax)]
    /// use allocators::{Allocator, Allocated};
    /// fn alloc_array<A: Allocator>(allocator: &A) -> Allocated<[u8; 1000], A> {
    ///     // if 1000 bytes were enough to smash the stack, this would still work.
    ///     in allocator.make_place().unwrap() { [0; 1000] }
    /// }
    /// ```
    fn make_place<T>(&self) -> Result<Place<T, Self>, AllocatorError> 
        where Self: Sized
    {
        let (size, align) = (mem::size_of::<T>(), mem::align_of::<T>());
        match unsafe { self.allocate_raw(size, align) } {
            Ok(blk) => {
                Ok(Place {
                    allocator: self,
                    block: blk,
                    _marker: PhantomData
                })
            }
            Err(e) => Err(e),
        }
    }

    /// Attempt to allocate a block of memory.
    ///
    /// Returns either a pointer to the block of memory allocated
    /// or an Error. If `size` is equal to 0, the pointer returned must
    /// be equal to `heap::EMPTY`
    ///
    /// # Safety
    /// Never use the pointer outside of the lifetime of the allocator.
    /// It must be deallocated with the same allocator as it was allocated with.
    /// It is undefined behavior to provide a non power-of-two align.
    unsafe fn allocate_raw(&self, size: usize, align: usize) -> Result<Block, AllocatorError>;

    /// Deallocate the memory referred to by this pointer.
    ///
    /// # Safety
    /// This block must have been allocated by this allocator.
    unsafe fn deallocate_raw(&self, blk: Block);
}

/// An allocator that knows which blocks have been issued by it.
pub trait BlockOwner: Allocator {
    /// Whether this allocator owns this allocated value. 
    fn owns<'a, T, A: Allocator>(&self, val: &Allocated<'a, T, A>) -> bool {
        self.owns_block(&val.block)
    }

    /// Whether this allocator owns the block passed to it.
    fn owns_block(&self, blk: &Block) -> bool;

    /// Joins this allocator with a fallback allocator.
    // TODO: Maybe not the right place for this?
    // Right now I've been more focused on shaking out the
    // specifics of allocation than crafting a fluent API.
    fn with_fallback<O: BlockOwner>(self, other: O) -> Fallback<Self, O>
    where Self: Sized {
        Fallback::new(self, other)
    }
}

/// A block of memory created by an allocator.
// TODO: should blocks be tied to lifetimes? seems good for safety!
pub struct Block<'a> {
    ptr: *mut u8,
    size: usize,
    align: usize,
    _marker: PhantomData<&'a [u8]>
}

impl<'a> Block<'a> {
    pub fn new(ptr: *mut u8, size: usize, align: usize) -> Self {
        Block {
            ptr: ptr,
            size: size,
            align: align,
            _marker: PhantomData,
        }
    }

    pub fn ptr(&self) -> *mut u8 { self.ptr }
    pub fn size(&self) -> usize { self.size }
    pub fn align(&self) -> usize { self.align }
}

impl<'a> Clone for Block<'a> {
    fn clone(&self) -> Block<'a> {
        Block {
            ptr: self.ptr(),
            size: self.size(),
            align: self.align(),
            _marker: self._marker
        }
    }
}

/// Errors that can occur while creating an allocator
/// or allocating from it.
#[derive(Debug, Eq, PartialEq)]
pub enum AllocatorError {
    /// The allocator failed to allocate the amount of memory requested of it.
    OutOfMemory,
    /// An allocator-specific error message.
    AllocatorSpecific(String),
}

impl fmt::Display for AllocatorError {
    fn fmt(&self, formatter: &mut fmt::Formatter) -> fmt::Result {
        formatter.write_str(self.description())
    }
}

impl Error for AllocatorError {
    fn description(&self) -> &str {
        use AllocatorError::*;

        match *self {
            OutOfMemory => {
                "Allocator out of memory."
            }
            AllocatorSpecific(ref reason) => {
                reason
            }
        }
    }
}

/// Allocator stub that just forwards to heap allocation.
/// It is recommended to use the `HEAP` constant instead
/// of creating a new instance of this, to benefit from
/// the static lifetime that it provides.
#[derive(Debug)]
pub struct HeapAllocator;

// A constant for allocators to use the heap as a root.
// Values allocated with this are effectively `Box`es.
pub const HEAP: &'static HeapAllocator = &HeapAllocator;

unsafe impl Allocator for HeapAllocator {
    unsafe fn allocate_raw(&self, size: usize, align: usize) -> Result<Block, AllocatorError> {
        let ptr = if size != 0 {
            heap::allocate(size, align)
        } else {
            heap::EMPTY as *mut u8
        };

        if ptr.is_null() {
            Err(AllocatorError::OutOfMemory)
        } else {
            Ok(Block::new(ptr, size, align))
        }
    }

    unsafe fn deallocate_raw(&self, blk: Block) {
        heap::deallocate(blk.ptr(), blk.size(), blk.align())
    }
}

/// An item allocated by a custom allocator.
pub struct Allocated<'a, T: 'a + ?Sized, A: 'a + Allocator> {
    item: *mut T,
    allocator: &'a A,
    block: Block<'a>,
}

impl<'a, T: ?Sized, A: Allocator> Deref for Allocated<'a, T, A> {
    type Target = T;

    fn deref(&self) -> &T {
        unsafe { mem::transmute(self.item) }
    }
}

impl<'a, T: ?Sized, A: Allocator> DerefMut for Allocated<'a, T, A> {
    fn deref_mut(&mut self) -> &mut T {
        unsafe { mem::transmute(self.item) }
    }
}

// Allocated can store trait objects!
impl<'a, T: ?Sized + Unsize<U>, U: ?Sized, A: Allocator> CoerceUnsized<Allocated<'a, U, A>> for Allocated<'a, T, A> {}

impl<'a, A: Allocator> Allocated<'a, Any, A> {
    /// Attempts to downcast this `Allocated` to a concrete type.
    pub fn downcast<T: Any>(self) -> Result<Allocated<'a, T, A>, Allocated<'a, Any, A>> {
        use std::raw::TraitObject;
        if self.is::<T>() {
            let obj: TraitObject = unsafe { mem::transmute(self.item as *mut Any) };
            let new_allocated = Allocated {
                item: obj.data as *mut T,
                allocator: self.allocator,
                block: self.block.clone(),
            };
            mem::forget(self);
            Ok(new_allocated)
        } else {
            Err(self)
        }
    }
}

impl<'a, T: ?Sized, A: Allocator> Borrow<T> for Allocated<'a, T, A> {
    fn borrow(&self) -> &T {
        &**self
    }
}

impl<'a, T: ?Sized, A: Allocator> BorrowMut<T> for Allocated<'a, T, A> {
    fn borrow_mut(&mut self) -> &mut T {
        &mut **self
    }
}

impl<'a, T: ?Sized, A: Allocator> Drop for Allocated<'a, T, A> {
    #[inline]
    fn drop(&mut self) {
        use std::intrinsics::drop_in_place;
        unsafe {
            drop_in_place(self.item);
;
            self.allocator.deallocate_raw(self.block.clone());
        }

    }
}

/// A place for allocating into.
/// This is only used for in-place allocation,
/// e.g. `let val = in (alloc.make_place().unwrap())`
pub struct Place<'a, T: 'a, A: 'a + Allocator> {
    allocator: &'a A,
    block: Block<'a>,
    _marker: PhantomData<T>,
}

impl<'a, T: 'a, A: 'a + Allocator> Placer<T> for Place<'a, T, A> {
    type Place = Self;
    fn make_place(self) -> Self { self }
}

impl<'a, T: 'a, A: 'a + Allocator> InPlace<T> for Place<'a, T, A> {
    type Owner = Allocated<'a, T, A>;
    unsafe fn finalize(self) -> Self::Owner {
        let allocated = Allocated {
            item: self.block.ptr() as *mut T,
            allocator: self.allocator,
            block: self.block.clone()
        };

        mem::forget(self);
        allocated
    }
}

impl<'a, T: 'a, A: 'a + Allocator> StdPlace<T> for Place<'a, T, A> {
    fn pointer(&mut self) -> *mut T {
        self.block.ptr() as *mut T
    }
}

impl<'a, T: 'a, A: 'a + Allocator> Drop for Place<'a, T, A> {
    #[inline]
    fn drop(&mut self) {
        // almost identical to Allocated::Drop, but we don't drop
        // the value in place. This is because if the finalize
        // method was never called, that means the expression
        // to create the value failed, and the memory at the
        // pointer is still uninitialized.
        unsafe {
            self.allocator.deallocate_raw(self.block.clone());
        }

    }
}

// aligns a pointer forward to the next value aligned with `align`.
#[inline(always)]
fn align_forward(ptr: *mut u8, align: usize) -> *mut u8 {
    ((ptr as usize + align - 1) & !(align - 1)) as *mut u8
}

#[cfg(test)]
mod tests {
        
    use std::any::Any;

    use super::*;

    #[test]
    fn heap_lifetime() {
        let my_int;
        {
            my_int = HEAP.allocate(0i32).unwrap(); 
        }

        assert_eq!(*my_int, 0);
    }
    #[test]
    fn heap_in_place() {
        let big = in HEAP.make_place().unwrap() { [0u8; 8_000_000] };
        assert_eq!(big.len(), 8_000_000);
    }

    #[test]
    fn unsizing() {
        #[derive(Debug)]
        struct Bomb;
        impl Drop for Bomb {
            fn drop(&mut self) {
                println!("Boom")
            }
        }

        let my_foo: Allocated<Any, _> = HEAP.allocate(Bomb).unwrap();
        let _: Allocated<Bomb, _> = my_foo.downcast().ok().unwrap();
    }
}