diff --git a/Cargo.toml b/Cargo.toml index f22eb9565..e433ee0c4 100644 --- a/Cargo.toml +++ b/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" diff --git a/hugr-core/src/envelope/reader.rs b/hugr-core/src/envelope/reader.rs index ad245479a..eee9cba4a 100644 --- a/hugr-core/src/envelope/reader.rs +++ b/hugr-core/src/envelope/reader.rs @@ -222,26 +222,33 @@ impl EnvelopeReader { 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::>() - .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::>>(); + 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::>>()?; 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)?; diff --git a/hugr-core/src/envelope/writer.rs b/hugr-core/src/envelope/writer.rs index 9638d07ec..e63fff565 100644 --- a/hugr-core/src/envelope/writer.rs +++ b/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(); diff --git a/hugr-core/src/extension.rs b/hugr-core/src/extension.rs index fbdcff833..e6f9241a1 100644 --- a/hugr-core/src/extension.rs +++ b/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 { - let extensions: Vec = serde_json::from_reader(reader)?; + let encoded_extensions: Vec> = + serde_json::from_reader(reader)?; + let extensions = encoded_extensions + .into_iter() + .map(|raw| Extension::from_raw_json(&raw)) + .collect::>>()?; // After deserialization, we need to update all the internal // `Weak` 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> { 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>, + /// The JSON object loaded from disk or generated by the envelope writer. + #[serde(skip)] + encoded_json: OnceLock>, +} + +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 { + 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::>(&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::*}; diff --git a/hugr-core/src/extension/op_def.rs b/hugr-core/src/extension/op_def.rs index c27719d68..c8b94b68a 100644 --- a/hugr-core/src/extension/op_def.rs +++ b/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, diff --git a/hugr-core/src/extension/type_def.rs b/hugr-core/src/extension/type_def.rs index 8c72a0f78..abdfc0cd3 100644 --- a/hugr-core/src/extension/type_def.rs +++ b/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)),