Add `MemoryUsageVisitor

Author: MichaReiser

components/salsa-macro-rules/src/setup_input_struct.rs

@@ -103,8 +103,8 @@ macro_rules! setup_input_struct {
                 type Durabilities = [$zalsa::Durability; $N];
 
                 $(
-                    fn heap_size(value: &Self::Fields) -> Option<usize> {
-                        Some($heap_size_fn(value))
+                    fn heap_size(value: &Self::Fields, visitor: &mut dyn salsa::MemoryUsageVisitor) -> Option<usize> {
+                        Some($heap_size_fn(value, visitor))
                     }
                 )?
             }

components/salsa-macro-rules/src/setup_interned_struct.rs

@@ -151,8 +151,8 @@ macro_rules! setup_interned_struct {
                 type Struct<'db> = $Struct< $($db_lt_arg)? >;
 
                 $(
-                    fn heap_size(value: &Self::Fields<'_>) -> Option<usize> {
-                        Some($heap_size_fn(value))
+                    fn heap_size(value: &Self::Fields<'_>, visitor: &mut dyn salsa::MemoryUsageVisitor) -> Option<usize> {
+                        Some($heap_size_fn(value, visitor))
                     }
                 )?
             }

components/salsa-macro-rules/src/setup_tracked_fn.rs

@@ -231,8 +231,8 @@ macro_rules! setup_tracked_fn {
                 $($values_equal)+
 
                 $(
-                    fn heap_size(value: &Self::Output<'_>) -> Option<usize> {
-                        Some($heap_size_fn(value))
+                    fn heap_size(value: &Self::Output<'_>, visitor: &mut dyn salsa::MemoryUsageVisitor) -> Option<usize> {
+                        Some($heap_size_fn(value, visitor))
                     }
                 )?
 

components/salsa-macro-rules/src/setup_tracked_struct.rs

@@ -190,8 +190,8 @@ macro_rules! setup_tracked_struct {
                 }
 
                 $(
-                    fn heap_size(value: &Self::Fields<'_>) -> Option<usize> {
-                        Some($heap_size_fn(value))
+                    fn heap_size(value: &Self::Fields<'_>, visitor: &mut dyn salsa::MemoryUsageVisitor) -> Option<usize> {
+                        Some($heap_size_fn(value, visitor))
                     }
                 )?
             }

src/database.rs

@@ -138,98 +138,66 @@ impl dyn Database {
 }
 
 #[cfg(feature = "salsa_unstable")]
-pub use memory_usage::IngredientInfo;
-
-#[cfg(feature = "salsa_unstable")]
-pub(crate) use memory_usage::{MemoInfo, SlotInfo};
+pub use memory_usage::{MemoMemoryInfo, MemoryUsageVisitor, StructMemoryInfo};
 
 #[cfg(feature = "salsa_unstable")]
 mod memory_usage {
     use crate::Database;
-    use hashbrown::HashMap;
 
-    impl dyn Database {
-        /// Returns information about any Salsa structs.
-        pub fn structs_info(&self) -> Vec<IngredientInfo> {
-            self.zalsa()
-                .ingredients()
-                .filter_map(|ingredient| {
-                    let mut size_of_fields = 0;
-                    let mut size_of_metadata = 0;
-                    let mut instances = 0;
-                    let mut heap_size_of_fields = None;
-
-                    for slot in ingredient.memory_usage(self)? {
-                        instances += 1;
-                        size_of_fields += slot.size_of_fields;
-                        size_of_metadata += slot.size_of_metadata;
-
-                        if let Some(slot_heap_size) = slot.heap_size_of_fields {
-                            heap_size_of_fields =
-                                Some(heap_size_of_fields.unwrap_or_default() + slot_heap_size);
-                        }
-                    }
-
-                    Some(IngredientInfo {
-                        count: instances,
-                        size_of_fields,
-                        size_of_metadata,
-                        heap_size_of_fields,
-                        debug_name: ingredient.debug_name(),
-                    })
-                })
-                .collect()
+    pub trait MemoryUsageVisitor: std::any::Any {
+        fn visit_tracked_struct(&mut self, info: StructMemoryInfo) {
+            self.visit_struct(info);
         }
 
-        /// Returns information about any memoized Salsa queries.
-        ///
-        /// The returned map holds memory usage information for memoized values of a given query, keyed
-        /// by the query function name.
-        pub fn queries_info(&self) -> HashMap<&'static str, IngredientInfo> {
-            let mut queries = HashMap::new();
-
-            for input_ingredient in self.zalsa().ingredients() {
-                let Some(input_info) = input_ingredient.memory_usage(self) else {
-                    continue;
-                };
-
-                for input in input_info {
-                    for memo in input.memos {
-                        let info = queries.entry(memo.debug_name).or_insert(IngredientInfo {
-                            debug_name: memo.output.debug_name,
-                            ..Default::default()
-                        });
-
-                        info.count += 1;
-                        info.size_of_fields += memo.output.size_of_fields;
-                        info.size_of_metadata += memo.output.size_of_metadata;
-
-                        if let Some(memo_heap_size) = memo.output.heap_size_of_fields {
-                            info.heap_size_of_fields =
-                                Some(info.heap_size_of_fields.unwrap_or_default() + memo_heap_size);
-                        }
-                    }
-                }
-            }
+        fn visit_memo(&mut self, info: MemoMemoryInfo) {
+            let _ = info;
+        }
+
+        fn visit_input_struct(&mut self, info: StructMemoryInfo) {
+            self.visit_struct(info);
+        }
+
+        fn visit_interned_struct(&mut self, info: StructMemoryInfo) {
+            self.visit_struct(info);
+        }
+
+        fn visit_struct(&mut self, info: StructMemoryInfo) {
+            let _ = info;
+        }
 
-            queries
+        fn add_detail(&mut self, name: &'static str, size: usize) {
+            let (_, _) = (name, size);
         }
     }
 
-    /// Information about instances of a particular Salsa ingredient.
-    #[derive(Default, Debug, PartialEq, Eq, PartialOrd, Ord)]
-    pub struct IngredientInfo {
-        debug_name: &'static str,
-        count: usize,
-        size_of_metadata: usize,
-        size_of_fields: usize,
-        heap_size_of_fields: Option<usize>,
+    impl dyn Database {
+        /// Collects information about the memory usage of salsa structs and query functions.
+        pub fn memory_usage(&self, visitor: &mut dyn MemoryUsageVisitor) {
+            for ingredient in self.zalsa().ingredients() {
+                ingredient.memory_usage(self, visitor);
+            }
+        }
+    }
+
+    /// Memory usage information about a particular instance of struct, input, output, or memo.
+    #[derive(Debug, PartialEq, Eq)]
+    pub struct MemoMemoryInfo {
+        pub(crate) query_debug_name: &'static str,
+        pub(crate) result_debug_name: &'static str,
+        pub(crate) size_of_metadata: usize,
+        pub(crate) size_of_fields: usize,
+        pub(crate) heap_size_of_fields: Option<usize>,
     }
 
-    impl IngredientInfo {
+    impl MemoMemoryInfo {
         /// Returns the debug name of the ingredient.
-        pub fn debug_name(&self) -> &'static str {
-            self.debug_name
+        pub fn query_debug_name(&self) -> &'static str {
+            self.query_debug_name
+        }
+
+        /// Returns the debug name of the ingredient.
+        pub fn result_debug_name(&self) -> &'static str {
+            self.result_debug_name
         }
 
         /// Returns the total stack size of the fields of any instances of this ingredient, in bytes.
@@ -248,25 +216,37 @@ mod memory_usage {
         pub fn size_of_metadata(&self) -> usize {
             self.size_of_metadata
         }
-
-        /// Returns the number of instances of this ingredient.
-        pub fn count(&self) -> usize {
-            self.count
-        }
     }
 
-    /// Memory usage information about a particular instance of struct, input or output.
-    pub struct SlotInfo {
+    #[derive(Debug, PartialEq, Eq)]
+    pub struct StructMemoryInfo {
         pub(crate) debug_name: &'static str,
         pub(crate) size_of_metadata: usize,
         pub(crate) size_of_fields: usize,
         pub(crate) heap_size_of_fields: Option<usize>,
-        pub(crate) memos: Vec<MemoInfo>,
     }
 
-    /// Memory usage information about a particular memo.
-    pub struct MemoInfo {
-        pub(crate) debug_name: &'static str,
-        pub(crate) output: SlotInfo,
+    impl StructMemoryInfo {
+        /// Returns the debug name of the ingredient.
+        pub fn debug_name(&self) -> &'static str {
+            self.debug_name
+        }
+
+        /// Returns the total stack size of the fields of any instances of this ingredient, in bytes.
+        pub fn size_of_fields(&self) -> usize {
+            self.size_of_fields
+        }
+
+        /// Returns the total heap size of the fields of any instances of this ingredient, in bytes.
+        ///
+        /// Returns `None` if the ingredient doesn't specify a `heap_size` function.
+        pub fn heap_size_of_fields(&self) -> Option<usize> {
+            self.heap_size_of_fields
+        }
+
+        /// Returns the total size of Salsa metadata of any instances of this ingredient, in bytes.
+        pub fn size_of_metadata(&self) -> usize {
+            self.size_of_metadata
+        }
     }
 }

src/function.rs

@@ -74,7 +74,10 @@ pub trait Configuration: Any {
     fn id_to_input(db: &Self::DbView, key: Id) -> Self::Input<'_>;
 
     /// Returns the size of any heap allocations in the output value, in bytes.
-    fn heap_size(_value: &Self::Output<'_>) -> Option<usize> {
+    fn heap_size(
+        _value: &Self::Output<'_>,
+        _visitor: &mut dyn crate::MemoryUsageVisitor,
+    ) -> Option<usize> {
         None
     }
 

src/function/memo.rs

@@ -14,6 +14,8 @@ use crate::sync::atomic::Ordering;
 use crate::table::memo::MemoTableWithTypesMut;
 use crate::zalsa::{MemoIngredientIndex, Zalsa};
 use crate::zalsa_local::{QueryOriginRef, QueryRevisions, ZalsaLocal};
+#[cfg(feature = "salsa_unstable")]
+use crate::MemoryUsageVisitor;
 use crate::{Event, EventKind, Id, Revision};
 
 impl<C: Configuration> IngredientImpl<C> {
@@ -316,24 +318,21 @@ where
     }
 
     #[cfg(feature = "salsa_unstable")]
-    fn memory_usage(&self) -> crate::database::MemoInfo {
+    fn memory_usage(&self, visitor: &mut dyn MemoryUsageVisitor) {
         let size_of = std::mem::size_of::<Memo<C>>() + self.revisions.allocation_size();
         let heap_size = if let Some(value) = self.value.as_ref() {
-            C::heap_size(value)
+            C::heap_size(value, visitor)
         } else {
             Some(0)
         };
 
-        crate::database::MemoInfo {
-            debug_name: C::DEBUG_NAME,
-            output: crate::database::SlotInfo {
-                size_of_metadata: size_of - std::mem::size_of::<C::Output<'static>>(),
-                debug_name: std::any::type_name::<C::Output<'static>>(),
-                size_of_fields: std::mem::size_of::<C::Output<'static>>(),
-                heap_size_of_fields: heap_size,
-                memos: Vec::new(),
-            },
-        }
+        visitor.visit_memo(crate::database::MemoMemoryInfo {
+            size_of_metadata: size_of - std::mem::size_of::<C::Output<'static>>(),
+            query_debug_name: C::DEBUG_NAME,
+            result_debug_name: std::any::type_name::<C::Output<'static>>(),
+            size_of_fields: std::mem::size_of::<C::Output<'static>>(),
+            heap_size_of_fields: heap_size,
+        })
     }
 }
 

src/ingredient.rs

@@ -12,6 +12,8 @@ use crate::table::memo::MemoTableTypes;
 use crate::table::Table;
 use crate::zalsa::{transmute_data_mut_ptr, transmute_data_ptr, IngredientIndex, Zalsa};
 use crate::zalsa_local::QueryOriginRef;
+#[cfg(feature = "salsa_unstable")]
+use crate::MemoryUsageVisitor;
 use crate::{Database, DatabaseKeyIndex, Id, Revision};
 
 /// A "jar" is a group of ingredients that are added atomically.
@@ -171,8 +173,8 @@ pub trait Ingredient: Any + std::fmt::Debug + Send + Sync {
     /// Returns memory usage information about any instances of the ingredient,
     /// if applicable.
     #[cfg(feature = "salsa_unstable")]
-    fn memory_usage(&self, _db: &dyn Database) -> Option<Vec<crate::database::SlotInfo>> {
-        None
+    fn memory_usage(&self, db: &dyn Database, visitor: &mut dyn MemoryUsageVisitor) {
+        let _ = (db, visitor);
     }
 }
 

src/input.rs

@@ -19,6 +19,8 @@ use crate::sync::Arc;
 use crate::table::memo::{MemoTable, MemoTableTypes};
 use crate::table::{Slot, Table};
 use crate::zalsa::{IngredientIndex, Zalsa};
+#[cfg(feature = "salsa_unstable")]
+use crate::MemoryUsageVisitor;
 use crate::{Database, Durability, Id, Revision, Runtime};
 
 pub trait Configuration: Any {
@@ -42,7 +44,10 @@ pub trait Configuration: Any {
     type Durabilities: Send + Sync + fmt::Debug + IndexMut<usize, Output = Durability>;
 
     /// Returns the size of any heap allocations in the output value, in bytes.
-    fn heap_size(_value: &Self::Fields) -> Option<usize> {
+    fn heap_size(
+        _value: &Self::Fields,
+        _visitor: &mut dyn crate::MemoryUsageVisitor,
+    ) -> Option<usize> {
         None
     }
 }
@@ -252,14 +257,12 @@ impl<C: Configuration> Ingredient for IngredientImpl<C> {
 
     /// Returns memory usage information about any inputs.
     #[cfg(feature = "salsa_unstable")]
-    fn memory_usage(&self, db: &dyn Database) -> Option<Vec<crate::database::SlotInfo>> {
-        let memory_usage = self
-            .entries(db)
+    fn memory_usage(&self, db: &dyn Database, visitor: &mut dyn MemoryUsageVisitor) {
+        for value in self.entries(db) {
             // SAFETY: The memo table belongs to a value that we allocated, so it
             // has the correct type.
-            .map(|value| unsafe { value.memory_usage(&self.memo_table_types) })
-            .collect();
-        Some(memory_usage)
+            unsafe { value.memory_usage(&self.memo_table_types, visitor) }
+        }
     }
 }
 
@@ -311,18 +314,23 @@ where
     ///
     /// The `MemoTable` must belong to a `Value` of the correct type.
     #[cfg(feature = "salsa_unstable")]
-    unsafe fn memory_usage(&self, memo_table_types: &MemoTableTypes) -> crate::database::SlotInfo {
-        let heap_size = C::heap_size(&self.fields);
+    unsafe fn memory_usage(
+        &self,
+        memo_table_types: &MemoTableTypes,
+        visitor: &mut dyn MemoryUsageVisitor,
+    ) {
+        let heap_size = C::heap_size(&self.fields, visitor);
         // SAFETY: The caller guarantees this is the correct types table.
         let memos = unsafe { memo_table_types.attach_memos(&self.memos) };
 
-        crate::database::SlotInfo {
+        visitor.visit_input_struct(crate::database::StructMemoryInfo {
             debug_name: C::DEBUG_NAME,
             size_of_metadata: std::mem::size_of::<Self>() - std::mem::size_of::<C::Fields>(),
             size_of_fields: std::mem::size_of::<C::Fields>(),
             heap_size_of_fields: heap_size,
-            memos: memos.memory_usage(),
-        }
+        });
+
+        memos.memory_usage(visitor);
     }
 }