Merge: Accurate memory statistics (#717)

Closes #608

Author: ashvardanian

include/usearch/index_dense.hpp

@@ -746,6 +746,37 @@ class index_dense_gt {
             vectors_tape_allocator_.total_allocated();
     }
 
+    /**
+     *  @brief  Aggregated memory statistics for the allocator tapes used by the dense index.
+     */
+    struct memory_stats_t {
+        /// Memory stats for the graph structure allocator.
+        std::size_t graph_allocated;
+        std::size_t graph_wasted;
+        std::size_t graph_reserved;
+        /// Memory stats for the vectors data allocator.
+        std::size_t vectors_allocated;
+        std::size_t vectors_wasted;
+        std::size_t vectors_reserved;
+    };
+
+    /**
+     *  @brief  Returns detailed memory statistics with separate breakdowns for the graph
+     *          and vectors allocator tapes.
+     *  @return A `memory_stats_t` struct with per-tape allocated, wasted, and reserved bytes.
+     */
+    memory_stats_t memory_stats() const {
+        auto const& graph_alloc = typed_->tape_allocator();
+        return {
+            graph_alloc.total_allocated(),
+            graph_alloc.total_wasted(),
+            graph_alloc.total_reserved(),
+            vectors_tape_allocator_.total_allocated(),
+            vectors_tape_allocator_.total_wasted(),
+            vectors_tape_allocator_.total_reserved(),
+        };
+    }
+
     static constexpr std::size_t any_thread() { return std::numeric_limits<std::size_t>::max(); }
     static constexpr distance_t infinite_distance() { return std::numeric_limits<distance_t>::max(); }
 

include/usearch/index_plugins.hpp

@@ -889,6 +889,7 @@ template <std::size_t alignment_ak = 1> class memory_mapping_allocator_gt {
     std::size_t last_usage_ = head_size();
     std::size_t last_capacity_ = min_capacity();
     std::size_t wasted_space_ = 0;
+    std::size_t total_allocated_ = 0;
 
   public:
     using value_type = byte_t;
@@ -899,13 +900,15 @@ template <std::size_t alignment_ak = 1> class memory_mapping_allocator_gt {
     memory_mapping_allocator_gt() = default;
     memory_mapping_allocator_gt(memory_mapping_allocator_gt&& other) noexcept
         : last_arena_(exchange(other.last_arena_, nullptr)), last_usage_(exchange(other.last_usage_, 0)),
-          last_capacity_(exchange(other.last_capacity_, 0)), wasted_space_(exchange(other.wasted_space_, 0)) {}
+          last_capacity_(exchange(other.last_capacity_, 0)), wasted_space_(exchange(other.wasted_space_, 0)),
+          total_allocated_(exchange(other.total_allocated_, 0)) {}
 
     memory_mapping_allocator_gt& operator=(memory_mapping_allocator_gt&& other) noexcept {
         std::swap(last_arena_, other.last_arena_);
         std::swap(last_usage_, other.last_usage_);
         std::swap(last_capacity_, other.last_capacity_);
         std::swap(wasted_space_, other.wasted_space_);
+        std::swap(total_allocated_, other.total_allocated_);
         return *this;
     }
 
@@ -930,6 +933,7 @@ template <std::size_t alignment_ak = 1> class memory_mapping_allocator_gt {
         last_usage_ = head_size();
         last_capacity_ = min_capacity();
         wasted_space_ = 0;
+        total_allocated_ = 0;
     }
 
     /**
@@ -968,6 +972,7 @@ template <std::size_t alignment_ak = 1> class memory_mapping_allocator_gt {
             last_arena_ = new_arena;
             last_capacity_ = new_cap;
             last_usage_ = head_size();
+            total_allocated_ += new_cap;
         }
 
         wasted_space_ += extended_bytes - count_bytes;
@@ -978,17 +983,7 @@ template <std::size_t alignment_ak = 1> class memory_mapping_allocator_gt {
      *  @brief Returns the amount of memory used by the allocator across all arenas.
      *  @return The amount of space in bytes.
      */
-    std::size_t total_allocated() const noexcept {
-        if (!last_arena_)
-            return 0;
-        std::size_t total_used = 0;
-        std::size_t last_capacity = last_capacity_;
-        do {
-            total_used += last_capacity;
-            last_capacity /= capacity_multiplier();
-        } while (last_capacity >= min_capacity());
-        return total_used;
-    }
+    std::size_t total_allocated() const noexcept { return total_allocated_; }
 
     /**
      *  @brief Returns the amount of wasted space due to alignment.

rust/lib.cpp

@@ -180,6 +180,19 @@ void NativeIndex::view(rust::Str path) const {
 
 void NativeIndex::reset() const { index_->reset(); }
 size_t NativeIndex::memory_usage() const { return index_->memory_usage(); }
+
+MemoryStats NativeIndex::memory_stats() const {
+    auto stats = index_->memory_stats();
+    MemoryStats result;
+    result.graph_allocated = stats.graph_allocated;
+    result.graph_wasted = stats.graph_wasted;
+    result.graph_reserved = stats.graph_reserved;
+    result.vectors_allocated = stats.vectors_allocated;
+    result.vectors_wasted = stats.vectors_wasted;
+    result.vectors_reserved = stats.vectors_reserved;
+    return result;
+}
+
 char const* NativeIndex::hardware_acceleration() const { return index_->metric().isa_name(); }
 
 void NativeIndex::save_to_buffer(rust::Slice<uint8_t> buffer) const {

rust/lib.hpp

@@ -4,6 +4,7 @@
 // We don't have to forward declare all of those:
 struct Matches;
 struct IndexOptions;
+struct MemoryStats;
 enum class MetricKind;
 enum class ScalarKind;
 
@@ -83,6 +84,7 @@ class NativeIndex {
     void view(rust::Str path) const;
     void reset() const;
     size_t memory_usage() const;
+    MemoryStats memory_stats() const;
     char const* hardware_acceleration() const;
 
     void save_to_buffer(rust::Slice<uint8_t> buffer) const;

rust/lib.rs

@@ -301,6 +301,24 @@ pub mod ffi {
         distances: Vec<f32>,
     }
 
+    /// Detailed memory statistics with separate breakdowns for the graph
+    /// and vectors allocator tapes.
+    #[derive(Debug, Clone, Copy, PartialEq, Eq)]
+    struct MemoryStats {
+        /// Total memory allocated by the graph structure allocator, in bytes.
+        graph_allocated: usize,
+        /// Memory wasted due to alignment in the graph allocator, in bytes.
+        graph_wasted: usize,
+        /// Reserved but unused memory in the graph allocator, in bytes.
+        graph_reserved: usize,
+        /// Total memory allocated by the vectors data allocator, in bytes.
+        vectors_allocated: usize,
+        /// Memory wasted due to alignment in the vectors allocator, in bytes.
+        vectors_wasted: usize,
+        /// Reserved but unused memory in the vectors allocator, in bytes.
+        vectors_reserved: usize,
+    }
+
     /// The index options used to configure the dense index during creation.
     /// It contains the number of dimensions, the metric kind, the scalar kind,
     /// the connectivity, the expansion values, and the multi-flag.
@@ -423,6 +441,7 @@ pub mod ffi {
         pub fn view(self: &NativeIndex, path: &str) -> Result<()>;
         pub fn reset(self: &NativeIndex) -> Result<()>;
         pub fn memory_usage(self: &NativeIndex) -> usize;
+        pub fn memory_stats(self: &NativeIndex) -> MemoryStats;
         pub fn hardware_acceleration(self: &NativeIndex) -> *const c_char;
 
         pub fn save_to_buffer(self: &NativeIndex, buffer: &mut [u8]) -> Result<()>;
@@ -432,7 +451,7 @@ pub mod ffi {
 }
 
 // Re-export the FFI structs and enums at the crate root for easy access
-pub use ffi::{IndexOptions, MetricKind, ScalarKind};
+pub use ffi::{IndexOptions, MemoryStats, MetricKind, ScalarKind};
 
 /// Represents custom metric functions for calculating distances between vectors in various formats.
 ///
@@ -1376,6 +1395,12 @@ impl Index {
         self.inner.memory_usage()
     }
 
+    /// Returns detailed memory statistics with separate breakdowns for the graph
+    /// and vectors allocator tapes.
+    pub fn memory_stats(self: &Index) -> ffi::MemoryStats {
+        self.inner.memory_stats()
+    }
+
     /// Saves the index to a specified file.
     ///
     /// # Arguments