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

use super::{Allocation, DeviceAllocator, MemUnit};

use anyhow::Result;
use ash::vk;

pub struct SizeSelector<
    SmallAllocator: DeviceAllocator,
    LargeAllocator: DeviceAllocator,
> {
    small_allocator: SmallAllocator,
    large_allocator: LargeAllocator,
    size: u64,
}

impl<SmallAllocator: DeviceAllocator, LargeAllocator: DeviceAllocator>
    SizeSelector<SmallAllocator, LargeAllocator>
{
    /// Create a new allocator which defers to the small allocator for
    /// allocations below size bytes, and otherwise uses the large allocator.
    pub fn new(
        small_allocator: SmallAllocator,
        size: MemUnit,
        large_allocator: LargeAllocator,
    ) -> Self {
        Self {
            small_allocator,
            large_allocator,
            size: size.to_bytes(),
        }
    }
}

impl<SmallAllocator: DeviceAllocator, LargeAllocator: DeviceAllocator>
    DeviceAllocator for SizeSelector<SmallAllocator, LargeAllocator>
{
    unsafe fn allocate(
        &mut self,
        allocate_info: vk::MemoryAllocateInfo,
    ) -> Result<Allocation> {
        if allocate_info.allocation_size < self.size {
            self.small_allocator.allocate(allocate_info)
        } else {
            self.large_allocator.allocate(allocate_info)
        }
    }

    unsafe fn free(&mut self, allocation: &Allocation) -> Result<()> {
        if allocation.byte_size < self.size {
            self.small_allocator.free(allocation)
        } else {
            self.large_allocator.free(allocation)
        }
    }
}