perf: Cache serialized extensions to speed up encoding

Cargo.toml

@@ -62,7 +62,7 @@ regex-syntax = "0.8.9"
 rstest = "0.26.1"
 semver = "1.0.28"
 serde = "1.0.228"
-serde_json = "1.0.150"
+serde_json = { version = "1.0.150", features = ["raw_value"] }
 serde_with = "3.20.0"
 serde_yaml = "0.9.34"
 smallvec = "1.15.1"

hugr-core/src/envelope/reader.rs

@@ -222,26 +222,33 @@ impl<R: BufRead> EnvelopeReader<R> {
         let format = self.header().format;
         check_model_version(format)?;
 
-        let packaged_extensions = if format == EnvelopeFormat::SExpressionWithExtensions {
-            let deserializer = serde_json::Deserializer::from_reader(&mut self.reader);
-            // Deserialize the first json object, leaving the rest of the reader unconsumed.
-            let extra_extensions = deserializer
-                .into_iter::<Vec<Extension>>()
-                .next()
-                .unwrap_or(Ok(vec![]))?;
+        // The S-expression parser already needs the complete model in memory. Reading
+        // the payload up front also lets us retain the extension JSON prefix.
+        let mut buffer = String::new();
+        self.reader.read_to_string(&mut buffer)?;
+
+        let (packaged_extensions, model_start) = if format
+            == EnvelopeFormat::SExpressionWithExtensions
+        {
+            let mut extensions = serde_json::Deserializer::from_str(&buffer)
+                .into_iter::<Vec<Box<serde_json::value::RawValue>>>();
+            let encoded_extensions = extensions.next().unwrap_or(Ok(vec![]))?;
+            let model_start = extensions.byte_offset();
+            let extra_extensions = encoded_extensions
+                .into_iter()
+                .map(|encoded| Extension::from_raw_json(&encoded))
+                .collect::<serde_json::Result<Vec<_>>>()?;
             let weak_registry: WeakExtensionRegistry = (&self.registry).into();
-            ExtensionRegistry::new_with_extension_resolution(extra_extensions, &weak_registry)
-                .map_err(ExtensionRegistryLoadError::from)?
+            let registry =
+                ExtensionRegistry::new_with_extension_resolution(extra_extensions, &weak_registry)
+                    .map_err(ExtensionRegistryLoadError::from)?;
+            (registry, model_start)
         } else {
-            ExtensionRegistry::new([])
+            (ExtensionRegistry::new([]), 0)
         };
 
-        // Read the package into a string, then parse it.
-        //
-        // Due to how `to_string` works, we cannot append extensions after the package.
-        let mut buffer = String::new();
-        self.reader.read_to_string(&mut buffer)?;
-        let ast_package = hugr_model::v0::ast::Package::from_str(&buffer)?;
+        // Due to how `to_string` works, extensions must precede the model.
+        let ast_package = hugr_model::v0::ast::Package::from_str(&buffer[model_start..])?;
 
         let bump = Bump::default();
         let model_package = ast_package.resolve(&bump)?;

hugr-core/src/envelope/writer.rs

@@ -1,6 +1,5 @@
 use std::io::Write;
 
-use itertools::Itertools as _;
 use thiserror::Error;
 
 use crate::Hugr;
@@ -82,7 +81,7 @@ fn encode_model_binary<'h>(
 
     // Append extensions for binary model.
     if format == EnvelopeFormat::ModelWithExtensions {
-        serde_json::to_writer(writer, &extensions.iter_all().collect_vec())?;
+        extensions.write_json(writer)?;
     }
 
     Ok(())
@@ -101,7 +100,7 @@ fn encode_model_text<'h>(
 
     // Prepend extensions for text model.
     if format == EnvelopeFormat::SExpressionWithExtensions {
-        serde_json::to_writer(&mut writer, &extensions.iter_all().collect_vec())?;
+        extensions.write_json(&mut writer)?;
     }
 
     let bump = Bump::default();

hugr-core/src/extension.rs

@@ -117,7 +117,7 @@ use std::collections::btree_map;
 use std::collections::{BTreeMap, BTreeSet};
 use std::fmt::Debug;
 use std::sync::atomic::{AtomicBool, Ordering};
-use std::sync::{Arc, Weak};
+use std::sync::{Arc, OnceLock, Weak};
 use std::{io, mem};
 
 use derive_more::Display;
@@ -398,7 +398,12 @@ impl ExtensionRegistry {
         reader: impl io::Read,
         other_extensions: &ExtensionRegistry,
     ) -> Result<Self, ExtensionRegistryLoadError> {
-        let extensions: Vec<Extension> = serde_json::from_reader(reader)?;
+        let encoded_extensions: Vec<Box<serde_json::value::RawValue>> =
+            serde_json::from_reader(reader)?;
+        let extensions = encoded_extensions
+            .into_iter()
+            .map(|raw| Extension::from_raw_json(&raw))
+            .collect::<serde_json::Result<Vec<_>>>()?;
         // After deserialization, we need to update all the internal
         // `Weak<Extension>` references.
         Ok(ExtensionRegistry::new_with_extension_resolution(
@@ -407,6 +412,18 @@ impl ExtensionRegistry {
         )?)
     }
 
+    /// Write the registry as a JSON array of individually cached extensions.
+    pub(crate) fn write_json(&self, mut writer: impl io::Write) -> serde_json::Result<()> {
+        writer.write_all(b"[").map_err(serde_json::Error::io)?;
+        for (index, extension) in self.iter_all().enumerate() {
+            if index > 0 {
+                writer.write_all(b",").map_err(serde_json::Error::io)?;
+            }
+            extension.write_json(&mut writer)?;
+        }
+        writer.write_all(b"]").map_err(serde_json::Error::io)
+    }
+
     /// Gets the Extension with the given name
     pub fn get(&self, name: &str) -> Option<&Arc<Extension>> {
         self.exts.get(name).map(ExtensionVersions::latest)
@@ -867,7 +884,7 @@ pub type ExtensionId = IdentList;
 ///     },
 /// );
 /// ```
-#[derive(Clone, Debug, serde::Serialize, serde::Deserialize)]
+#[derive(serde::Serialize, serde::Deserialize)]
 pub struct Extension {
     /// Extension version, follows semver.
     pub version: Version,
@@ -882,6 +899,33 @@ pub struct Extension {
     // and the other references to the OpDef are from ExternalOp's in the Hugr
     // (which are serialized as OpaqueOp's i.e. Strings).
     operations: BTreeMap<OpName, Arc<op_def::OpDef>>,
+    /// The JSON object loaded from disk or generated by the envelope writer.
+    #[serde(skip)]
+    encoded_json: OnceLock<Arc<[u8]>>,
+}
+
+impl Clone for Extension {
+    fn clone(&self) -> Self {
+        Self {
+            version: self.version.clone(),
+            name: self.name.clone(),
+            types: self.types.clone(),
+            operations: self.operations.clone(),
+            // A mutable clone may diverge from the original definition.
+            encoded_json: OnceLock::new(),
+        }
+    }
+}
+
+impl Debug for Extension {
+    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+        f.debug_struct("Extension")
+            .field("version", &self.version)
+            .field("name", &self.name)
+            .field("types", &self.types)
+            .field("operations", &self.operations)
+            .finish()
+    }
 }
 
 impl Extension {
@@ -899,6 +943,7 @@ impl Extension {
             version,
             types: Default::default(),
             operations: Default::default(),
+            encoded_json: OnceLock::new(),
         }
     }
 
@@ -1005,6 +1050,46 @@ impl Extension {
         }
         Ok(())
     }
+
+    /// Write this extension as JSON.
+    ///
+    /// Reuses its cached encoding when available.
+    fn write_json(&self, mut writer: impl io::Write) -> serde_json::Result<()> {
+        if let Some(encoded) = self.encoded_json.get() {
+            return writer.write_all(encoded).map_err(serde_json::Error::io);
+        }
+
+        let encoded: Arc<[u8]> = serde_json::to_vec(self)?.into();
+        let _ = self.encoded_json.set(encoded);
+        let encoded = self
+            .encoded_json
+            .get()
+            .expect("the extension JSON cache was just initialized");
+        writer.write_all(encoded).map_err(serde_json::Error::io)
+    }
+
+    /// Deserialize an extension from a raw json value.
+    ///
+    /// Retains the encoded JSON in the extension's cache, to speed up future
+    /// serialization. The in-memory extension may be updated during resolution
+    /// to point to updated versions of transitive extension dependencies, but
+    /// the original serialization always remains valid.
+    ///
+    /// If the encoded JSON contains additional fields not recognized by the
+    /// current version of the extension, they will be ignored. These will
+    /// remain in the encoded JSON cache.
+    pub(crate) fn from_raw_json(encoded: &serde_json::value::RawValue) -> serde_json::Result<Self> {
+        let extension: Self = serde_json::from_str(encoded.get())?;
+        let _ = extension
+            .encoded_json
+            .set(encoded.get().as_bytes().to_vec().into());
+        Ok(extension)
+    }
+
+    /// Discard an encoding before changing this extension's definition.
+    fn clear_json_cache(&mut self) {
+        let _ = self.encoded_json.take();
+    }
 }
 
 impl PartialEq for Extension {
@@ -1287,6 +1372,70 @@ pub mod test {
         assert_eq!(reg.len(), 1);
     }
 
+    #[rstest::rstest]
+    fn loaded_json_is_reused() {
+        let json = br#"[{ "version" : "1.0.0", "name":"cached","types":{},"operations":{}}]"#;
+        let registry = ExtensionRegistry::load_json(json.as_slice(), &EMPTY_REG).unwrap();
+        let mut encoded = Vec::new();
+
+        registry.write_json(&mut encoded).unwrap();
+
+        assert_eq!(encoded, json);
+    }
+
+    #[rstest::rstest]
+    fn registry_mutation_preserves_extension_json() {
+        let json = br#"[{ "version" : "1.0.0", "name":"cached","types":{},"operations":{}}]"#;
+        let mut registry = ExtensionRegistry::load_json(json.as_slice(), &EMPTY_REG).unwrap();
+        let added_id = ExtensionId::new("added").unwrap();
+        registry.register(Arc::new(Extension::new(
+            added_id.clone(),
+            Version::new(1, 0, 0),
+        )));
+        let mut encoded = Vec::new();
+
+        registry.write_json(&mut encoded).unwrap();
+
+        assert_ne!(encoded, json);
+        let cached_extension = &json[1..json.len() - 1];
+        assert!(
+            encoded
+                .windows(cached_extension.len())
+                .any(|window| window == cached_extension)
+        );
+        assert_eq!(
+            serde_json::from_slice::<Vec<Extension>>(&encoded)
+                .unwrap()
+                .len(),
+            2
+        );
+
+        registry.remove_extension(&added_id);
+        encoded.clear();
+        registry.write_json(&mut encoded).unwrap();
+        assert_eq!(encoded, json);
+    }
+
+    #[rstest::rstest]
+    fn definition_mutation_invalidates_extension_json() {
+        let mut extension =
+            Extension::new(ExtensionId::new("cached").unwrap(), Version::new(1, 0, 0));
+        extension.write_json(Vec::new()).unwrap();
+        assert!(extension.encoded_json.get().is_some());
+
+        extension
+            .add_type(
+                "new_type".into(),
+                vec![],
+                String::new(),
+                TypeDefBound::any(),
+                &Weak::new(),
+            )
+            .unwrap();
+
+        assert!(extension.encoded_json.get().is_none());
+    }
+
     mod proptest {
 
         use ::proptest::{collection::hash_set, prelude::*};

hugr-core/src/extension/op_def.rs

@@ -631,6 +631,7 @@ impl Extension {
             constant_folder: Default::default(),
         };
 
+        self.clear_json_cache();
         match self.operations.entry(op.name.clone()) {
             Entry::Occupied(_) => Err(ExtensionBuildError::OpDefExists(op.name)),
             // Just made the arc so should only be one reference to it, can get_mut,

hugr-core/src/extension/type_def.rs

@@ -259,6 +259,7 @@ impl Extension {
             description,
             bound,
         };
+        self.clear_json_cache();
         match self.types.entry(ty.name.clone()) {
             Entry::Occupied(_) => Err(ExtensionBuildError::TypeDefExists(ty.name)),
             Entry::Vacant(ve) => Ok(ve.insert(ty)),