feat: Packages in hugr-model and envelope support. (#2026)

This PR adds packages to hugr-model with support for envelopes.

 - Packages for both the AST and table representation of hugr-model.
 - Packages in the Python model classes.
 - Reading and writing model envelopes in Rust.
 - Writing model envelopes in Python.

Author: zrho

hugr-core/src/envelope.rs

@@ -34,6 +34,12 @@
 //! - Bit 7,6: Constant "01" to make some headers ascii-printable.
 //!
 
+#![allow(deprecated)]
+// TODO: Due to a bug in `derive_more`
+// (https://github.com/JelteF/derive_more/issues/419) we need to deactivate
+// deprecation warnings here. We can reactivate them once the bug is fixed by
+// https://github.com/JelteF/derive_more/pull/454.
+
 mod header;
 
 pub use header::{EnvelopeConfig, EnvelopeFormat, ZstdConfig, MAGIC_NUMBERS};
@@ -161,9 +167,11 @@ pub enum EnvelopeError {
     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."
-    )]
+            "Packages with multiple HUGRs are currently unsupported. Tried to encode {count} HUGRs, when 1 was expected."
+        )]
     #[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 {
         /// The number of HUGRs in the package.
         count: usize,
@@ -244,7 +252,7 @@ fn decode_model(
     extension_registry: &ExtensionRegistry,
     format: EnvelopeFormat,
 ) -> Result<Package, EnvelopeError> {
-    use crate::{import::import_hugr, Extension};
+    use crate::{import::import_package, Extension};
     use hugr_model::v0::bumpalo::Bump;
 
     if format.model_version() != Some(0) {
@@ -255,7 +263,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 +274,7 @@ fn decode_model(
         }
     }
 
-    // TODO: Import multiple hugrs from the model?
-    let hugr = import_hugr(&module_list, &extension_registry)?;
-    Ok(Package::new([hugr])?)
+    Ok(import_package(&model_package, &extension_registry)?)
 }
 
 /// Internal implementation of [`write_envelope`] to call with/without the zstd compression wrapper.
@@ -301,25 +307,20 @@ fn encode_model(
     package: &Package,
     format: EnvelopeFormat,
 ) -> Result<(), EnvelopeError> {
-    use crate::export::export_hugr;
     use hugr_model::v0::{binary::write_to_writer, bumpalo::Bump};
 
+    use crate::export::export_package;
+
     if format.model_version() != Some(0) {
         return Err(EnvelopeError::FormatUnsupported {
             format,
             feature: None,
         });
     }
 
-    // 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 = export_package(package, &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

@@ -5,6 +5,7 @@ use crate::{
     ops::{
         constant::CustomSerialized, DataflowBlock, DataflowOpTrait, OpName, OpTrait, OpType, Value,
     },
+    package::Package,
     std_extensions::{
         arithmetic::{float_types::ConstF64, int_types::ConstInt},
         collections::array::ArrayValue,
@@ -27,6 +28,16 @@ use hugr_model::v0::{
 use petgraph::unionfind::UnionFind;
 use std::fmt::Write;
 
+/// Export a [`Package`] to its representation in the model.
+pub fn export_package<'a>(package: &'a Package, bump: &'a Bump) -> table::Package<'a> {
+    let modules = package
+        .modules
+        .iter()
+        .map(|module| export_hugr(module, bump))
+        .collect();
+    table::Package { modules }
+}
+
 /// Export a [`Hugr`] graph to its representation in the model.
 pub fn export_hugr<'a>(hugr: &'a Hugr, bump: &'a Bump) -> table::Module<'a> {
     let mut ctx = Context::new(hugr, bump);

hugr-core/src/import.rs

@@ -14,6 +14,7 @@ use crate::{
         ExitBlock, FuncDecl, FuncDefn, Input, LoadConstant, LoadFunction, Module, OpType, OpaqueOp,
         Output, Tag, TailLoop, Value, CFG, DFG,
     },
+    package::Package,
     std_extensions::{
         arithmetic::{float_types::ConstF64, int_types::ConstInt},
         collections::array::ArrayValue,
@@ -79,7 +80,23 @@ macro_rules! error_uninferred {
     ($($e:expr),*) => { ImportError::Uninferred(format!($($e),*)) }
 }
 
-/// Import a `hugr` module from its model representation.
+/// Import a [`Package`] from its model representation.
+pub fn import_package(
+    package: &table::Package,
+    extensions: &ExtensionRegistry,
+) -> Result<Package, ImportError> {
+    let modules = package
+        .modules
+        .iter()
+        .map(|module| import_hugr(module, extensions))
+        .collect::<Result<Vec<_>, _>>()?;
+
+    // This does not panic since the import already requires a module root.
+    let package = Package::new(modules).expect("non-module root");
+    Ok(package)
+}
+
+/// Import a [`Hugr`] module from its model representation.
 pub fn import_hugr(
     module: &table::Module,
     extensions: &ExtensionRegistry,

hugr-core/tests/model.rs

@@ -3,15 +3,15 @@
 use std::str::FromStr;
 
 use hugr::std_extensions::std_reg;
-use hugr_core::{export::export_hugr, import::import_hugr};
+use hugr_core::{export::export_package, import::import_package};
 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 = import_package(&package_table, &std_reg()).unwrap();
+    let exported_table = export_package(&core, &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)

hugr-core/tests/snapshots/model__roundtrip_cond.snap

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

hugr-core/tests/snapshots/model__roundtrip_const.snap

@@ -4,6 +4,8 @@ expression: "roundtrip(include_str!(\"../../hugr-model/tests/fixtures/model-cons
 ---
 (hugr 0)
 
+(mod)
+
 (import collections.array.array)
 
 (import compat.const_json)

hugr-core/tests/snapshots/model__roundtrip_constraints.snap

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

hugr-core/tests/snapshots/model__roundtrip_loop.snap

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

hugr-core/tests/snapshots/model__roundtrip_params.snap

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

hugr-model/capnp/hugr-v0.capnp

@@ -18,6 +18,10 @@ using LinkId = UInt32;
 # The index of a `Link`.
 using LinkIndex = UInt32;
 
+struct Package {
+    modules @0 :List(Module);
+}
+
 struct Module {
     root @0 :RegionId;
     nodes @1 :List(Node);