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>

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(),
             },
         }

src/input.rs

@@ -40,6 +40,11 @@ pub trait Configuration: Any {
 
     /// A array of [`Durability`], one per each of the value fields.
     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> {
+        None
+    }
 }
 
 pub struct JarImpl<C: Configuration> {
@@ -307,6 +312,7 @@ 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);
         // SAFETY: The caller guarantees this is the correct types table.
         let memos = unsafe { memo_table_types.attach_memos(&self.memos) };
 
@@ -314,6 +320,7 @@ where
             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(),
         }
     }

src/interned.rs

@@ -44,6 +44,11 @@ pub trait Configuration: Sized + 'static {
 
     /// The end user struct
     type Struct<'db>: Copy + FromId + AsId;
+
+    /// Returns the size of any heap allocations in the output value, in bytes.
+    fn heap_size(_value: &Self::Fields<'_>) -> Option<usize> {
+        None
+    }
 }
 
 pub trait InternedData: Sized + Eq + Hash + Clone + Sync + Send {}
@@ -199,6 +204,7 @@ where
     /// lock must be held for the shard containing the value.
     #[cfg(all(not(feature = "shuttle"), feature = "salsa_unstable"))]
     unsafe fn memory_usage(&self, memo_table_types: &MemoTableTypes) -> crate::database::SlotInfo {
+        let heap_size = C::heap_size(self.fields());
         // SAFETY: The caller guarantees we hold the lock for the shard containing the value, so we
         // have at-least read-only access to the value's memos.
         let memos = unsafe { &*self.memos.get() };
@@ -209,6 +215,7 @@ where
             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(),
         }
     }

src/table/memo.rs

@@ -127,6 +127,7 @@ impl Memo for DummyMemo {
                 debug_name: "dummy",
                 size_of_metadata: 0,
                 size_of_fields: 0,
+                heap_size_of_fields: None,
                 memos: Vec::new(),
             },
         }

src/tracked_struct.rs

@@ -90,6 +90,11 @@ pub trait Configuration: Sized + 'static {
         old_fields: *mut Self::Fields<'db>,
         new_fields: Self::Fields<'db>,
     ) -> bool;
+
+    /// Returns the size of any heap allocations in the output value, in bytes.
+    fn heap_size(_value: &Self::Fields<'_>) -> Option<usize> {
+        None
+    }
 }
 // ANCHOR_END: Configuration
 
@@ -935,6 +940,7 @@ 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());
         // SAFETY: The caller guarantees this is the correct types table.
         let memos = unsafe { memo_table_types.attach_memos(&self.memos) };
 
@@ -942,6 +948,7 @@ where
             debug_name: C::DEBUG_NAME,
             size_of_metadata: mem::size_of::<Self>() - mem::size_of::<C::Fields<'_>>(),
             size_of_fields: mem::size_of::<C::Fields<'_>>(),
+            heap_size_of_fields: heap_size,
             memos: memos.memory_usage(),
         }
     }

tests/compile-fail/input_struct_incompatibles.rs

@@ -25,7 +25,4 @@ struct InputWithTrackedField {
     field: u32,
 }
 
-#[salsa::input(heap_size = size)]
-struct InputWithHeapSize(u32);
-
 fn main() {}

tests/compile-fail/input_struct_incompatibles.stderr

@@ -47,12 +47,6 @@ error: `#[tracked]` cannot be used with `#[salsa::input]`
 25 | |     field: u32,
    | |______________^
 
-error: `heap_size` option not allowed here
-  --> tests/compile-fail/input_struct_incompatibles.rs:28:16
-   |
-28 | #[salsa::input(heap_size = size)]
-   |                ^^^^^^^^^
-
 error: cannot find attribute `tracked` in this scope
   --> tests/compile-fail/input_struct_incompatibles.rs:24:7
    |

tests/compile-fail/interned_struct_incompatibles.rs

@@ -39,9 +39,4 @@ struct InternedWithZeroRevisions {
     field: u32,
 }
 
-#[salsa::interned(heap_size = size)]
-struct AccWithHeapSize {
-    field: u32,
-}
-
 fn main() {}

tests/compile-fail/interned_struct_incompatibles.stderr

@@ -41,12 +41,6 @@ error: `#[tracked]` cannot be used with `#[salsa::interned]`
 34 | |     field: u32,
    | |______________^
 
-error: `heap_size` option not allowed here
-  --> tests/compile-fail/interned_struct_incompatibles.rs:42:19
-   |
-42 | #[salsa::interned(heap_size = size)]
-   |                   ^^^^^^^^^
-
 error: cannot find attribute `tracked` in this scope
   --> tests/compile-fail/interned_struct_incompatibles.rs:33:7
    |

tests/compile-fail/tracked_struct_incompatibles.rs

@@ -33,9 +33,4 @@ struct TrackedStructWithRevisions {
     field: u32,
 }
 
-#[salsa::tracked(heap_size = size)]
-struct TrackedStructWithHeapSize {
-    field: u32,
-}
-
 fn main() {}

