Skip serializing empty type headers

Author: matt-cornell

cobalt-ast/src/context.rs

@@ -26,6 +26,7 @@ pub struct Flags {
     pub all_move_metadata: bool,
     pub private_syms: bool,
     pub skip_header_version_check: bool,
+    pub add_type_map: bool,
 }
 impl Default for Flags {
     fn default() -> Self {
@@ -37,6 +38,7 @@ impl Default for Flags {
             all_move_metadata: false,
             private_syms: true,
             skip_header_version_check: false,
+            add_type_map: false,
         }
     }
 }
@@ -331,9 +333,13 @@ impl Serialize for CtxTypeSerde<'_, '_, '_> {
     {
         use ser::*;
         let tsr = TYPE_SERIAL_REGISTRY.pin();
-        let mut map = serializer.serialize_map(Some(tsr.len()))?;
+        let mut map = serializer.serialize_map(Some(
+            tsr.iter().filter(|(_, info)| (info.has_header)()).count(),
+        ))?;
         for (id, info) in &tsr {
-            map.serialize_entry(&HexArray(id.to_le_bytes()), &(info.erased_header)())?;
+            if (info.has_header)() {
+                map.serialize_entry(&HexArray(id.to_le_bytes()), &(info.erased_header)())?;
+            }
         }
         map.end()
     }
@@ -384,8 +390,19 @@ impl Serialize for CompCtx<'_, '_> {
                 }
             }
         });
-        let mut map = serializer.serialize_struct("Context", 3)?;
+        let mut map =
+            serializer.serialize_struct("Context", 3 + usize::from(self.flags.add_type_map))?;
         map.serialize_field("version", &HEADER_FMT_VERSION)?;
