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

use {
    crate::{
        Allocation, AllocationRequirements, AllocatorError,
        ComposableAllocator, MemoryProperties, MemoryTypePoolAllocator,
    },
    std::{cell::RefCell, collections::HashMap, rc::Rc},
};

type SharedAllocator<T> = Rc<RefCell<T>>;

pub struct PoolAllocator<A: ComposableAllocator> {
    typed_pools: HashMap<usize, MemoryTypePoolAllocator<SharedAllocator<A>>>,
}

impl<A: ComposableAllocator> PoolAllocator<A> {
    pub fn new(
        memory_properties: MemoryProperties,
        chunk_size: u64,
        page_size: u64,
        allocator: A,
    ) -> Self {
        let allocator = SharedAllocator::new(RefCell::new(allocator));
        let typed_pools = memory_properties
            .types()
            .iter()
            .enumerate()
            .map(|(memory_type_index, _memory_type)| {
                (
                    memory_type_index,
                    MemoryTypePoolAllocator::new(
                        memory_type_index,
                        chunk_size,
                        page_size,
                        allocator.clone(),
                    ),
                )
            })
            .collect::<HashMap<_, _>>();
        Self { typed_pools }
    }
}

impl<A: ComposableAllocator> ComposableAllocator for PoolAllocator<A> {
    unsafe fn allocate(
        &mut self,
        allocation_requirements: AllocationRequirements,
    ) -> Result<Allocation, AllocatorError> {
        let pool = self
            .typed_pools
            .get_mut(&allocation_requirements.memory_type_index)
            .unwrap();
        pool.allocate(allocation_requirements)
    }

    unsafe fn free(&mut self, allocation: Allocation) {
        let pool = self
            .typed_pools
            .get_mut(&allocation.memory_type_index())
            .unwrap();
        pool.free(allocation)
    }
}