feat!: Only expose envelope serialization of hugrs and packages

Author: aborgna-q

Cargo.lock

@@ -60,6 +60,7 @@ rstest = "0.24.0"
 semver = "1.0.26"
 serde = "1.0.219"
 serde_json = "1.0.140"
+serde_with = "3.12.0"
 serde_yaml = "0.9.34"
 smol_str = "0.3.1"
 static_assertions = "1.1.0"

Cargo.toml

@@ -5,7 +5,7 @@ use hugr::envelope::{read_envelope, EnvelopeError};
 use hugr::extension::ExtensionRegistry;
 use hugr::package::Package;
 use hugr::{Extension, Hugr};
-use std::io::BufReader;
+use std::io::{BufReader, Read};
 use std::path::PathBuf;
 
 use crate::CliError;
@@ -57,12 +57,23 @@ impl HugrInputArgs {
 
     /// Read a hugr JSON file from the input.
     ///
-    /// If [HugrInputArgs::hugr_json] is `false`, [HugrInputArgs::get_package] should be called instead as
-    /// reading an input envelope as a HUGR json will fail.
+    /// This is a legacy option for reading old HUGR JSON files when the
+    /// [HugrInputArgs::hugr_json] flag is used.
+    ///
+    /// For most cases, [HugrInputArgs::get_package] should be called instead.
     pub fn get_hugr(&mut self) -> Result<Hugr, CliError> {
         let extensions = self.load_extensions()?;
         let mut buffer = BufReader::new(&mut self.input);
-        let hugr = Hugr::load_json(&mut buffer, &extensions)?;
+
+        /// Wraps the hugr JSON so that it defines a valid envelope.
+        const PREPEND: &str = r#"HUGRiHJv?@{"modules": ["#;
+        const APPEND: &str = r#"],"extensions": []}"#;
+
+        let mut envelope = PREPEND.to_string();
+        buffer.read_to_string(&mut envelope)?;
+        envelope.push_str(APPEND);
+
+        let hugr = Hugr::load_str(envelope, Some(&extensions))?;
         Ok(hugr)
     }
 

hugr-cli/src/hugr_io.rs

@@ -60,8 +60,7 @@ use clap::{crate_version, Parser};
 use clap_verbosity_flag::log::Level;
 use clap_verbosity_flag::{InfoLevel, Verbosity};
 use hugr::envelope::EnvelopeError;
-use hugr::hugr::LoadHugrError;
-use hugr::package::{PackageEncodingError, PackageValidationError};
+use hugr::package::PackageValidationError;
 use std::ffi::OsString;
 
 pub mod extensions;
@@ -97,12 +96,6 @@ pub enum CliError {
     /// Error parsing input.
     #[display("Error parsing package: {_0}")]
     Parse(serde_json::Error),
-    /// Package load error.
-    #[display("Error parsing package: {_0}")]
-    PackageLoad(PackageEncodingError),
-    /// Hugr load error.
-    #[display("Error loading hugr: {_0}")]
-    HugrLoad(LoadHugrError),
     #[display("Error validating HUGR: {_0}")]
     /// Errors produced by the `validate` subcommand.
     Validate(PackageValidationError),

hugr-cli/src/lib.rs

@@ -110,30 +110,28 @@ fn bad_hugr_string() -> String {
     let df = DFGBuilder::new(Signature::new_endo(vec![qb_t()])).unwrap();
     let bad_hugr = df.hugr().clone();
 
-    serde_json::to_string(&bad_hugr).unwrap()
+    bad_hugr.store_str(EnvelopeConfig::text()).unwrap()
 }
 
 #[rstest]
 fn test_mermaid_invalid(bad_hugr_string: String, mut cmd: Command) {
     cmd.arg("mermaid");
     cmd.arg("--validate");
-    cmd.arg("--hugr-json");
     cmd.write_stdin(bad_hugr_string);
     cmd.assert()
         .failure()
-        .stderr(contains("Error loading hugr"));
+        .stderr(contains("Error validating HUGR"));
 }
 
 #[rstest]
 fn test_bad_hugr(bad_hugr_string: String, mut val_cmd: Command) {
     val_cmd.write_stdin(bad_hugr_string);
-    val_cmd.arg("--hugr-json");
     val_cmd.arg("-");
 
     val_cmd
         .assert()
         .failure()
-        .stderr(contains("Error loading hugr"));
+        .stderr(contains("Error validating HUGR"));
 }
 
 #[rstest]

hugr-cli/tests/validate.rs

@@ -51,6 +51,7 @@ regex = { workspace = true }
 # Rc used here for Extension, but unfortunately we must turn the feature on globally
 serde = { workspace = true, features = ["derive", "rc"] }
 serde_json = { workspace = true }
+serde_with = { workspace = true }
 serde_yaml = { workspace = true, optional = true }
 smol_str = { workspace = true, features = ["serde"] }
 static_assertions = { workspace = true }

hugr-core/Cargo.toml

@@ -41,13 +41,14 @@
 // https://github.com/JelteF/derive_more/pull/454.
 
 mod header;
+mod package_json;
+pub mod serde_with;
 
 pub use header::{EnvelopeConfig, EnvelopeFormat, ZstdConfig, MAGIC_NUMBERS};
+pub use package_json::PackageEncodingError;
 
-use crate::{
-    extension::ExtensionRegistry,
-    package::{Package, PackageEncodingError},
-};
+use crate::Hugr;
+use crate::{extension::ExtensionRegistry, package::Package};
 use header::EnvelopeHeader;
 use std::io::BufRead;
 use std::io::Write;
@@ -90,9 +91,22 @@ pub fn read_envelope(
 /// It is recommended to use a buffered writer for better performance.
 /// See [`std::io::BufWriter`] for more information.
 pub fn write_envelope(
-    mut writer: impl Write,
+    writer: impl Write,
     package: &Package,
     config: EnvelopeConfig,
+) -> Result<(), EnvelopeError> {
+    write_envelope_impl(writer, &package.modules, &package.extensions, config)
+}
+
+/// Write a deconstructed HUGR package into an envelope, using the specified configuration.
+///
+/// It is recommended to use a buffered writer for better performance.
+/// See [`std::io::BufWriter`] for more information.
+pub(crate) fn write_envelope_impl<'h>(
+    mut writer: impl Write,
+    hugrs: impl IntoIterator<Item = &'h Hugr>,
+    extensions: &ExtensionRegistry,
+    config: EnvelopeConfig,
 ) -> Result<(), EnvelopeError> {
     let header = config.make_header();
     header.write(&mut writer)?;
@@ -101,11 +115,11 @@ pub fn write_envelope(
         #[cfg(feature = "zstd")]
         Some(zstd) => {
             let writer = zstd::Encoder::new(writer, zstd.level())?.auto_finish();
-            write_impl(writer, package, config)?;
+            write_impl(writer, hugrs, extensions, config)?;
         }
         #[cfg(not(feature = "zstd"))]
         Some(_) => return Err(EnvelopeError::ZstdUnsupported),
-        None => write_impl(writer, package, config)?,
+        None => write_impl(writer, hugrs, extensions, config)?,
     }
 
     Ok(())
@@ -164,14 +178,14 @@ pub enum EnvelopeError {
     #[display("Zstd compression is not supported. This requires the 'zstd' feature for `hugr`.")]
     #[from(ignore)]
     ZstdUnsupported,
-    /// Tried to encode a package with multiple HUGRs, when only 1 was expected.
-    #[display(
-            "Packages with multiple HUGRs are currently unsupported. Tried to encode {count} HUGRs, when 1 was expected."
-        )]
+    /// Expected the envelope to contain a single HUGR.
+    #[display("Expected an envelope containing a single hugr, but it contained {}.", if *count == 0 {
+        "none".to_string()
+    } else {
+        count.to_string()
+    })]
     #[from(ignore)]