tests/compile-fail/tracked_struct_incompatibles.stderr

@@ -39,9 +39,3 @@ error: `revisions` option not allowed here
    |
 31 | #[salsa::tracked(revisions = 12)]
    |                  ^^^^^^^^^
-
-error: `heap_size` option not allowed here
-  --> tests/compile-fail/tracked_struct_incompatibles.rs:36:18
-   |
-36 | #[salsa::tracked(heap_size = size)]
-   |                  ^^^^^^^^^

tests/memory-usage.rs

@@ -2,19 +2,19 @@
 
 use expect_test::expect;
 
-#[salsa::input]
+#[salsa::input(heap_size = string_tuple_size_of)]
 struct MyInput {
-    field: u32,
+    field: String,
 }
 
-#[salsa::tracked]
+#[salsa::tracked(heap_size = string_tuple_size_of)]
 struct MyTracked<'db> {
-    field: u32,
+    field: String,
 }
 
-#[salsa::interned]
+#[salsa::interned(heap_size = string_tuple_size_of)]
 struct MyInterned<'db> {
-    field: u32,
+    field: String,
 }
 
 #[salsa::tracked]
@@ -32,12 +32,16 @@ fn input_to_string<'db>(_db: &'db dyn salsa::Database) -> String {
     "a".repeat(1000)
 }
 
-#[salsa::tracked(heap_size = string_heap_size)]
+#[salsa::tracked(heap_size = string_size_of)]
 fn input_to_string_get_size<'db>(_db: &'db dyn salsa::Database) -> String {
     "a".repeat(1000)
 }
 
-fn string_heap_size(x: &String) -> usize {
+fn string_size_of(x: &String) -> usize {
+    x.capacity()
+}
+
+fn string_tuple_size_of((x,): &(String,)) -> usize {
     x.capacity()
 }
 
@@ -56,9 +60,9 @@ fn input_to_tracked_tuple<'db>(
 fn test() {
     let db = salsa::DatabaseImpl::new();
 
-    let input1 = MyInput::new(&db, 1);
-    let input2 = MyInput::new(&db, 2);
-    let input3 = MyInput::new(&db, 3);
+    let input1 = MyInput::new(&db, "a".repeat(50));
+    let input2 = MyInput::new(&db, "a".repeat(150));
+    let input3 = MyInput::new(&db, "a".repeat(250));
 
     let _tracked1 = input_to_tracked(&db, input1);
     let _tracked2 = input_to_tracked(&db, input2);
@@ -79,32 +83,43 @@ fn test() {
             IngredientInfo {
                 debug_name: "MyInput",
                 count: 3,
-                size_of_metadata: 84,
-                size_of_fields: 12,
+                size_of_metadata: 96,
+                size_of_fields: 72,
+                heap_size_of_fields: Some(
+                    450,
+                ),
             },
             IngredientInfo {
                 debug_name: "MyTracked",
                 count: 4,
-                size_of_metadata: 112,
-                size_of_fields: 16,
+                size_of_metadata: 128,
+                size_of_fields: 96,
+                heap_size_of_fields: Some(
+                    300,
+                ),
             },
             IngredientInfo {
                 debug_name: "MyInterned",
                 count: 3,
-                size_of_metadata: 156,
-                size_of_fields: 12,
+                size_of_metadata: 168,
+                size_of_fields: 72,
+                heap_size_of_fields: Some(
+                    450,
+                ),
             },
             IngredientInfo {
                 debug_name: "input_to_string::interned_arguments",
                 count: 1,
                 size_of_metadata: 56,
                 size_of_fields: 0,
+                heap_size_of_fields: None,
             },
             IngredientInfo {
                 debug_name: "input_to_string_get_size::interned_arguments",
                 count: 1,
                 size_of_metadata: 56,
                 size_of_fields: 0,
+                heap_size_of_fields: None,
             },
         ]"#]];
 
@@ -124,6 +139,7 @@ fn test() {
                     count: 3,
                     size_of_metadata: 192,
                     size_of_fields: 24,
+                    heap_size_of_fields: None,
                 },
             ),
             (
@@ -133,6 +149,7 @@ fn test() {
                     count: 1,
                     size_of_metadata: 40,
                     size_of_fields: 24,
+                    heap_size_of_fields: None,
                 },
             ),
             (
@@ -141,7 +158,10 @@ fn test() {
                     debug_name: "alloc::string::String",
                     count: 1,
                     size_of_metadata: 40,
-                    size_of_fields: 1024,
+                    size_of_fields: 24,
+                    heap_size_of_fields: Some(
+                        1000,
+                    ),
                 },
             ),
             (
@@ -151,6 +171,7 @@ fn test() {
                     count: 2,
                     size_of_metadata: 192,
                     size_of_fields: 16,
+                    heap_size_of_fields: None,
                 },
             ),
             (
@@ -160,6 +181,7 @@ fn test() {
                     count: 1,
                     size_of_metadata: 132,
                     size_of_fields: 16,
+                    heap_size_of_fields: None,
                 },
             ),
         ]"#]];