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

use crate::{grow, Owns};
use core::{
    alloc::{AllocErr, AllocInit, AllocRef, Layout, MemoryBlock, ReallocPlacement},
    ptr::NonNull,
};

/// An allocator equivalent of an "or" operator in algebra.
///
/// An allocation request is first attempted with the `Primary` allocator. If that fails, the
/// request is forwarded to the `Fallback` allocator. All other requests are dispatched
/// appropriately to one of the two allocators.
///
/// A `FallbackAlloc` is useful for fast, special-purpose allocators backed up by general-purpose
/// allocators like [`Global`] or [`System`].
///
/// [`Global`]: https://doc.rust-lang.org/alloc/alloc/struct.Global.html
/// [`System`]: https://doc.rust-lang.org/std/alloc/struct.System.html
///
/// # Example
///
/// ```rust
/// #![feature(allocator_api)]
///
/// use alloc_compose::{FallbackAlloc, Owns, Region};
/// use std::alloc::{AllocInit, AllocRef, Layout, System};
///
/// let mut data = [0; 32];
/// let mut alloc = FallbackAlloc {
///     primary: Region::new(&mut data),
///     fallback: System,
/// };
///
/// let small_memory = alloc.alloc(Layout::new::<u32>(), AllocInit::Uninitialized)?;
/// let big_memory = alloc.alloc(Layout::new::<[u32; 64]>(), AllocInit::Uninitialized)?;
///
/// assert!(alloc.primary.owns(small_memory));
/// assert!(!alloc.primary.owns(big_memory));
///
/// unsafe {
///     // `big_memory` was allocated from `System`, we can dealloc it directly
///     System.dealloc(big_memory.ptr, Layout::new::<[u32; 64]>());
///     alloc.dealloc(small_memory.ptr, Layout::new::<u32>());
/// };
/// # Ok::<(), core::alloc::AllocErr>(())
/// ```
#[derive(Debug, Copy, Clone)]
pub struct FallbackAlloc<Primary, Fallback> {
    /// The primary allocator
    pub primary: Primary,
    /// The fallback allocator
    pub fallback: Fallback,
}

unsafe impl<Primary, Fallback> AllocRef for FallbackAlloc<Primary, Fallback>
where
    Primary: AllocRef + Owns,
    Fallback: AllocRef,
{
    fn alloc(&mut self, layout: Layout, init: AllocInit) -> Result<MemoryBlock, AllocErr> {
        match self.primary.alloc(layout, init) {
            primary @ Ok(_) => primary,
            Err(_) => self.fallback.alloc(layout, init),
        }
    }

    unsafe fn dealloc(&mut self, ptr: NonNull<u8>, layout: Layout) {
        if self.primary.owns(MemoryBlock {
            ptr,
            size: layout.size(),
        }) {
            self.primary.dealloc(ptr, layout)
        } else {
            self.fallback.dealloc(ptr, layout)
        }
    }

    unsafe fn grow(
        &mut self,
        ptr: NonNull<u8>,
        layout: Layout,
        new_size: usize,
        placement: ReallocPlacement,
        init: AllocInit,
    ) -> Result<MemoryBlock, AllocErr> {
        if self.primary.owns(MemoryBlock {
            ptr,
            size: layout.size(),
        }) {
            if let Ok(memory) = self.primary.grow(ptr, layout, new_size, placement, init) {
                Ok(memory)
            } else {
                grow(
                    &mut self.primary,
                    &mut self.fallback,
                    ptr,
                    layout,
                    new_size,
                    placement,
                    init,
                )
            }
        } else {
            self.fallback.grow(ptr, layout, new_size, placement, init)
        }
    }

    unsafe fn shrink(
        &mut self,
        ptr: NonNull<u8>,
        layout: Layout,
        new_size: usize,
        placement: ReallocPlacement,
    ) -> Result<MemoryBlock, AllocErr> {
        if self.primary.owns(MemoryBlock {
            ptr,
            size: layout.size(),
        }) {
            self.primary.shrink(ptr, layout, new_size, placement)
        } else {
            self.fallback.shrink(ptr, layout, new_size, placement)
        }
    }
}

impl<Primary, Fallback> Owns for FallbackAlloc<Primary, Fallback>
where
    Primary: Owns,
    Fallback: Owns,
{
    fn owns(&self, memory: MemoryBlock) -> bool {
        self.primary.owns(memory) || self.fallback.owns(memory)
    }
}