From 0588a06cfe2f771ce68d51afb28d666524637edf Mon Sep 17 00:00:00 2001
From: Nicolas Silva <nical@fastmail.com>
Date: Thu, 18 Jul 2024 18:43:02 +0200
Subject: [PATCH] Avoid allocating strings when formatting allocation sizes

---
 src/allocator/mod.rs                 | 30 +++++++++++++++-------------
 src/visualizer/allocation_reports.rs |  6 +++---
 2 files changed, 19 insertions(+), 17 deletions(-)

diff --git a/src/allocator/mod.rs b/src/allocator/mod.rs
index 5d40bbf..e7435ea 100644
--- a/src/allocator/mod.rs
+++ b/src/allocator/mod.rs
@@ -72,7 +72,7 @@ impl fmt::Debug for AllocationReport {
         } else {
             "--"
         };
-        write!(f, "{name:?}: {}", fmt_bytes(self.size))
+        write!(f, "{name:?}: {}", FmtBytes(self.size))
     }
 }
 
@@ -89,8 +89,8 @@ impl fmt::Debug for AllocatorReport {
                 "summary",
                 &std::format_args!(
                     "{} / {}",
-                    fmt_bytes(self.total_allocated_bytes),
-                    fmt_bytes(self.total_reserved_bytes)
+                    FmtBytes(self.total_allocated_bytes),
+                    FmtBytes(self.total_reserved_bytes)
                 ),
             )
             .field("blocks", &self.blocks.len())
@@ -145,18 +145,20 @@ pub(crate) trait SubAllocator: SubAllocatorBase + fmt::Debug + Sync + Send {
     }
 }
 
-pub(crate) fn fmt_bytes(mut amount: u64) -> String {
-    const SUFFIX: [&str; 5] = ["B", "KB", "MB", "GB", "TB"];
+pub struct FmtBytes(pub u64);
 
-    let mut idx = 0;
-    let mut print_amount = amount as f64;
-    loop {
-        if amount < 1024 {
-            return format!("{:.2} {}", print_amount, SUFFIX[idx]);
+impl fmt::Display for FmtBytes {
+    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
+        const SUFFIX: [&str; 5] = ["B", "KB", "MB", "GB", "TB"];
+        let mut idx = 0;
+        let mut amount = self.0 as f64;
+        loop {
+            if amount < 1024.0 || idx >= SUFFIX.len() - 1 {
+                return write!(f, "{:.2} {}", amount, SUFFIX[idx]);
+            }
+
+            amount /= 1024.0;
+            idx += 1;
         }
-
-        print_amount = amount as f64 / 1024.0;
-        amount /= 1024;
-        idx += 1;
     }
 }
diff --git a/src/visualizer/allocation_reports.rs b/src/visualizer/allocation_reports.rs
index c95b51a..9daf810 100644
--- a/src/visualizer/allocation_reports.rs
+++ b/src/visualizer/allocation_reports.rs
@@ -3,7 +3,7 @@ use std::backtrace::BacktraceStatus;
 use egui::{Label, Response, Sense, Ui, WidgetText};
 use egui_extras::{Column, TableBuilder};
 
-use crate::allocator::{fmt_bytes, AllocationReport};
+use crate::allocator::{AllocationReport, FmtBytes};
 
 #[derive(Clone, Copy, Debug, Default, PartialEq, Eq)]
 pub(crate) enum AllocationReportVisualizeSorting {
@@ -39,7 +39,7 @@ pub(crate) fn render_allocation_reports_ui(
         .collect::<Vec<_>>();
     let total_size_under_filter: u64 = allocations.iter().map(|a| a.1.size).sum();
 
-    ui.label(format!("Total: {}", fmt_bytes(total_size_under_filter)));
+    ui.label(format!("Total: {}", FmtBytes(total_size_under_filter)));
 
     let row_height = ui.text_style_height(&egui::TextStyle::Body);
     let table = TableBuilder::new(ui)
@@ -133,7 +133,7 @@ pub(crate) fn render_allocation_reports_ui(
                 }
 
                 row.col(|ui| {
-                    ui.label(fmt_bytes(size));
+                    ui.label(format!("{}", FmtBytes(size)));
                 });
             });
         }