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

use crate::allocators::allocator::Allocator;
use crate::allocators::global::current_allocator_in_use::CurrentAllocatorInUse;
use crate::allocators::global::local_allocator::LocalAllocator;
use std::alloc::{Alloc, GlobalAlloc};

/// A trait that all such allocators implement.
///
/// Create a new instance using the macro `switchable_allocator`.
pub trait GlobalSwitchableAllocator: Sync + GlobalAlloc + Alloc + Allocator {
    /// Type of the coroutine local allocator.
    type CoroutineLocalAllocator: LocalAllocator;

    /// Type of the thread local allocator.
    type ThreadLocalAllocator: LocalAllocator;

    /// Type of the global allocator.
    type GlobalAllocator: Allocator;

    /// Swaps the coroutine local allocator.
    ///
    /// Used before calling a coroutine.
    ///
    /// Used after calling a coroutine.
    fn replace_coroutine_local_allocator(
        &self,
        replacement: Option<Self::CoroutineLocalAllocator>,
    ) -> Option<Self::CoroutineLocalAllocator>;

    /// Initializes the thread local allocator.
    fn initialize_thread_local_allocator(&self, thread_local_allocator: Self::ThreadLocalAllocator);

    /// Drops the thread local allocator.
    ///
    /// Panics in debug if no thread local allocator has been initialized with `initialize_thread_local_allocator()`.
    ///
    /// Could be made hidden by using a destructor with `libc::pthread_key_create()` for an otherwise unused key.
    fn drop_thread_local_allocator(&self);

    /// Save the current allocator in use.
    fn save_current_allocator_in_use(&self) -> CurrentAllocatorInUse;

    /// Restore the current allocator in use.
    fn restore_current_allocator_in_use(&self, restore_to: CurrentAllocatorInUse);

    /// Replace the current allocator in use.
    #[inline(always)]
    fn replace_current_allocator_in_use(
        &self,
        replacement: CurrentAllocatorInUse,
    ) -> CurrentAllocatorInUse {
        let was = self.save_current_allocator_in_use();
        self.restore_current_allocator_in_use(replacement);
        was
    }

    /// Switch the current allocator in use to coroutine local and execute the callback; restore it after calling the callback unless a panic occurs.
    #[inline(always)]
    fn callback_with_coroutine_local_allocator<R>(&self, callback: impl FnOnce() -> R) -> R {
        self.callback_with_different_current_allocator(
            CurrentAllocatorInUse::CoroutineLocal,
            callback,
        )
    }

    /// Switch the current allocator in use to thread local and execute the callback; restore it after calling the callback unless a panic occurs.
    #[inline(always)]
    fn callback_with_thread_local_allocator<R>(&self, callback: impl FnOnce() -> R) -> R {
        self.callback_with_different_current_allocator(CurrentAllocatorInUse::ThreadLocal, callback)
    }

    /// Switch the current allocator in use to global and execute the callback; restore it after calling the callback unless a panic occurs.
    #[inline(always)]
    fn callback_with_global_allocator<R>(&self, callback: impl FnOnce() -> R) -> R {
        self.callback_with_different_current_allocator(CurrentAllocatorInUse::Global, callback)
    }

    /// Switch the current allocator in use and execute the callback; restore it after calling the callback unless a panic occurs.
    #[inline(always)]
    fn callback_with_different_current_allocator<R>(
        &self,
        different: CurrentAllocatorInUse,
        callback: impl FnOnce() -> R,
    ) -> R {
        let restore_to = self.save_current_allocator_in_use();
        self.restore_current_allocator_in_use(different);
        let result = callback();
        self.restore_current_allocator_in_use(restore_to);
        result
    }

    /// Obtain the current coroutine local allocator, if any.
    fn coroutine_local_allocator(&self) -> Option<&Self::CoroutineLocalAllocator>;

    /// Obtain the coroutine local allocator.
    ///
    /// Panics if no coroutine local allocator has been assigned with `replace_coroutine_local_allocator()`.
    #[inline(always)]
    fn coroutine_local_allocator_unchecked(&self) -> &Self::CoroutineLocalAllocator {
        self.coroutine_local_allocator().expect("Assign the coroutine local allocator first using `replace_coroutine_local_allocator()`")
    }

    /// Obtain the thread local allocator.
    ///
    /// None if no thread local allocator has been initialized with `initialize_thread_local_allocator()`.
    fn thread_local_allocator(&self) -> Option<&Self::ThreadLocalAllocator>;

    /// Obtain the thread local allocator.
    ///
    /// Panics if no thread local allocator has been initialized with `initialize_thread_local_allocator()`.
    #[inline(always)]
    fn thread_local_allocator_unchecked(&self) -> &Self::ThreadLocalAllocator {
        self.thread_local_allocator().expect("Initialize the thread local allocator first using `initialize_thread_local_allocator()`")
    }

    /// Obtain the global allocator.
    fn global_allocator(&self) -> &Self::GlobalAllocator;
}