Skip to content

Commit d0ceeb4

Browse files
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.
1 parent d28d66b commit d0ceeb4

23 files changed

+231
-86
lines changed

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

Lines changed: 9 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -50,6 +50,9 @@ macro_rules! setup_input_struct {
5050
// If true, generate a debug impl.
5151
generate_debug_impl: $generate_debug_impl:tt,
5252

53+
// The function used to implement `C::heap_size`.
54+
heap_size_fn: $($heap_size_fn:path)?,
55+
5356
// Annoyingly macro-rules hygiene does not extend to items defined in the macro.
5457
// We have the procedural macro generate names for those items that are
5558
// not used elsewhere in the user's code.
@@ -89,6 +92,12 @@ macro_rules! setup_input_struct {
8992

9093
type Revisions = [$zalsa::Revision; $N];
9194
type Durabilities = [$zalsa::Durability; $N];
95+
96+
$(
97+
fn heap_size(value: &Self::Fields, _panic_if_missing: $zalsa::PanicIfHeapSizeMissing) -> usize {
98+
$heap_size_fn(value)
99+
}
100+
)?
92101
}
93102

94103
impl $Configuration {

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

Lines changed: 9 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -66,6 +66,9 @@ macro_rules! setup_interned_struct {
6666
// If true, generate a debug impl.
6767
generate_debug_impl: $generate_debug_impl:tt,
6868

69+
// The function used to implement `C::heap_size`.
70+
heap_size_fn: $($heap_size_fn:path)?,
71+
6972
// Annoyingly macro-rules hygiene does not extend to items defined in the macro.
7073
// We have the procedural macro generate names for those items that are
7174
// not used elsewhere in the user's code.
@@ -137,6 +140,12 @@ macro_rules! setup_interned_struct {
137140
)?
138141
type Fields<'a> = $StructDataIdent<'a>;
139142
type Struct<'db> = $Struct< $($db_lt_arg)? >;
143+
144+
$(
145+
fn heap_size(value: &Self::Fields<'_>, _panic_if_missing: $zalsa::PanicIfHeapSizeMissing) -> usize {
146+
$heap_size_fn(value)
147+
}
148+
)?
140149
}
141150

142151
impl $Configuration {

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

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -211,7 +211,7 @@ macro_rules! setup_tracked_fn {
211211
$($values_equal)+
212212

213213
$(
214-
fn heap_size(value: &Self::Output<'_>) -> usize {
214+
fn heap_size(value: &Self::Output<'_>, _panic_if_missing: $zalsa::PanicIfHeapSizeMissing) -> usize {
215215
$heap_size_fn(value)
216216
}
217217
)?

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

Lines changed: 9 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -88,6 +88,9 @@ macro_rules! setup_tracked_struct {
8888
// If true, generate a debug impl.
8989
generate_debug_impl: $generate_debug_impl:tt,
9090

91+
// The function used to implement `C::heap_size`.
92+
heap_size_fn: $($heap_size_fn:path)?,
93+
9194
// Annoyingly macro-rules hygiene does not extend to items defined in the macro.
9295
// We have the procedural macro generate names for those items that are
9396
// not used elsewhere in the user's code.
@@ -176,6 +179,12 @@ macro_rules! setup_tracked_struct {
176179
)* false
177180
}
178181
}
182+
183+
$(
184+
fn heap_size(value: &Self::Fields<'_>, _panic_if_missing: $zalsa::PanicIfHeapSizeMissing) -> usize {
185+
$heap_size_fn(value)
186+
}
187+
)?
179188
}
180189

181190
impl $Configuration {

components/salsa-macros/src/input.rs

Lines changed: 3 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -65,7 +65,7 @@ impl crate::options::AllowedOptions for InputStruct {
6565

6666
const REVISIONS: bool = false;
6767

68-
const HEAP_SIZE: bool = false;
68+
const HEAP_SIZE: bool = true;
6969

7070
const SELF_TY: bool = false;
7171
}
@@ -112,6 +112,7 @@ impl Macro {
112112
let field_attrs = salsa_struct.field_attrs();
113113
let is_singleton = self.args.singleton.is_some();
114114
let generate_debug_impl = salsa_struct.generate_debug_impl();
115+
let heap_size_fn = self.args.heap_size_fn.iter();
115116

116117
let zalsa = self.hygiene.ident("zalsa");
117118
let zalsa_struct = self.hygiene.ident("zalsa_struct");
@@ -140,6 +141,7 @@ impl Macro {
140141
num_fields: #num_fields,
141142
is_singleton: #is_singleton,
142143
generate_debug_impl: #generate_debug_impl,
144+
heap_size_fn: #(#heap_size_fn)*,
143145
unused_names: [
144146
#zalsa,
145147
#zalsa_struct,

components/salsa-macros/src/interned.rs

Lines changed: 4 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -65,7 +65,7 @@ impl crate::options::AllowedOptions for InternedStruct {
6565

6666
const REVISIONS: bool = true;
6767

68-
const HEAP_SIZE: bool = false;
68+
const HEAP_SIZE: bool = true;
6969

7070
const SELF_TY: bool = false;
7171
}
@@ -131,6 +131,8 @@ impl Macro {
131131
(None, quote!(#struct_ident), static_lifetime)
132132
};
133133

134+
let heap_size_fn = self.args.heap_size_fn.iter();
135+
134136
let zalsa = self.hygiene.ident("zalsa");
135137
let zalsa_struct = self.hygiene.ident("zalsa_struct");
136138
let Configuration = self.hygiene.ident("Configuration");
@@ -161,6 +163,7 @@ impl Macro {
161163
field_attrs: [#([#(#field_unused_attrs),*]),*],
162164
num_fields: #num_fields,
163165
generate_debug_impl: #generate_debug_impl,
166+
heap_size_fn: #(#heap_size_fn)*,
164167
unused_names: [
165168
#zalsa,
166169
#zalsa_struct,

components/salsa-macros/src/tracked_struct.rs

Lines changed: 6 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -61,7 +61,7 @@ impl crate::options::AllowedOptions for TrackedStruct {
6161

6262
const REVISIONS: bool = false;
6363

64-
const HEAP_SIZE: bool = false;
64+
const HEAP_SIZE: bool = true;
6565

6666
const SELF_TY: bool = false;
6767
}
@@ -141,6 +141,8 @@ impl Macro {
141141
}
142142
});
143143

144+
let heap_size_fn = self.args.heap_size_fn.iter();
145+
144146
let num_tracked_fields = salsa_struct.num_tracked_fields();
145147
let generate_debug_impl = salsa_struct.generate_debug_impl();
146148

@@ -188,6 +190,9 @@ impl Macro {
188190

189191
num_tracked_fields: #num_tracked_fields,
190192
generate_debug_impl: #generate_debug_impl,
193+
194+
heap_size_fn: #(#heap_size_fn)*,
195+
191196
unused_names: [
192197
#zalsa,
193198
#zalsa_struct,

src/database.rs

Lines changed: 24 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -140,22 +140,35 @@ pub use memory_usage::IngredientInfo;
140140
#[cfg(feature = "salsa_unstable")]
141141
pub(crate) use memory_usage::{MemoInfo, SlotInfo};
142142

143+
/// Whether to panic on the default `heap_size()`, to ensure the user has provided a custom one.
144+
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
145+
pub enum PanicIfHeapSizeMissing {
146+
Yes,
147+
No,
148+
}
149+
143150
#[cfg(feature = "salsa_unstable")]
144151
mod memory_usage {
145-
use crate::Database;
152+
use crate::{Database, PanicIfHeapSizeMissing};
146153
use hashbrown::HashMap;
147154

148155
impl dyn Database {
149156
/// Returns information about any Salsa structs.
150-
pub fn structs_info(&self) -> Vec<IngredientInfo> {
157+
///
158+
/// If `panic_is_missing` is [`PanicIfHeapSizeMissing::Yes`], and there is an ingredient with no `heap_size()` function,
159+
/// this function will panic. This can be used to ensure coverage.
160+
pub fn structs_info(
161+
&self,
162+
panic_if_missing: PanicIfHeapSizeMissing,
163+
) -> Vec<IngredientInfo> {
151164
self.zalsa()
152165
.ingredients()
153166
.filter_map(|ingredient| {
154167
let mut size_of_fields = 0;
155168
let mut size_of_metadata = 0;
156169
let mut instances = 0;
157170

158-
for slot in ingredient.memory_usage(self)? {
171+
for slot in ingredient.memory_usage(self, panic_if_missing)? {
159172
instances += 1;
160173
size_of_fields += slot.size_of_fields;
161174
size_of_metadata += slot.size_of_metadata;
@@ -175,11 +188,17 @@ mod memory_usage {
175188
///
176189
/// The returned map holds memory usage information for memoized values of a given query, keyed
177190
/// by the query function name.
178-
pub fn queries_info(&self) -> HashMap<&'static str, IngredientInfo> {
191+
///
192+
/// If `panic_is_missing` is [`PanicIfHeapSizeMissing::Yes`], and there is an ingredient with no `heap_size()` function,
193+
/// this function will panic. This can be used to ensure coverage.
194+
pub fn queries_info(
195+
&self,
196+
panic_if_missing: PanicIfHeapSizeMissing,
197+
) -> HashMap<&'static str, IngredientInfo> {
179198
let mut queries = HashMap::new();
180199

181200
for input_ingredient in self.zalsa().ingredients() {
182-
let Some(input_info) = input_ingredient.memory_usage(self) else {
201+
let Some(input_info) = input_ingredient.memory_usage(self, panic_if_missing) else {
183202
continue;
184203
};
185204

src/function.rs

Lines changed: 12 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -73,7 +73,18 @@ pub trait Configuration: Any {
7373
fn id_to_input(db: &Self::DbView, key: Id) -> Self::Input<'_>;
7474

7575
/// Returns the size of any heap allocations in the output value, in bytes.
76-
fn heap_size(_value: &Self::Output<'_>) -> usize {
76+
fn heap_size(
77+
_value: &Self::Output<'_>,
78+
panic_if_missing: crate::PanicIfHeapSizeMissing,
79+
) -> usize {
80+
if panic_if_missing == crate::PanicIfHeapSizeMissing::Yes {
81+
panic!(
82+
"tried to estimate sizes but `size_of()` was not defined.\n\
83+
ingredient: {}\ntype: {}",
84+
Self::DEBUG_NAME,
85+
std::any::type_name::<Self::Output<'_>>()
86+
);
87+
}
7788
0
7889
}
7990

src/function/memo.rs

Lines changed: 9 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -314,9 +314,16 @@ where
314314
}
315315

316316
#[cfg(feature = "salsa_unstable")]
317-
fn memory_usage(&self) -> crate::database::MemoInfo {
317+
fn memory_usage(
318+
&self,
319+
panic_if_missing: crate::PanicIfHeapSizeMissing,
320+
) -> crate::database::MemoInfo {
318321
let size_of = std::mem::size_of::<Memo<C>>() + self.revisions.allocation_size();
319-
let heap_size = self.value.as_ref().map(C::heap_size).unwrap_or(0);
322+
let heap_size = self
323+
.value
324+
.as_ref()
325+
.map(|value| C::heap_size(value, panic_if_missing))
326+
.unwrap_or(0);
320327

321328
crate::database::MemoInfo {
322329
debug_name: C::DEBUG_NAME,

0 commit comments

Comments
 (0)