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
use std::{backtrace::Backtrace, fmt, ops::Range, sync::Arc};

use log::*;

use crate::result::*;

pub(crate) mod dedicated_block_allocator;
pub(crate) use dedicated_block_allocator::DedicatedBlockAllocator;

pub(crate) mod free_list_allocator;
pub(crate) use free_list_allocator::FreeListAllocator;

#[derive(PartialEq, Copy, Clone, Debug)]
#[repr(u8)]
pub(crate) enum AllocationType {
    Free,
    Linear,
    NonLinear,
}

impl AllocationType {
    #[cfg(feature = "visualizer")]
    pub fn as_str(self) -> &'static str {
        match self {
            Self::Free => "Free",
            Self::Linear => "Linear",
            Self::NonLinear => "Non-Linear",
        }
    }
}

/// Describes an allocation in the [`AllocatorReport`].
#[derive(Clone)]
pub struct AllocationReport {
    /// The name provided to the `allocate()` function.
    pub name: String,
    /// The offset in bytes of the allocation in its memory block.
    pub offset: u64,
    /// The size in bytes of the allocation.
    pub size: u64,
    #[cfg(feature = "visualizer")]
    pub(crate) backtrace: Arc<Backtrace>,
}

/// Describes a memory block in the [`AllocatorReport`].
#[derive(Clone)]
pub struct MemoryBlockReport {
    /// The size in bytes of this memory block.
    pub size: u64,
    /// The range of allocations in [`AllocatorReport::allocations`] that are associated
    /// to this memory block.
    pub allocations: Range<usize>,
}

/// A report that can be generated for informational purposes using `Allocator::generate_report()`.
#[derive(Clone)]
pub struct AllocatorReport {
    /// All live allocations, sub-allocated from memory blocks.
    pub allocations: Vec<AllocationReport>,
    /// All memory blocks.
    pub blocks: Vec<MemoryBlockReport>,
    /// Sum of the memory used by all allocations, in bytes.
    pub total_allocated_bytes: u64,
    /// Sum of the memory reserved by all memory blocks including unallocated regions, in bytes.
    pub total_reserved_bytes: u64,
}

impl fmt::Debug for AllocationReport {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        let name = if !self.name.is_empty() {
            self.name.as_str()
        } else {
            "--"
        };
        write!(f, "{name:?}: {}", fmt_bytes(self.size))
    }
}

impl fmt::Debug for AllocatorReport {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        let mut allocations = self.allocations.clone();
        allocations.sort_by_key(|alloc| std::cmp::Reverse(alloc.size));

        let max_num_allocations_to_print = f.precision().unwrap_or(usize::MAX);
        allocations.truncate(max_num_allocations_to_print);

        f.debug_struct("AllocatorReport")
            .field(
                "summary",
                &std::format_args!(
                    "{} / {}",
                    fmt_bytes(self.total_allocated_bytes),
                    fmt_bytes(self.total_reserved_bytes)
                ),
            )
            .field("blocks", &self.blocks.len())
            .field("allocations", &self.allocations.len())
            .field("largest", &allocations.as_slice())
            .finish()
    }
}

#[cfg(feature = "visualizer")]
pub(crate) trait SubAllocatorBase: crate::visualizer::SubAllocatorVisualizer {}
#[cfg(not(feature = "visualizer"))]
pub(crate) trait SubAllocatorBase {}

pub(crate) trait SubAllocator: SubAllocatorBase + fmt::Debug + Sync + Send {
    fn allocate(
        &mut self,
        size: u64,
        alignment: u64,
        allocation_type: AllocationType,
        granularity: u64,
        name: &str,
        backtrace: Arc<Backtrace>,
    ) -> Result<(u64, std::num::NonZeroU64)>;

    fn free(&mut self, chunk_id: Option<std::num::NonZeroU64>) -> Result<()>;

    fn rename_allocation(
        &mut self,
        chunk_id: Option<std::num::NonZeroU64>,
        name: &str,
    ) -> Result<()>;

    fn report_memory_leaks(
        &self,
        log_level: Level,
        memory_type_index: usize,
        memory_block_index: usize,
    );

    fn report_allocations(&self) -> Vec<AllocationReport>;

    #[must_use]
    fn supports_general_allocations(&self) -> bool;
    #[must_use]
    fn allocated(&self) -> u64;

    /// Helper function: reports if the suballocator is empty (meaning, having no allocations).
    #[must_use]
    fn is_empty(&self) -> bool {
        self.allocated() == 0
    }
}

pub(crate) fn fmt_bytes(mut amount: u64) -> String {
    const SUFFIX: [&str; 5] = ["B", "KB", "MB", "GB", "TB"];

    let mut idx = 0;
    let mut print_amount = amount as f64;
    loop {
        if amount < 1024 {
            return format!("{:.2} {}", print_amount, SUFFIX[idx]);
        }

        print_amount = amount as f64 / 1024.0;
        amount /= 1024;
        idx += 1;
    }
}