Add heap size support for salsa structs (#943)

* Improve unstable size analysis support

 1. Include an option `panic_if_missing` that will panic if there is an ingredient with no `heap_size()` defined, to ensure coverage.
 2. Add `heap_size()` to tracked structs, interneds an inputs.

* Make heap size a separate field, remove panic argument

* Remove stale comment

---------

Co-authored-by: Chayim Refael Friedman <chayimfr@gmail.com>

Author: MichaReiser

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

@@ -50,6 +50,9 @@ macro_rules! setup_input_struct {
         // If true, generate a debug impl.
         generate_debug_impl: $generate_debug_impl:tt,
 
+        // The function used to implement `C::heap_size`.
+        heap_size_fn: $($heap_size_fn:path)?,
+
         // Annoyingly macro-rules hygiene does not extend to items defined in the macro.
         // We have the procedural macro generate names for those items that are
         // not used elsewhere in the user's code.
@@ -98,6 +101,12 @@ macro_rules! setup_input_struct {
 
                 type Revisions = [$zalsa::Revision; $N];
                 type Durabilities = [$zalsa::Durability; $N];
+
+                $(
+                    fn heap_size(value: &Self::Fields) -> Option<usize> {
+                        Some($heap_size_fn(value))
+                    }
+                )?
             }
 
             impl $Configuration {

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

@@ -66,6 +66,9 @@ macro_rules! setup_interned_struct {
         // If true, generate a debug impl.
         generate_debug_impl: $generate_debug_impl:tt,
 
+        // The function used to implement `C::heap_size`.
+        heap_size_fn: $($heap_size_fn:path)?,
+
         // Annoyingly macro-rules hygiene does not extend to items defined in the macro.
         // We have the procedural macro generate names for those items that are
         // not used elsewhere in the user's code.
@@ -146,6 +149,12 @@ macro_rules! setup_interned_struct {
                 )?
                 type Fields<'a> = $StructDataIdent<'a>;
                 type Struct<'db> = $Struct< $($db_lt_arg)? >;
+
+                $(
+                    fn heap_size(value: &Self::Fields<'_>) -> Option<usize> {
+                        Some($heap_size_fn(value))
+                    }
+                )?
             }
 
             impl $Configuration {

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

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

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

@@ -88,6 +88,9 @@ macro_rules! setup_tracked_struct {
         // If true, generate a debug impl.
         generate_debug_impl: $generate_debug_impl:tt,
 
+        // The function used to implement `C::heap_size`.
+        heap_size_fn: $($heap_size_fn:path)?,
+
         // Annoyingly macro-rules hygiene does not extend to items defined in the macro.
         // We have the procedural macro generate names for those items that are
         // not used elsewhere in the user's code.
@@ -185,6 +188,12 @@ macro_rules! setup_tracked_struct {
                         )* false
                     }
                 }
+
+                $(
+                    fn heap_size(value: &Self::Fields<'_>) -> Option<usize> {
+                        Some($heap_size_fn(value))
+                    }
+                )?
             }
 
             impl $Configuration {

components/salsa-macros/src/input.rs

@@ -65,7 +65,7 @@ impl crate::options::AllowedOptions for InputStruct {
 
     const REVISIONS: bool = false;
 
-    const HEAP_SIZE: bool = false;
+    const HEAP_SIZE: bool = true;
 
     const SELF_TY: bool = false;
 }
@@ -112,6 +112,7 @@ impl Macro {
         let field_attrs = salsa_struct.field_attrs();
         let is_singleton = self.args.singleton.is_some();
         let generate_debug_impl = salsa_struct.generate_debug_impl();
+        let heap_size_fn = self.args.heap_size_fn.iter();
 
         let zalsa = self.hygiene.ident("zalsa");
         let zalsa_struct = self.hygiene.ident("zalsa_struct");
@@ -140,6 +141,7 @@ impl Macro {
                     num_fields: #num_fields,
                     is_singleton: #is_singleton,
                     generate_debug_impl: #generate_debug_impl,
+                    heap_size_fn: #(#heap_size_fn)*,
                     unused_names: [
                         #zalsa,
                         #zalsa_struct,

components/salsa-macros/src/interned.rs

@@ -65,7 +65,7 @@ impl crate::options::AllowedOptions for InternedStruct {
 
     const REVISIONS: bool = true;
 
-    const HEAP_SIZE: bool = false;
+    const HEAP_SIZE: bool = true;
 
     const SELF_TY: bool = false;
 }
@@ -131,6 +131,8 @@ impl Macro {
             (None, quote!(#struct_ident), static_lifetime)
         };
 
+        let heap_size_fn = self.args.heap_size_fn.iter();
+
         let zalsa = self.hygiene.ident("zalsa");
         let zalsa_struct = self.hygiene.ident("zalsa_struct");
         let Configuration = self.hygiene.ident("Configuration");
@@ -161,6 +163,7 @@ impl Macro {
                     field_attrs: [#([#(#field_unused_attrs),*]),*],
                     num_fields: #num_fields,
                     generate_debug_impl: #generate_debug_impl,
+                    heap_size_fn: #(#heap_size_fn)*,
                     unused_names: [
                         #zalsa,
                         #zalsa_struct,

components/salsa-macros/src/tracked_struct.rs

@@ -61,7 +61,7 @@ impl crate::options::AllowedOptions for TrackedStruct {
 
     const REVISIONS: bool = false;
 
-    const HEAP_SIZE: bool = false;
+    const HEAP_SIZE: bool = true;
 
     const SELF_TY: bool = false;
 }
@@ -141,6 +141,8 @@ impl Macro {
             }
         });
 
+        let heap_size_fn = self.args.heap_size_fn.iter();
+
         let num_tracked_fields = salsa_struct.num_tracked_fields();
         let generate_debug_impl = salsa_struct.generate_debug_impl();
 
@@ -188,6 +190,9 @@ impl Macro {
 
                     num_tracked_fields: #num_tracked_fields,
                     generate_debug_impl: #generate_debug_impl,
+
+                    heap_size_fn: #(#heap_size_fn)*,
+
                     unused_names: [
                         #zalsa,
                         #zalsa_struct,

src/database.rs

@@ -172,17 +172,24 @@ mod memory_usage {
                     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(),
                     })
                 })
@@ -211,6 +218,11 @@ mod memory_usage {
                         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);
+                        }
                     }
                 }
             }
@@ -226,6 +238,7 @@ mod memory_usage {
         count: usize,
         size_of_metadata: usize,
         size_of_fields: usize,
+        heap_size_of_fields: Option<usize>,
     }
 
     impl IngredientInfo {
@@ -234,11 +247,18 @@ mod memory_usage {
             self.debug_name
         }
 
-        /// Returns the total size of the fields of any instances of this ingredient, in bytes.
+        /// 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
@@ -255,6 +275,7 @@ mod memory_usage {
         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>,
     }
 

src/function.rs

@@ -75,8 +75,8 @@ pub trait Configuration: Any {
     fn id_to_input(zalsa: &Zalsa, 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
+    fn heap_size(_value: &Self::Output<'_>) -> Option<usize> {
+        None
     }
 
     /// Invoked when we need to compute the value for the given key, either because we've never

src/function/memo.rs

@@ -321,14 +321,19 @@ where
     #[cfg(feature = "salsa_unstable")]
     fn memory_usage(&self) -> crate::database::MemoInfo {
         let size_of = std::mem::size_of::<Memo<C>>() + self.revisions.allocation_size();
-        let heap_size = self.value.as_ref().map(C::heap_size).unwrap_or(0);
+        let heap_size = if let Some(value) = self.value.as_ref() {
+            C::heap_size(value)
+        } 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,
+                size_of_fields: std::mem::size_of::<C::Output<'static>>(),
+                heap_size_of_fields: heap_size,
                 memos: Vec::new(),
             },
         }