+        if self.flags.add_type_map {
+            map.serialize_field(
+                "names",
+                &TYPE_SERIAL_REGISTRY
+                    .pin()
+                    .iter()
+                    .map(|(k, v)| (hex::encode(k.to_le_bytes()), v.name))
+                    .collect::<hashbrown::HashMap<_, _>>(),
+            )?;
+        }
         map.serialize_field("types", &CtxTypeSerde(self))?;
         self.with_vars(|v| map.serialize_field("vars", v))?;
         SERIALIZATION_CONTEXT.with(|c| {
@@ -399,6 +416,8 @@ impl Serialize for CompCtx<'_, '_> {
 #[serde(bound = "'a: 'de")]
 struct ContextDeProxy<'a> {
     version: u16,
+    #[serde(rename = "names")]
+    _names: Option<serde::de::IgnoredAny>, // if it gets into the serialization, ignore it for deserialization - it should be stable
     #[serde(borrow = "'a")]
     types: serde::__private::de::Content<'a>,
     #[serde(borrow = "'a")]

cobalt-ast/src/types.rs

@@ -87,12 +87,18 @@ impl<T: ConcreteType> TypeKind for T {
 pub trait TypeSerde: Sized {
     type Header: Serialize + DeserializeOwned;
     type Proxy: Serialize + DeserializeOwned;
+    fn has_header() -> bool {
+        true
+    }
     fn get_header() -> Self::Header;
     fn set_header(header: Self::Header);
     fn get_proxy(&'static self) -> Self::Proxy;
     fn from_proxy(body: Self::Proxy) -> &'static Self;
 }
 pub trait TypeSerdeFns {
+    fn has_header() -> bool
+    where
+        Self: Sized;
     fn erased_header() -> Box<dyn erased_serde::Serialize>
     where
         Self: Sized;
@@ -109,6 +115,9 @@ pub trait TypeSerdeFns {
         Self: Sized;
 }
 impl<T: Type + TypeSerde> TypeSerdeFns for T {
+    fn has_header() -> bool {
+        <Self as TypeSerde>::has_header()
+    }
     fn erased_header() -> Box<dyn erased_serde::Serialize> {
         Box::new(Self::get_header())
     }
@@ -136,6 +145,9 @@ pub trait NullType: Type {
 impl<T: NullType> TypeSerde for T {
     type Header = ();
     type Proxy = ();
+    fn has_header() -> bool {
+        false
+    }
     fn get_header() -> Self::Header {}
     fn set_header(header: Self::Header) {}
     fn get_proxy(&'static self) -> Self::Proxy {}
@@ -146,6 +158,9 @@ impl<T: NullType> TypeSerde for T {
 macro_rules! no_type_header {
     () => {
         type Header = ();
+        fn has_header() -> bool {
+            false
+        }
         fn get_header() -> Self::Header {}
         fn set_header(header: Self::Header) {}
     };
@@ -732,6 +747,8 @@ impl<'de> Deserialize<'de> for TypeRef {
 }
 pub struct TypeLoader {
     pub kind: u64,
+    pub name: fn() -> &'static str,
+    pub has_header: fn() -> bool,
     pub erased_header: fn() -> Box<dyn erased_serde::Serialize>,
     pub load_header: fn(&mut dyn erased_serde::Deserializer) -> Result<(), erased_serde::Error>,
     pub load: fn(&mut dyn erased_serde::Deserializer) -> Result<TypeRef, erased_serde::Error>,
@@ -741,6 +758,11 @@ macro_rules! type_loader {
     ($T:ty) => {
         $crate::types::TypeLoader {
             kind: <$T as $crate::types::ConcreteType>::KIND,
+            name: || {
+                let s = ::std::any::type_name::<$T>();
+                s.rsplit_once("::").map_or(s, |x| x.1)
+            },
+            has_header: <$T as $crate::types::TypeSerdeFns>::has_header,
             erased_header: <$T as $crate::types::TypeSerdeFns>::erased_header,
             load_header: <$T as $crate::types::TypeSerdeFns>::load_header,
             load: <$T as $crate::types::TypeSerdeFns>::load,
@@ -758,23 +780,29 @@ macro_rules! submit_types {
 }
 inventory::collect!(TypeLoader);
 pub struct LoadInfo {
+    pub has_header: fn() -> bool,
     pub erased_header: fn() -> Box<dyn erased_serde::Serialize>,
     pub load_header: fn(&mut dyn erased_serde::Deserializer) -> Result<(), erased_serde::Error>,
     pub load: fn(&mut dyn erased_serde::Deserializer) -> Result<TypeRef, erased_serde::Error>,
+    pub name: &'static str,
 }
 impl LoadInfo {
     pub fn new(
         &TypeLoader {
             kind,
+            has_header,
             erased_header,
             load_header,
             load,
+            name,
             ..
         }: &'static TypeLoader,
     ) -> (u64, Self) {
         (
             kind,
             Self {
+                name: name(),
+                has_header,
                 erased_header,
                 load_header,
                 load,

cobalt-ast/src/types/custom.rs

@@ -243,6 +243,9 @@ impl<'de> Deserialize<'de> for CustomHeader {
 }
 impl TypeSerde for Custom {
     type Header = CustomHeader;
+    fn has_header() -> bool {
+        !CUSTOM_DATA.is_empty() // may return a false positive but much faster than iterating
+    }
     fn get_header() -> Self::Header {
         CustomHeader
     }

cobalt-cli/src/lib.rs

@@ -1393,6 +1393,7 @@ pub fn driver(cli: Cli) -> anyhow::Result<()> {
             {
                 flags.word_size = size as u16;
             }
+            flags.add_type_map = dump_header;
             let ctx = CompCtx::with_flags(&ink_ctx, &input_name, flags);
             let trip = TargetMachine::get_default_triple();
             let triple = trip.as_str().to_string_lossy();
@@ -2360,6 +2361,7 @@ pub fn driver(cli: Cli) -> anyhow::Result<()> {
                 {
                     flags.word_size = size as u16;
                 }
+                flags.add_type_map = dump_header;
                 let ctx = CompCtx::with_flags(&ink_ctx, "multi-check", flags);
                 ctx.module.set_triple(&trip);
                 let mut cc = cc::CompileCommand::new();