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,