diff --git a/components/salsa-macro-rules/src/setup_input_struct.rs b/components/salsa-macro-rules/src/setup_input_struct.rs index d6d131ab..cc387136 100644 --- a/components/salsa-macro-rules/src/setup_input_struct.rs +++ b/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 { + Some($heap_size_fn(value)) + } + )? } impl $Configuration { diff --git a/components/salsa-macro-rules/src/setup_interned_struct.rs b/components/salsa-macro-rules/src/setup_interned_struct.rs index b637586e..ebf7f23c 100644 --- a/components/salsa-macro-rules/src/setup_interned_struct.rs +++ b/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 { + Some($heap_size_fn(value)) + } + )? } impl $Configuration { diff --git a/components/salsa-macro-rules/src/setup_tracked_fn.rs b/components/salsa-macro-rules/src/setup_tracked_fn.rs index 50ca6a03..de0c11a1 100644 --- a/components/salsa-macro-rules/src/setup_tracked_fn.rs +++ b/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 { + Some($heap_size_fn(value)) } )? diff --git a/components/salsa-macro-rules/src/setup_tracked_struct.rs b/components/salsa-macro-rules/src/setup_tracked_struct.rs index f92b1ac5..cb69a08e 100644 --- a/components/salsa-macro-rules/src/setup_tracked_struct.rs +++ b/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 { + Some($heap_size_fn(value)) + } + )? } impl $Configuration { diff --git a/components/salsa-macros/src/input.rs b/components/salsa-macros/src/input.rs index 71799f2b..b04176d6 100644 --- a/components/salsa-macros/src/input.rs +++ b/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, diff --git a/components/salsa-macros/src/interned.rs b/components/salsa-macros/src/interned.rs index 606701c6..dd064af1 100644 --- a/components/salsa-macros/src/interned.rs +++ b/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, diff --git a/components/salsa-macros/src/tracked_struct.rs b/components/salsa-macros/src/tracked_struct.rs index fe077da5..5768eb9c 100644 --- a/components/salsa-macros/src/tracked_struct.rs +++ b/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, diff --git a/src/database.rs b/src/database.rs index a253ac5a..933e137a 100644 --- a/src/database.rs +++ b/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, } 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 { + 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, pub(crate) memos: Vec, } diff --git a/src/function.rs b/src/function.rs index e9ab5793..6d4ec536 100644 --- a/src/function.rs +++ b/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 { + None } /// Invoked when we need to compute the value for the given key, either because we've never diff --git a/src/function/memo.rs b/src/function/memo.rs index 810e5b26..a6107060 100644 --- a/src/function/memo.rs +++ b/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::>() + 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::>(), debug_name: std::any::type_name::>(), - size_of_fields: std::mem::size_of::>() + heap_size, + size_of_fields: std::mem::size_of::>(), + heap_size_of_fields: heap_size, memos: Vec::new(), }, } diff --git a/src/input.rs b/src/input.rs index 58d769a6..464f79f1 100644 --- a/src/input.rs +++ b/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; + + /// Returns the size of any heap allocations in the output value, in bytes. + fn heap_size(_value: &Self::Fields) -> Option { + None + } } pub struct JarImpl { @@ -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::() - std::mem::size_of::(), size_of_fields: std::mem::size_of::(), + heap_size_of_fields: heap_size, memos: memos.memory_usage(), } } diff --git a/src/interned.rs b/src/interned.rs index 49663ee6..cdf4d61b 100644 --- a/src/interned.rs +++ b/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 { + 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::() - std::mem::size_of::>(), size_of_fields: std::mem::size_of::>(), + heap_size_of_fields: heap_size, memos: memos.memory_usage(), } } diff --git a/src/table/memo.rs b/src/table/memo.rs index b7bc5fb7..7e4837aa 100644 --- a/src/table/memo.rs +++ b/src/table/memo.rs @@ -14,7 +14,7 @@ pub struct MemoTable { } impl MemoTable { - /// Create a `MemoTable` with slots for memos from the provided `MemoTableTypes`. + /// Create a `MemoTable` with slots for memos from the provided `MemoTableTypes`. /// /// # Safety /// @@ -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(), }, } diff --git a/src/tracked_struct.rs b/src/tracked_struct.rs index 2ec3e24d..00986107 100644 --- a/src/tracked_struct.rs +++ b/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 { + 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::() - mem::size_of::>(), size_of_fields: mem::size_of::>(), + heap_size_of_fields: heap_size, memos: memos.memory_usage(), } } diff --git a/tests/compile-fail/input_struct_incompatibles.rs b/tests/compile-fail/input_struct_incompatibles.rs index 98cdb916..31ca9abb 100644 --- a/tests/compile-fail/input_struct_incompatibles.rs +++ b/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() {} diff --git a/tests/compile-fail/input_struct_incompatibles.stderr b/tests/compile-fail/input_struct_incompatibles.stderr index 9fe02527..a1b94e9a 100644 --- a/tests/compile-fail/input_struct_incompatibles.stderr +++ b/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 | diff --git a/tests/compile-fail/interned_struct_incompatibles.rs b/tests/compile-fail/interned_struct_incompatibles.rs index b8d50428..435335b1 100644 --- a/tests/compile-fail/interned_struct_incompatibles.rs +++ b/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() {} diff --git a/tests/compile-fail/interned_struct_incompatibles.stderr b/tests/compile-fail/interned_struct_incompatibles.stderr index 76ccc7f8..482e38b4 100644 --- a/tests/compile-fail/interned_struct_incompatibles.stderr +++ b/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 | diff --git a/tests/compile-fail/tracked_struct_incompatibles.rs b/tests/compile-fail/tracked_struct_incompatibles.rs index eff1eebd..5abd62dc 100644 --- a/tests/compile-fail/tracked_struct_incompatibles.rs +++ b/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() {} diff --git a/tests/compile-fail/tracked_struct_incompatibles.stderr b/tests/compile-fail/tracked_struct_incompatibles.stderr index e27777ca..928bbb12 100644 --- a/tests/compile-fail/tracked_struct_incompatibles.stderr +++ b/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)] - | ^^^^^^^^^ diff --git a/tests/memory-usage.rs b/tests/memory-usage.rs index f9fca29a..30a669e7 100644 --- a/tests/memory-usage.rs +++ b/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, }, ), ]"#]];