-    /// Deprecated: Packages with multiple HUGRs is a legacy feature that is no longer supported.
-    #[deprecated(since = "0.15.2", note = "Multiple HUGRs are supported via packages.")]
-    MultipleHugrs {
+    ExpectedSingleHugr {
         /// The number of HUGRs in the package.
         count: usize,
     },
@@ -215,7 +229,7 @@ fn read_impl(
 ) -> Result<Package, EnvelopeError> {
     match header.format {
         #[allow(deprecated)]
-        EnvelopeFormat::PackageJson => Ok(Package::from_json_reader(payload, registry)?),
+        EnvelopeFormat::PackageJson => Ok(package_json::from_json_reader(payload, registry)?),
         EnvelopeFormat::Model | EnvelopeFormat::ModelWithExtensions => {
             decode_model(payload, registry, header.format)
         }
@@ -260,24 +274,26 @@ fn decode_model(
 }
 
 /// Internal implementation of [`write_envelope`] to call with/without the zstd compression wrapper.
-fn write_impl(
+fn write_impl<'h>(
     writer: impl Write,
-    package: &Package,
+    hugrs: impl IntoIterator<Item = &'h Hugr>,
+    extensions: &ExtensionRegistry,
     config: EnvelopeConfig,
 ) -> Result<(), EnvelopeError> {
     match config.format {
         #[allow(deprecated)]
-        EnvelopeFormat::PackageJson => package.to_json_writer(writer)?,
+        EnvelopeFormat::PackageJson => package_json::to_json_writer(hugrs, extensions, writer)?,
         EnvelopeFormat::Model | EnvelopeFormat::ModelWithExtensions => {
-            encode_model(writer, package, config.format)?
+            encode_model(writer, hugrs, extensions, config.format)?
         }
     }
     Ok(())
 }
 
-fn encode_model(
+fn encode_model<'h>(
     mut writer: impl Write,
-    package: &Package,
+    hugrs: impl IntoIterator<Item = &'h Hugr>,
+    extensions: &ExtensionRegistry,
     format: EnvelopeFormat,
 ) -> Result<(), EnvelopeError> {
     use hugr_model::v0::{binary::write_to_writer, bumpalo::Bump};
@@ -292,11 +308,11 @@ fn encode_model(
     }
 
     let bump = Bump::default();
-    let model_package = export_package(package, &bump);
+    let model_package = export_package(hugrs, extensions, &bump);
     write_to_writer(&model_package, &mut writer)?;
 
     if format.append_extensions() {
-        serde_json::to_writer(writer, &package.extensions.iter().collect_vec())?;
+        serde_json::to_writer(writer, &extensions.iter().collect_vec())?;
     }
 
     Ok(())

hugr-core/src/envelope.rs

@@ -0,0 +1,119 @@
+//! Encoding / decoding of Package json, used in the `PackageJson` envelope format.
+use derive_more::{Display, Error, From};
+use itertools::Itertools;
+use std::io;
+
+use crate::extension::resolution::ExtensionResolutionError;
+use crate::extension::{ExtensionRegistry, PRELUDE_REGISTRY};
+use crate::hugr::ExtensionError;
+use crate::package::Package;
+use crate::{Extension, Hugr, HugrView};
+
+/// Read a Package in json format from an io reader.
+pub(super) fn from_json_reader(
+    reader: impl io::Read,
+    extension_registry: &ExtensionRegistry,
+) -> Result<Package, PackageEncodingError> {
+    let val: serde_json::Value = serde_json::from_reader(reader)?;
+
+    let PackageDeser {
+        modules,
+        extensions: pkg_extensions,
+    } = serde_json::from_value::<PackageDeser>(val.clone())?;
+    let mut modules = modules.into_iter().map(|h| h.0).collect_vec();
+
+    // TODO: We don't currently store transitive extension dependencies in the
+    // package's extensions. For example, if we use a `collections.list` const
+    // value but don't use anything `prelude` we would not include `prelude` in
+    // the package's extensions. But this would then fail when loading the
+    // extensions, as we _need_ the prelude to load the `collections.list` op
+    // definitions here.
+    //
+    // The current fix is to always include the prelude when decoding, but this
+    // only works for transitive `prelude` dependencies.
+    //
+    // Chains of custom extensions will cause this to fail.
+    let extension_registry = if PRELUDE_REGISTRY
+        .iter()
+        .any(|e| !extension_registry.contains(&e.name))
+    {
+        let mut reg_with_prelude = extension_registry.clone();
+        reg_with_prelude.extend(PRELUDE_REGISTRY.iter().cloned());
+        reg_with_prelude
+    } else {
+        extension_registry.clone()
+    };
+
+    let mut pkg_extensions = ExtensionRegistry::new_with_extension_resolution(
+        pkg_extensions,
+        &(&extension_registry).into(),
+    )?;
+
+    // Resolve the operations in the modules using the defined registries.
+    let mut combined_registry = extension_registry.clone();
+    combined_registry.extend(&pkg_extensions);
+
+    for module in &mut modules {
+        module.resolve_extension_defs(&combined_registry)?;
+        pkg_extensions.extend(module.extensions());
+    }
+
+    Ok(Package {
+        modules,
+        extensions: pkg_extensions,
+    })
+}
+
+/// Write the Package in json format into an io writer.
+pub(super) fn to_json_writer<'h>(
+    hugrs: impl IntoIterator<Item = &'h Hugr>,
+    extensions: &ExtensionRegistry,
+    writer: impl io::Write,
+) -> Result<(), PackageEncodingError> {
+    let pkg_ser = PackageSer {
+        modules: hugrs.into_iter().map(HugrSer).collect(),
+        extensions: extensions.iter().map(|e| e.as_ref()).collect(),
+    };
+    serde_json::to_writer(writer, &pkg_ser)?;
+    Ok(())
+}
+
+/// Error raised while loading a package.
+#[derive(Debug, Display, Error, From)]
+#[non_exhaustive]
+pub enum PackageEncodingError {
+    /// Error raised while parsing the package json.
+    JsonEncoding(serde_json::Error),
+    /// Error raised while reading from a file.
+    IOError(io::Error),
+    /// Could not resolve the extension needed to encode the hugr.
+    ExtensionResolution(ExtensionResolutionError),
+    /// Could not resolve the runtime extensions for the hugr.
+    RuntimeExtensionResolution(ExtensionError),
+}
+
+/// A private package structure implementing the serde traits.
+///
+/// We use this to avoid exposing a public implementation of Serialize/Deserialize,
+/// as the json definition is not stable, and should always be wrapped in an Envelope.
+#[derive(Debug, serde::Serialize)]
+struct PackageSer<'h> {
+    pub modules: Vec<HugrSer<'h>>,
+    pub extensions: Vec<&'h Extension>,
+}
+#[derive(Debug, serde::Serialize)]
+#[serde(transparent)]
+struct HugrSer<'h>(#[serde(serialize_with = "Hugr::serde_serialize")] pub &'h Hugr);
+
+/// A private package structure implementing the serde traits.
+///
+/// We use this to avoid exposing a public implementation of Serialize/Deserialize,
+/// as the json definition is not stable, and should always be wrapped in an Envelope.
+#[derive(Debug, serde::Deserialize)]
+struct PackageDeser {
+    pub modules: Vec<HugrDeser>,
+    pub extensions: Vec<Extension>,
+}
+#[derive(Debug, serde::Deserialize)]
+#[serde(transparent)]
+struct HugrDeser(#[serde(deserialize_with = "Hugr::serde_deserialize")] pub Hugr);