Add option to track heap memory usage of memos

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

@@ -55,6 +55,9 @@ macro_rules! setup_tracked_fn {
         // If true, the input needs an interner (because it has >1 argument).
         needs_interner: $needs_interner:tt,
 
+        // The function used to implement `C::heap_size`.
+        heap_size_fn: $($heap_size_fn:path)?,
+
         // LRU capacity (a literal, maybe 0)
         lru: $lru:tt,
 
@@ -196,6 +199,12 @@ macro_rules! setup_tracked_fn {
 
                 $($values_equal)+
 
+                $(
+                    fn heap_size(value: &Self::Output<'_>) -> usize {
+                        $heap_size_fn(value)
+                    }
+                )?
+
                 fn execute<$db_lt>($db: &$db_lt Self::DbView, ($($input_id),*): ($($interned_input_ty),*)) -> Self::Output<$db_lt> {
                     $($assert_return_type_is_update)*
 

components/salsa-macros/src/accumulator.rs

@@ -45,6 +45,7 @@ impl AllowedOptions for Accumulator {
     const CONSTRUCTOR_NAME: bool = false;
     const ID: bool = false;
     const REVISIONS: bool = false;
+    const HEAP_SIZE: bool = false;
 }
 
 struct StructMacro {

components/salsa-macros/src/input.rs

@@ -64,6 +64,8 @@ impl crate::options::AllowedOptions for InputStruct {
     const ID: bool = false;
 
     const REVISIONS: bool = false;
+
+    const HEAP_SIZE: bool = false;
 }
 
 impl SalsaStructAllowedOptions for InputStruct {

components/salsa-macros/src/interned.rs

@@ -64,6 +64,8 @@ impl crate::options::AllowedOptions for InternedStruct {
     const ID: bool = true;
 
     const REVISIONS: bool = true;
+
+    const HEAP_SIZE: bool = false;
 }
 
 impl SalsaStructAllowedOptions for InternedStruct {

components/salsa-macros/src/options.rs

@@ -99,6 +99,12 @@ pub(crate) struct Options<A: AllowedOptions> {
     /// This is stored as a `syn::Expr` to support `usize::MAX`.
     pub revisions: Option<syn::Expr>,
 
+    /// The `heap_size = <path>` option can be used to track heap memory usage of memoized
+    /// values.
+    ///
+    /// If this is `Some`, the value is the provided `heap_size` function.
+    pub heap_size_fn: Option<syn::Path>,
+
     /// Remember the `A` parameter, which plays no role after parsing.
     phantom: PhantomData<A>,
 }
@@ -123,6 +129,7 @@ impl<A: AllowedOptions> Default for Options<A> {
             singleton: Default::default(),
             id: Default::default(),
             revisions: Default::default(),
+            heap_size_fn: Default::default(),
         }
     }
 }
@@ -145,6 +152,7 @@ pub(crate) trait AllowedOptions {
     const CONSTRUCTOR_NAME: bool;
     const ID: bool;
     const REVISIONS: bool;
+    const HEAP_SIZE: bool;
 }
 
 type Equals = syn::Token![=];
@@ -392,6 +400,22 @@ impl<A: AllowedOptions> syn::parse::Parse for Options<A> {
                         "`revisions` option not allowed here",
                     ));
                 }
+            } else if ident == "heap_size" {
+                if A::HEAP_SIZE {
+                    let _eq = Equals::parse(input)?;
+                    let path = syn::Path::parse(input)?;
+                    if let Some(old) = options.heap_size_fn.replace(path) {
+                        return Err(syn::Error::new(
+                            old.span(),
+                            "option `heap_size` provided twice",
+                        ));
+                    }
+                } else {
+                    return Err(syn::Error::new(
+                        ident.span(),
+                        "`heap_size` option not allowed here",
+                    ));
+                }
             } else {
                 return Err(syn::Error::new(
                     ident.span(),

components/salsa-macros/src/tracked_fn.rs

@@ -57,6 +57,8 @@ impl crate::options::AllowedOptions for TrackedFn {
     const ID: bool = false;
 
     const REVISIONS: bool = false;
+
+    const HEAP_SIZE: bool = true;
 }
 
 struct Macro {
@@ -97,6 +99,7 @@ impl Macro {
             self.cycle_recovery()?;
         let is_specifiable = self.args.specify.is_some();
         let requires_update = self.args.non_update_return_type.is_none();
+        let heap_size_fn = self.args.heap_size_fn.iter();
         let eq = if let Some(token) = &self.args.no_eq {
             if self.args.cycle_fn.is_some() {
                 return Err(syn::Error::new_spanned(
@@ -217,6 +220,7 @@ impl Macro {
                 is_specifiable: #is_specifiable,
                 values_equal: {#eq},
                 needs_interner: #needs_interner,
+                heap_size_fn: #(#heap_size_fn)*,
                 lru: #lru,
                 return_mode: #return_mode,
                 assert_return_type_is_update: { #assert_return_type_is_update },

components/salsa-macros/src/tracked_struct.rs

@@ -60,6 +60,8 @@ impl crate::options::AllowedOptions for TrackedStruct {
     const ID: bool = false;
 
     const REVISIONS: bool = false;
+
+    const HEAP_SIZE: bool = false;
 }
 
 impl SalsaStructAllowedOptions for TrackedStruct {

src/database.rs

@@ -135,7 +135,10 @@ impl dyn Database {
 }
 
 #[cfg(feature = "salsa_unstable")]
-pub use memory_usage::{IngredientInfo, SlotInfo};
+pub use memory_usage::IngredientInfo;
+
+#[cfg(feature = "salsa_unstable")]
+pub(crate) use memory_usage::{MemoInfo, SlotInfo};
 
 #[cfg(feature = "salsa_unstable")]
 mod memory_usage {
@@ -171,8 +174,8 @@ mod memory_usage {
         /// Returns information about any memoized Salsa queries.
         ///
         /// The returned map holds memory usage information for memoized values of a given query, keyed
-        /// by its `(input, output)` type names.
-        pub fn queries_info(&self) -> HashMap<(&'static str, &'static str), IngredientInfo> {
+        /// 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() {
@@ -181,17 +184,15 @@ mod memory_usage {
                 };
 
                 for input in input_info {
-                    for output in input.memos {
-                        let info = queries
-                            .entry((input.debug_name, output.debug_name))
-                            .or_insert(IngredientInfo {
-                                debug_name: output.debug_name,
-                                ..Default::default()
-                            });
+                    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 += output.size_of_fields;
-                        info.size_of_metadata += output.size_of_metadata;
+                        info.size_of_fields += memo.output.size_of_fields;
+                        info.size_of_metadata += memo.output.size_of_metadata;
                     }
                 }
             }
@@ -236,6 +237,12 @@ mod memory_usage {
         pub(crate) debug_name: &'static str,
         pub(crate) size_of_metadata: usize,
         pub(crate) size_of_fields: usize,
-        pub(crate) memos: Vec<SlotInfo>,
+        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,
     }
 }

src/function.rs

@@ -36,7 +36,7 @@ mod memo;
 mod specify;
 mod sync;
 
-pub type Memo<C> = memo::Memo<<C as Configuration>::Output<'static>>;
+pub type Memo<C> = memo::Memo<'static, C>;
 
 pub trait Configuration: Any {
     const DEBUG_NAME: &'static str;
@@ -72,6 +72,11 @@ pub trait Configuration: Any {
     /// This is a no-op if the input to the function is a salsa struct.
     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<'_>) -> usize {
+        0
+    }
+
     /// Invoked when we need to compute the value for the given key, either because we've never
     /// computed it before or because the old one relied on inputs that have changed.
     ///
@@ -181,8 +186,8 @@ where
     /// only cleared with `&mut self`.
     unsafe fn extend_memo_lifetime<'this>(
         &'this self,
-        memo: &memo::Memo<C::Output<'this>>,
-    ) -> &'this memo::Memo<C::Output<'this>> {
+        memo: &memo::Memo<'this, C>,
+    ) -> &'this memo::Memo<'this, C> {
         // SAFETY: the caller must guarantee that the memo will not be released before `&self`
         unsafe { std::mem::transmute(memo) }
     }
@@ -191,9 +196,9 @@ where
         &'db self,
         zalsa: &'db Zalsa,
         id: Id,
-        mut memo: memo::Memo<C::Output<'db>>,
+        mut memo: memo::Memo<'db, C>,
         memo_ingredient_index: MemoIngredientIndex,
-    ) -> &'db memo::Memo<C::Output<'db>> {
+    ) -> &'db memo::Memo<'db, C> {
         if let Some(tracked_struct_ids) = memo.revisions.tracked_struct_ids_mut() {
             tracked_struct_ids.shrink_to_fit();
         }

src/function/backdate.rs

@@ -12,7 +12,7 @@ where
     /// on an old memo when a new memo has been produced to check whether there have been changed.
     pub(super) fn backdate_if_appropriate<'db>(
         &self,
-        old_memo: &Memo<C::Output<'db>>,
+        old_memo: &Memo<'db, C>,
         index: DatabaseKeyIndex,
         revisions: &mut QueryRevisions,
         value: &C::Output<'db>,

src/function/delete.rs

@@ -7,7 +7,7 @@ use crate::function::Configuration;
 /// once the next revision starts. See the comment on the field
 /// `deleted_entries` of [`FunctionIngredient`][] for more details.
 pub(super) struct DeletedEntries<C: Configuration> {
-    memos: boxcar::Vec<SharedBox<Memo<C::Output<'static>>>>,
+    memos: boxcar::Vec<SharedBox<Memo<'static, C>>>,
 }
 
 #[allow(clippy::undocumented_unsafe_blocks)] // TODO(#697) document safety
@@ -27,13 +27,10 @@ impl<C: Configuration> DeletedEntries<C> {
     /// # Safety
     ///
     /// The memo must be valid and safe to free when the `DeletedEntries` list is cleared or dropped.
-    pub(super) unsafe fn push(&self, memo: NonNull<Memo<C::Output<'_>>>) {
+    pub(super) unsafe fn push(&self, memo: NonNull<Memo<'_, C>>) {
         // Safety: The memo must be valid and safe to free when the `DeletedEntries` list is cleared or dropped.
-        let memo = unsafe {
-            std::mem::transmute::<NonNull<Memo<C::Output<'_>>>, NonNull<Memo<C::Output<'static>>>>(
-                memo,
-            )
-        };
+        let memo =
+            unsafe { std::mem::transmute::<NonNull<Memo<'_, C>>, NonNull<Memo<'static, C>>>(memo) };
 
         self.memos.push(SharedBox(memo));
     }

src/function/diff_outputs.rs

@@ -18,7 +18,7 @@ where
         &self,
         zalsa: &Zalsa,
         key: DatabaseKeyIndex,
-        old_memo: &Memo<C::Output<'_>>,
+        old_memo: &Memo<'_, C>,
         revisions: &mut QueryRevisions,
     ) {
         let (QueryOriginRef::Derived(edges) | QueryOriginRef::DerivedUntracked(edges)) =

src/function/execute.rs

@@ -24,8 +24,8 @@ where
         &'db self,
         db: &'db C::DbView,
         active_query: ActiveQueryGuard<'db>,
-        opt_old_memo: Option<&Memo<C::Output<'db>>>,
-    ) -> &'db Memo<C::Output<'db>> {
+        opt_old_memo: Option<&Memo<'db, C>>,
+    ) -> &'db Memo<'db, C> {
         let database_key_index = active_query.database_key_index;
         let id = database_key_index.key_index();
 
@@ -121,7 +121,7 @@ where
         &'db self,
         db: &'db C::DbView,
         mut active_query: ActiveQueryGuard<'db>,
-        opt_old_memo: Option<&Memo<C::Output<'db>>>,
+        opt_old_memo: Option<&Memo<'db, C>>,
         zalsa: &'db Zalsa,
         id: Id,
         memo_ingredient_index: MemoIngredientIndex,
@@ -133,7 +133,7 @@ where
         // Our provisional value from the previous iteration, when doing fixpoint iteration.
         // Initially it's set to None, because the initial provisional value is created lazily,
         // only when a cycle is actually encountered.
-        let mut opt_last_provisional: Option<&Memo<<C as Configuration>::Output<'db>>> = None;
+        let mut opt_last_provisional: Option<&Memo<'db, C>> = None;
         loop {
             let previous_memo = opt_last_provisional.or(opt_old_memo);
             let (mut new_value, mut revisions) = Self::execute_query(
@@ -257,7 +257,7 @@ where
     fn execute_query<'db>(
         db: &'db C::DbView,
         active_query: ActiveQueryGuard<'db>,
-        opt_old_memo: Option<&Memo<C::Output<'db>>>,
+        opt_old_memo: Option<&Memo<'db, C>>,
         current_revision: Revision,
         id: Id,
     ) -> (C::Output<'db>, QueryRevisions) {

src/function/fetch.rs

@@ -43,7 +43,7 @@ where
         zalsa: &'db Zalsa,
         zalsa_local: &'db ZalsaLocal,
         id: Id,
-    ) -> &'db Memo<C::Output<'db>> {
+    ) -> &'db Memo<'db, C> {
         let memo_ingredient_index = self.memo_ingredient_index(zalsa, id);
         loop {
             if let Some(memo) = self
@@ -63,7 +63,7 @@ where
         zalsa: &'db Zalsa,
         id: Id,
         memo_ingredient_index: MemoIngredientIndex,
-    ) -> Option<&'db Memo<C::Output<'db>>> {
+    ) -> Option<&'db Memo<'db, C>> {
         let memo = self.get_memo_from_table_for(zalsa, id, memo_ingredient_index)?;
 
         memo.value.as_ref()?;
@@ -91,7 +91,7 @@ where
         db: &'db C::DbView,
         id: Id,
         memo_ingredient_index: MemoIngredientIndex,
-    ) -> Option<&'db Memo<C::Output<'db>>> {
+    ) -> Option<&'db Memo<'db, C>> {
         let memo = self.fetch_cold(zalsa, zalsa_local, db, id, memo_ingredient_index)?;
 
         // If we get back a provisional cycle memo, and it's provisional on any cycle heads
@@ -117,7 +117,7 @@ where
         db: &'db C::DbView,
         id: Id,
         memo_ingredient_index: MemoIngredientIndex,
-    ) -> Option<&'db Memo<C::Output<'db>>> {
+    ) -> Option<&'db Memo<'db, C>> {
         let database_key_index = self.database_key_index(id);
         // Try to claim this query: if someone else has claimed it already, go back and start again.
         let claim_guard = match self.sync_table.try_claim(zalsa, id) {