Packages and envelopes for hugr-model.

Author: zrho

hugr-core/src/envelope.rs

@@ -159,15 +159,6 @@ 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."
-    )]
-    #[from(ignore)]
-    MultipleHugrs {
-        /// The number of HUGRs in the package.
-        count: usize,
-    },
     /// JSON serialization error.
     SerdeError {
         /// The source error.
@@ -244,7 +235,7 @@ fn decode_model(
     extension_registry: &ExtensionRegistry,
     format: EnvelopeFormat,
 ) -> Result<Package, EnvelopeError> {
-    use crate::{import::import_hugr, Extension};
+    use crate::Extension;
     use hugr_model::v0::bumpalo::Bump;
 
     if format.model_version() != Some(0) {
@@ -255,7 +246,7 @@ fn decode_model(
     }
 
     let bump = Bump::default();
-    let module_list = hugr_model::v0::binary::read_from_reader(&mut stream, &bump)?;
+    let model_package = hugr_model::v0::binary::read_from_reader(&mut stream, &bump)?;
 
     let mut extension_registry = extension_registry.clone();
     if format.append_extensions() {
@@ -266,9 +257,7 @@ fn decode_model(
         }
     }
 
-    // TODO: Import multiple hugrs from the model?
-    let hugr = import_hugr(&module_list, &extension_registry)?;
-    Ok(Package::new([hugr])?)
+    Ok(Package::from_model(&model_package, &extension_registry)?)
 }
 
 /// Internal implementation of [`write_envelope`] to call with/without the zstd compression wrapper.
@@ -301,7 +290,6 @@ fn encode_model(
     package: &Package,
     format: EnvelopeFormat,
 ) -> Result<(), EnvelopeError> {
-    use crate::export::export_hugr;
     use hugr_model::v0::{binary::write_to_writer, bumpalo::Bump};
 
     if format.model_version() != Some(0) {
@@ -311,15 +299,9 @@ fn encode_model(
         });
     }
 
-    // TODO: Export multiple hugrs to the model?
-    if package.modules.len() != 1 {
-        return Err(EnvelopeError::MultipleHugrs {
-            count: package.modules.len(),
-        });
-    }
     let bump = Bump::default();
-    let module = export_hugr(&package.modules[0], &bump);
-    write_to_writer(&module, &mut writer)?;
+    let model_package = package.to_model(&bump);
+    write_to_writer(&model_package, &mut writer)?;
 
     if format.append_extensions() {
         serde_json::to_writer(writer, &package.extensions.iter().collect_vec())?;

hugr-core/src/export.rs

@@ -28,7 +28,7 @@ use petgraph::unionfind::UnionFind;
 use std::fmt::Write;
 
 /// Export a [`Hugr`] graph to its representation in the model.
-pub fn export_hugr<'a>(hugr: &'a Hugr, bump: &'a Bump) -> table::Module<'a> {
+pub(crate) fn export_hugr<'a>(hugr: &'a Hugr, bump: &'a Bump) -> table::Module<'a> {
     let mut ctx = Context::new(hugr, bump);
     ctx.export_root();
     ctx.module

hugr-core/src/hugr.rs

@@ -34,6 +34,13 @@ use crate::ops::{OpTag, OpTrait};
 pub use crate::ops::{OpType, DEFAULT_OPTYPE};
 use crate::{Direction, Node};
 
+#[cfg(feature = "model_unstable")]
+use crate::export::export_hugr;
+#[cfg(feature = "model_unstable")]
+use crate::import::{import_hugr, ImportError};
+#[cfg(feature = "model_unstable")]
+use hugr_model::v0 as model;
+
 /// The Hugr data structure.
 #[derive(Clone, Debug, PartialEq)]
 pub struct Hugr {
@@ -256,6 +263,21 @@ impl Hugr {
         self.extensions = used_extensions;
         Ok(())
     }
+
+    /// Export this module to the model representation.
+    #[cfg(feature = "model_unstable")]
+    pub fn to_model<'a>(&'a self, bump: &'a model::bumpalo::Bump) -> model::table::Module<'a> {
+        export_hugr(self, bump)
+    }
+
+    /// Import a module from the model representation.
+    #[cfg(feature = "model_unstable")]
+    pub fn from_model<'a>(
+        module: &'a model::table::Module<'a>,
+        extensions: &ExtensionRegistry,
+    ) -> Result<Self, ImportError> {
+        import_hugr(module, extensions)
+    }
 }
 
 /// Internal API for HUGRs, not intended for use by users.

hugr-core/src/import.rs

@@ -80,7 +80,7 @@ macro_rules! error_uninferred {
 }
 
 /// Import a `hugr` module from its model representation.
-pub fn import_hugr(
+pub(crate) fn import_hugr(
     module: &table::Module,
     extensions: &ExtensionRegistry,
 ) -> Result<Hugr, ImportError> {

hugr-core/src/package.rs

@@ -14,6 +14,11 @@ use crate::hugr::{ExtensionError, HugrView, ValidationError};
 use crate::ops::{FuncDefn, Module, NamedOp, OpTag, OpTrait, OpType};
 use crate::{Extension, Hugr};
 
+#[cfg(feature = "model_unstable")]
+use crate::import::ImportError;
+#[cfg(feature = "model_unstable")]
+use hugr_model::v0 as model;
+
 #[derive(Debug, Default, Clone, PartialEq, serde::Serialize, serde::Deserialize)]
 /// Package of module HUGRs.
 pub struct Package {
@@ -308,6 +313,34 @@ impl Package {
         #[allow(deprecated)]
         self.to_json_writer(writer)
     }
+
+    /// Export this package to the model representation.
+    #[cfg(feature = "model_unstable")]
+    pub fn to_model<'a>(&'a self, bump: &'a model::bumpalo::Bump) -> model::table::Package<'a> {
+        let modules = self
+            .modules
+            .iter()
+            .map(|module| module.to_model(bump))
+            .collect();
+        model::table::Package { modules }
+    }
+
+    /// Import a package from the model representation.
+    #[cfg(feature = "model_unstable")]
+    pub fn from_model<'a>(
+        package: &'a model::table::Package<'a>,
+        extensions: &ExtensionRegistry,
+    ) -> Result<Self, ImportError> {
+        let modules = package
+            .modules
+            .iter()
+            .map(|module| Hugr::from_model(module, extensions))
+            .collect::<Result<Vec<_>, _>>()?;
+
+        // This does not panic since the import already requires a module root.
+        let package = Self::new(modules).expect("non-module root");
+        Ok(package)
+    }
 }
 
 impl AsRef<[Hugr]> for Package {

hugr-core/tests/model.rs

@@ -2,16 +2,15 @@
 
 use std::str::FromStr;
 
-use hugr::std_extensions::std_reg;
-use hugr_core::{export::export_hugr, import::import_hugr};
+use hugr::{package::Package, std_extensions::std_reg};
 use hugr_model::v0 as model;
 
 fn roundtrip(source: &str) -> String {
     let bump = model::bumpalo::Bump::new();
-    let module_ast = model::ast::Module::from_str(source).unwrap();
-    let module_table = module_ast.resolve(&bump).unwrap();
-    let hugr = import_hugr(&module_table, &std_reg()).unwrap();
-    let exported_table = export_hugr(&hugr, &bump);
+    let package_ast = model::ast::Package::from_str(source).unwrap();
+    let package_table = package_ast.resolve(&bump).unwrap();
+    let core = Package::from_model(&package_table, &std_reg()).unwrap();
+    let exported_table = core.to_model(&bump);
     let exported_ast = exported_table.as_ast().unwrap();
     exported_ast.to_string()
 }

hugr-core/tests/snapshots/model__roundtrip_add.snap

@@ -4,6 +4,8 @@ expression: "roundtrip(include_str!(\"../../hugr-model/tests/fixtures/model-add.
 ---
 (hugr 0)
 
+(mod)
+
 (import core.fn)
 
 (import arithmetic.int.iadd)

hugr-core/tests/snapshots/model__roundtrip_alias.snap

@@ -4,6 +4,8 @@ expression: "roundtrip(include_str!(\"../../hugr-model/tests/fixtures/model-alia
 ---
 (hugr 0)
 
+(mod)
+
 (import core.fn)
 
 (import arithmetic.int.types.int)

hugr-core/tests/snapshots/model__roundtrip_call.snap

@@ -4,6 +4,8 @@ expression: "roundtrip(include_str!(\"../../hugr-model/tests/fixtures/model-call
 ---
 (hugr 0)
 
+(mod)
+
 (import core.call)
 
 (import core.load_const)

hugr-core/tests/snapshots/model__roundtrip_cfg.snap

@@ -4,6 +4,8 @@ expression: "roundtrip(include_str!(\"../../hugr-model/tests/fixtures/model-cfg.
 ---
 (hugr 0)
 
+(mod)
+
 (import core.make_adt)
 
 (import core.ctrl)