From df1042499ddc36f08ba268c77bc7be7bd7d616ce Mon Sep 17 00:00:00 2001
From: Nicolas Silva <nical@fastmail.com>
Date: Tue, 14 May 2024 17:13:38 +0200
Subject: [PATCH] Add a public API to expose allocation reports

---
 .../dedicated_block_allocator/mod.rs          |  1 +
 src/allocator/free_list_allocator/mod.rs      |  1 +
 src/allocator/mod.rs                          | 23 +++++++++++---
 src/d3d12/mod.rs                              | 28 +++++++++++++++++
 src/visualizer/allocation_reports.rs          |  1 +
 src/vulkan/mod.rs                             | 31 +++++++++++++++++--
 6 files changed, 79 insertions(+), 6 deletions(-)

diff --git a/src/allocator/dedicated_block_allocator/mod.rs b/src/allocator/dedicated_block_allocator/mod.rs
index 79ec8ef..4262291 100644
--- a/src/allocator/dedicated_block_allocator/mod.rs
+++ b/src/allocator/dedicated_block_allocator/mod.rs
@@ -116,6 +116,7 @@ impl SubAllocator for DedicatedBlockAllocator {
                 .name
                 .clone()
                 .unwrap_or_else(|| "<Unnamed Dedicated allocation>".to_owned()),
+            offset: 0,
             size: self.size,
             #[cfg(feature = "visualizer")]
             backtrace: self.backtrace.clone(),
diff --git a/src/allocator/free_list_allocator/mod.rs b/src/allocator/free_list_allocator/mod.rs
index 33437c8..e04bb78 100644
--- a/src/allocator/free_list_allocator/mod.rs
+++ b/src/allocator/free_list_allocator/mod.rs
@@ -398,6 +398,7 @@ impl SubAllocator for FreeListAllocator {
                     .name
                     .clone()
                     .unwrap_or_else(|| "<Unnamed FreeList allocation>".to_owned()),
+                offset: chunk.offset,
                 size: chunk.size,
                 #[cfg(feature = "visualizer")]
                 backtrace: chunk.backtrace.clone(),
diff --git a/src/allocator/mod.rs b/src/allocator/mod.rs
index a84c047..ab6dfa0 100644
--- a/src/allocator/mod.rs
+++ b/src/allocator/mod.rs
@@ -1,4 +1,4 @@
-use std::{backtrace::Backtrace, sync::Arc};
+use std::{backtrace::Backtrace, ops::Range, sync::Arc};
 
 use log::*;
 
@@ -30,13 +30,28 @@ impl AllocationType {
 }
 
 #[derive(Clone)]
-pub(crate) struct AllocationReport {
-    pub(crate) name: String,
-    pub(crate) size: u64,
+pub struct AllocationReport {
+    pub name: String,
+    pub offset: u64,
+    pub size: u64,
     #[cfg(feature = "visualizer")]
     pub(crate) backtrace: Arc<Backtrace>,
 }
 
+#[derive(Clone)]
+pub struct MemoryBlockReport {
+    pub size: u64,
+    pub allocations: Range<usize>,
+}
+
+#[derive(Clone)]
+pub struct AllocatorReport {
+    pub allocations: Vec<AllocationReport>,
+    pub blocks: Vec<MemoryBlockReport>,
+    pub total_allocated_bytes: u64,
+    pub total_reserved_bytes: u64,
+}
+
 #[cfg(feature = "visualizer")]
 pub(crate) trait SubAllocatorBase: crate::visualizer::SubAllocatorVisualizer {}
 #[cfg(not(feature = "visualizer"))]
diff --git a/src/d3d12/mod.rs b/src/d3d12/mod.rs
index 75b0d75..fe68364 100644
--- a/src/d3d12/mod.rs
+++ b/src/d3d12/mod.rs
@@ -4,6 +4,7 @@ use std::{backtrace::Backtrace, fmt, sync::Arc};
 
 use log::{debug, warn, Level};
 
+use crate::allocator::AllocatorReport;
 use windows::Win32::{
     Foundation::E_OUTOFMEMORY,
     Graphics::{Direct3D12::*, Dxgi::Common::DXGI_FORMAT},
@@ -1099,6 +1100,33 @@ impl Allocator {
             Ok(())
         }
     }
+
+    pub fn genetate_report(&self) -> AllocatorReport {
+        let mut allocations = vec![];
+        let mut blocks = vec![];
+        let mut total_reserved_bytes = 0;
+
+        for memory_type in &self.memory_types {
+            for block in memory_type.memory_blocks.iter().flatten() {
+                total_reserved_bytes += block.size;
+                let first_allocation = allocations.len();
+                allocations.extend(block.sub_allocator.report_allocations());
+                blocks.push(MemoryBlockReport {
+                    size: block.size,
+                    allocations: first_allocation..allocations.len(),
+                });
+            }
+        }
+
+        let total_allocated_bytes = allocations.iter().map(|report| report.size).sum();
+
+        AllocatorReport {
+            allocations,
+            blocks,
+            total_allocated_bytes,
+            total_reserved_bytes,
+        }
+    }
 }
 
 impl fmt::Debug for Allocator {
diff --git a/src/visualizer/allocation_reports.rs b/src/visualizer/allocation_reports.rs
index 9d9bc0e..c95b51a 100644
--- a/src/visualizer/allocation_reports.rs
+++ b/src/visualizer/allocation_reports.rs
@@ -115,6 +115,7 @@ pub(crate) fn render_allocation_reports_ui(
                     name,
                     size,
                     backtrace,
+                    ..
                 } = alloc;
 
                 row.col(|ui| {
diff --git a/src/vulkan/mod.rs b/src/vulkan/mod.rs
index abf42f5..5fd6560 100644
--- a/src/vulkan/mod.rs
+++ b/src/vulkan/mod.rs
@@ -12,8 +12,8 @@ use log::{debug, Level};
 
 use super::allocator;
 use crate::{
-    allocator::fmt_bytes, AllocationError, AllocationSizes, AllocatorDebugSettings, MemoryLocation,
-    Result,
+    allocator::{fmt_bytes, AllocatorReport, MemoryBlockReport},
+    AllocationError, AllocationSizes, AllocatorDebugSettings, MemoryLocation, Result,
 };
 
 #[derive(Copy, Clone, Debug, Eq, PartialEq)]
@@ -972,6 +972,33 @@ impl Allocator {
             })
             .map(|memory_type| memory_type.memory_type_index as _)
     }
+
+    pub fn generate_report(&self) -> AllocatorReport {
+        let mut allocations = vec![];
+        let mut blocks = vec![];
+        let mut total_reserved_bytes = 0;
+
+        for memory_type in &self.memory_types {
+            for block in memory_type.memory_blocks.iter().flatten() {
+                total_reserved_bytes += block.size;
+                let first_allocation = allocations.len();
+                allocations.extend(block.sub_allocator.report_allocations());
+                blocks.push(MemoryBlockReport {
+                    size: block.size,
+                    allocations: first_allocation..allocations.len(),
+                });
+            }
+        }
+
+        let total_allocated_bytes = allocations.iter().map(|report| report.size).sum();
+
+        AllocatorReport {
+            allocations,
+            blocks,
+            total_allocated_bytes,
+            total_reserved_bytes,
+        }
+    }
 }
 
 impl Drop for Allocator {