From 95c77f9d8243ba799062ee21fa64ced988eacf6c Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Maximilian=20R=C3=BCsch?= Date: Mon, 9 Mar 2026 13:55:26 +0000 Subject: [PATCH 01/16] hugr-py linking --- hugr-py/rust/lib.rs | 3 ++ hugr-py/rust/linking.rs | 44 ++++++++++++++++++++++++++++++ hugr-py/src/hugr/_hugr/linking.pyi | 1 + hugr-py/src/hugr/package.py | 27 ++++++++++++++++++ uv.lock | 2 +- 5 files changed, 76 insertions(+), 1 deletion(-) create mode 100644 hugr-py/rust/linking.rs create mode 100644 hugr-py/src/hugr/_hugr/linking.pyi diff --git a/hugr-py/rust/lib.rs b/hugr-py/rust/lib.rs index 3b65ec8dc..3606f10b9 100644 --- a/hugr-py/rust/lib.rs +++ b/hugr-py/rust/lib.rs @@ -1,5 +1,6 @@ //! Supporting Rust library for the hugr Python bindings. +mod linking; mod metadata; mod model; mod zstd_util; @@ -8,6 +9,8 @@ use pyo3::pymodule; #[pymodule] mod _hugr { + #[pymodule_export] + use super::linking::linking; #[pymodule_export] use super::metadata::metadata; #[pymodule_export] diff --git a/hugr-py/rust/linking.rs b/hugr-py/rust/linking.rs new file mode 100644 index 000000000..cbf84f33f --- /dev/null +++ b/hugr-py/rust/linking.rs @@ -0,0 +1,44 @@ +//! Bindings for linking utilities defined in the hugr-core crate + +use pyo3::pymodule; + +#[pymodule(submodule)] +#[pyo3(module = "hugr._hugr.linking")] +pub mod linking { + /// Hack: workaround for + #[pymodule_init] + fn init(m: &Bound<'_, PyModule>) -> PyResult<()> { + Python::attach(|py| { + py.import("sys")? + .getattr("modules")? + .set_item("hugr._hugr.linking", m) + }) + } + + use hugr_core; + use hugr_core::envelope::EnvelopeConfig; + use hugr_core::hugr::linking::{HugrLinking, NameLinkingPolicy}; + use hugr_core::std_extensions::std_reg; + use pyo3::exceptions::{PyRuntimeError, PyValueError}; + use pyo3::types::{PyAnyMethods, PyModule}; + use pyo3::{Bound, PyResult, Python, pyfunction}; + + #[pyfunction] + fn link_modules(module_into: Option<&[u8]>, module_from: Option<&[u8]>) -> PyResult> { + let reg = std_reg(); + let mut hugr_into = hugr_core::Hugr::load(module_into.unwrap(), Some(®)) + .map_err(|err| PyValueError::new_err(err.to_string()))?; + let hugr_from = hugr_core::Hugr::load(module_from.unwrap(), Some(®)) + .map_err(|err| PyValueError::new_err(err.to_string()))?; + + hugr_into + .link_module(hugr_from, &NameLinkingPolicy::default()) + .map_err(|err| PyRuntimeError::new_err(err.to_string()))?; // TODO Add a proper error + + let mut result = Vec::new(); + hugr_into + .store(&mut result, EnvelopeConfig::default()) + .map_err(|err| PyRuntimeError::new_err(err.to_string()))?; // TODO Add a proper error + Ok(result) + } +} diff --git a/hugr-py/src/hugr/_hugr/linking.pyi b/hugr-py/src/hugr/_hugr/linking.pyi new file mode 100644 index 000000000..885f78859 --- /dev/null +++ b/hugr-py/src/hugr/_hugr/linking.pyi @@ -0,0 +1 @@ +def link_modules(module_into: bytes, module_from: bytes) -> bytes: ... diff --git a/hugr-py/src/hugr/package.py b/hugr-py/src/hugr/package.py index 7c392a5f8..9ea642e64 100644 --- a/hugr-py/src/hugr/package.py +++ b/hugr-py/src/hugr/package.py @@ -10,6 +10,7 @@ import hugr._serialization.extension as ext_s import hugr.model as model from hugr import ext +from hugr._hugr.linking import link_modules from hugr.envelope import ( EnvelopeConfig, _make_envelope, @@ -188,6 +189,32 @@ def used_extensions( return result + def link(self, *other: Package): + """Link this package with other packages, returning a new package containing the + extensions of all packages, as well as a single module created from linking the + modules from all packages. + + Args: + *other: Other packages to link with. + + Returns: + A new package containing the modules and extensions of all packages. + """ + modules = self.modules[:] + extensions = self.extensions[:] + for pkg in other: + modules.extend(pkg.modules) + extensions.extend(pkg.extensions) + + if len(modules) == 0: + return Package([], extensions) + + result_module_bytes = modules[0].to_bytes() + for module in modules[1:]: + result_module_bytes = link_modules(result_module_bytes, module.to_bytes()) + + return Package([Hugr.from_bytes(result_module_bytes)], extensions) + @dataclass(frozen=True) class PackagePointer: diff --git a/uv.lock b/uv.lock index 55a44a0b1..88323a0e3 100644 --- a/uv.lock +++ b/uv.lock @@ -399,7 +399,7 @@ requires-dist = [ { name = "semver", specifier = "~=3.0" }, { name = "sphinx", marker = "extra == 'docs'", specifier = ">=8.1.3,<10" }, { name = "sphinx-favicon", marker = "extra == 'docs'", specifier = "~=1.1.0" }, - { name = "sphinxcontrib-mermaid", marker = "extra == 'docs'", specifier = "~=1.0" }, + { name = "sphinxcontrib-mermaid", marker = "extra == 'docs'", specifier = ">=1,<3" }, { name = "typing-extensions", specifier = "~=4.12" }, ] provides-extras = ["docs", "pytket"] From da1482bc6d6a5c0c2ef6080e2c57f594b181a5cd Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Maximilian=20R=C3=BCsch?= Date: Mon, 9 Mar 2026 16:27:26 +0000 Subject: [PATCH 02/16] Pitchit hard --- hugr-py/rust/linking.rs | 50 ++++++++++++++++++++++++----------------- 1 file changed, 30 insertions(+), 20 deletions(-) diff --git a/hugr-py/rust/linking.rs b/hugr-py/rust/linking.rs index cbf84f33f..bec567ad0 100644 --- a/hugr-py/rust/linking.rs +++ b/hugr-py/rust/linking.rs @@ -5,6 +5,13 @@ use pyo3::pymodule; #[pymodule(submodule)] #[pyo3(module = "hugr._hugr.linking")] pub mod linking { + use hugr_core; + use hugr_core::envelope::EnvelopeConfig; + use hugr_core::hugr::linking::{HugrLinking, NameLinkingPolicy}; + use pyo3::exceptions::{PyRuntimeError, PyValueError}; + use pyo3::types::{PyAnyMethods, PyModule}; + use pyo3::{Bound, PyResult, Python, pyfunction}; + /// Hack: workaround for #[pymodule_init] fn init(m: &Bound<'_, PyModule>) -> PyResult<()> { @@ -15,30 +22,33 @@ pub mod linking { }) } - use hugr_core; - use hugr_core::envelope::EnvelopeConfig; - use hugr_core::hugr::linking::{HugrLinking, NameLinkingPolicy}; - use hugr_core::std_extensions::std_reg; - use pyo3::exceptions::{PyRuntimeError, PyValueError}; - use pyo3::types::{PyAnyMethods, PyModule}; - use pyo3::{Bound, PyResult, Python, pyfunction}; - #[pyfunction] - fn link_modules(module_into: Option<&[u8]>, module_from: Option<&[u8]>) -> PyResult> { - let reg = std_reg(); - let mut hugr_into = hugr_core::Hugr::load(module_into.unwrap(), Some(®)) - .map_err(|err| PyValueError::new_err(err.to_string()))?; - let hugr_from = hugr_core::Hugr::load(module_from.unwrap(), Some(®)) - .map_err(|err| PyValueError::new_err(err.to_string()))?; - - hugr_into - .link_module(hugr_from, &NameLinkingPolicy::default()) + fn link_modules(module_into: &[u8], module_from: &[u8]) -> PyResult> { + let mut pkg_into = hugr_core::package::Package::load(module_into, None).map_err(|err| { + PyValueError::new_err(format!("Loading of first envelope failed: {}", err)) + })?; // TODO need a combination of loading a package, and expecting a single hugr in it. + + let pkg_from = hugr_core::package::Package::load(module_from, None).map_err(|err| { + PyValueError::new_err(format!("Loading of second envelope failed: {}", err)) + })?; + + pkg_into.modules[0] + .link_module( + pkg_from.modules.into_iter().next().unwrap(), + &NameLinkingPolicy::default(), + ) .map_err(|err| PyRuntimeError::new_err(err.to_string()))?; // TODO Add a proper error + pkg_into.extensions.extend(pkg_from.extensions); + let mut result = Vec::new(); - hugr_into - .store(&mut result, EnvelopeConfig::default()) - .map_err(|err| PyRuntimeError::new_err(err.to_string()))?; // TODO Add a proper error + pkg_into + .store(&mut result, EnvelopeConfig::binary()) + .unwrap(); + + hugr_core::package::Package::load(&result[..], None) + .map_err(|err| PyValueError::new_err(format!("Roundtrip failed: {:?}", err)))?; + Ok(result) } } From 4a8a4a9f57cf1749514e56f1d1a38aab6c1dad3a Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Maximilian=20R=C3=BCsch?= Date: Mon, 9 Mar 2026 17:14:58 +0000 Subject: [PATCH 03/16] Better linking --- hugr-core/src/hugr.rs | 18 ++++++++++++++++++ hugr-py/rust/linking.rs | 24 +++++++++++------------- 2 files changed, 29 insertions(+), 13 deletions(-) diff --git a/hugr-core/src/hugr.rs b/hugr-core/src/hugr.rs index f89940cc4..ac38caa69 100644 --- a/hugr-core/src/hugr.rs +++ b/hugr-core/src/hugr.rs @@ -166,6 +166,24 @@ impl Hugr { } } + /// Read a HUGR from an Envelope, and return the enclosed extensions. + /// + /// To load a HUGR, all the extensions used in its definition must be + /// available. The Envelope may include some of the extensions, but any + /// additional extensions must be provided in the `extensions` parameter. If + /// `extensions` is `None`, the default [`crate::std_extensions::STD_REG`] + /// is used. + pub fn load_with_exts( + reader: impl io::BufRead, + extensions: Option<&ExtensionRegistry>, + ) -> Result<(Self, ExtensionRegistry), ReadError> { + let pkg = Package::load(reader, extensions)?; + match pkg.modules.into_iter().exactly_one() { + Ok(hugr) => Ok((hugr, pkg.extensions)), + Err(e) => Err(ReadError::ExpectedSingleHugr { count: e.count() }), + } + } + /// Read a HUGR from an Envelope encoded in a string. /// /// Note that not all Envelopes are valid strings. In the general case, diff --git a/hugr-py/rust/linking.rs b/hugr-py/rust/linking.rs index bec567ad0..a0c69808c 100644 --- a/hugr-py/rust/linking.rs +++ b/hugr-py/rust/linking.rs @@ -5,7 +5,7 @@ use pyo3::pymodule; #[pymodule(submodule)] #[pyo3(module = "hugr._hugr.linking")] pub mod linking { - use hugr_core; + use hugr_core::Hugr; use hugr_core::envelope::EnvelopeConfig; use hugr_core::hugr::linking::{HugrLinking, NameLinkingPolicy}; use pyo3::exceptions::{PyRuntimeError, PyValueError}; @@ -24,26 +24,24 @@ pub mod linking { #[pyfunction] fn link_modules(module_into: &[u8], module_from: &[u8]) -> PyResult> { - let mut pkg_into = hugr_core::package::Package::load(module_into, None).map_err(|err| { - PyValueError::new_err(format!("Loading of first envelope failed: {}", err)) - })?; // TODO need a combination of loading a package, and expecting a single hugr in it. + let (mut hugr_into, mut exts_into) = + Hugr::load_with_exts(module_into, None).map_err(|err| { + PyValueError::new_err(format!("Loading of first envelope failed: {}", err)) + })?; - let pkg_from = hugr_core::package::Package::load(module_from, None).map_err(|err| { + let (hugr_from, exts_from) = Hugr::load_with_exts(module_from, None).map_err(|err| { PyValueError::new_err(format!("Loading of second envelope failed: {}", err)) })?; - pkg_into.modules[0] - .link_module( - pkg_from.modules.into_iter().next().unwrap(), - &NameLinkingPolicy::default(), - ) + hugr_into + .link_module(hugr_from, &NameLinkingPolicy::default()) .map_err(|err| PyRuntimeError::new_err(err.to_string()))?; // TODO Add a proper error - pkg_into.extensions.extend(pkg_from.extensions); + exts_into.extend(exts_from); let mut result = Vec::new(); - pkg_into - .store(&mut result, EnvelopeConfig::binary()) + hugr_into + .store_with_exts(&mut result, EnvelopeConfig::binary(), &exts_into) .unwrap(); hugr_core::package::Package::load(&result[..], None) From 87cd79f48bbe1223151f0069e5fef803e3ca33c1 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Maximilian=20R=C3=BCsch?= Date: Mon, 9 Mar 2026 17:20:19 +0000 Subject: [PATCH 04/16] Add proper errorrage --- hugr-py/rust/linking.rs | 14 +++++++++++--- hugr-py/src/hugr/_hugr/linking.pyi | 3 +++ 2 files changed, 14 insertions(+), 3 deletions(-) diff --git a/hugr-py/rust/linking.rs b/hugr-py/rust/linking.rs index a0c69808c..6584d3027 100644 --- a/hugr-py/rust/linking.rs +++ b/hugr-py/rust/linking.rs @@ -1,6 +1,7 @@ //! Bindings for linking utilities defined in the hugr-core crate -use pyo3::pymodule; +use pyo3::exceptions::PyException; +use pyo3::{create_exception, pymodule}; #[pymodule(submodule)] #[pyo3(module = "hugr._hugr.linking")] @@ -8,7 +9,7 @@ pub mod linking { use hugr_core::Hugr; use hugr_core::envelope::EnvelopeConfig; use hugr_core::hugr::linking::{HugrLinking, NameLinkingPolicy}; - use pyo3::exceptions::{PyRuntimeError, PyValueError}; + use pyo3::exceptions::PyValueError; use pyo3::types::{PyAnyMethods, PyModule}; use pyo3::{Bound, PyResult, Python, pyfunction}; @@ -35,7 +36,7 @@ pub mod linking { hugr_into .link_module(hugr_from, &NameLinkingPolicy::default()) - .map_err(|err| PyRuntimeError::new_err(err.to_string()))?; // TODO Add a proper error + .map_err(|err| super::HugrLinkingError::new_err(err.to_string()))?; exts_into.extend(exts_from); @@ -50,3 +51,10 @@ pub mod linking { Ok(result) } } + +create_exception!( + _hugr.linking, + HugrLinkingError, + PyException, + "Base exception for HUGR linking errors." +); diff --git a/hugr-py/src/hugr/_hugr/linking.pyi b/hugr-py/src/hugr/_hugr/linking.pyi index 885f78859..ebf7931cc 100644 --- a/hugr-py/src/hugr/_hugr/linking.pyi +++ b/hugr-py/src/hugr/_hugr/linking.pyi @@ -1 +1,4 @@ +class HugrLinkingError(Exception): ... + + def link_modules(module_into: bytes, module_from: bytes) -> bytes: ... From c519d207d20aaab199e5502dfed18d65d14e8e83 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Maximilian=20R=C3=BCsch?= Date: Mon, 9 Mar 2026 17:20:56 +0000 Subject: [PATCH 05/16] Formatting --- hugr-py/rust/linking.rs | 1 - 1 file changed, 1 deletion(-) diff --git a/hugr-py/rust/linking.rs b/hugr-py/rust/linking.rs index 6584d3027..c269bc915 100644 --- a/hugr-py/rust/linking.rs +++ b/hugr-py/rust/linking.rs @@ -29,7 +29,6 @@ pub mod linking { Hugr::load_with_exts(module_into, None).map_err(|err| { PyValueError::new_err(format!("Loading of first envelope failed: {}", err)) })?; - let (hugr_from, exts_from) = Hugr::load_with_exts(module_from, None).map_err(|err| { PyValueError::new_err(format!("Loading of second envelope failed: {}", err)) })?; From 136ce79bb4b2d8520e3f739d8f6aeeaaf3d616c2 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Maximilian=20R=C3=BCsch?= Date: Mon, 16 Mar 2026 11:11:47 +0000 Subject: [PATCH 06/16] Add entrypoint support --- hugr-py/rust/linking.rs | 20 +++++++++++++++++--- 1 file changed, 17 insertions(+), 3 deletions(-) diff --git a/hugr-py/rust/linking.rs b/hugr-py/rust/linking.rs index c269bc915..b7c585bd5 100644 --- a/hugr-py/rust/linking.rs +++ b/hugr-py/rust/linking.rs @@ -6,9 +6,10 @@ use pyo3::{create_exception, pymodule}; #[pymodule(submodule)] #[pyo3(module = "hugr._hugr.linking")] pub mod linking { - use hugr_core::Hugr; use hugr_core::envelope::EnvelopeConfig; + use hugr_core::hugr::hugrmut::HugrMut; use hugr_core::hugr::linking::{HugrLinking, NameLinkingPolicy}; + use hugr_core::{Hugr, HugrView}; use pyo3::exceptions::PyValueError; use pyo3::types::{PyAnyMethods, PyModule}; use pyo3::{Bound, PyResult, Python, pyfunction}; @@ -32,11 +33,24 @@ pub mod linking { let (hugr_from, exts_from) = Hugr::load_with_exts(module_from, None).map_err(|err| { PyValueError::new_err(format!("Loading of second envelope failed: {}", err)) })?; + let into_executable = hugr_into.entrypoint() != hugr_into.module_root(); + let from_executable = hugr_from.entrypoint() != hugr_from.module_root(); + let replacement_entrypoint = if into_executable && from_executable { + return Err(PyValueError::new_err( + "Cannot link two executable modules together.", + )); + } else if !into_executable && from_executable { + Some(hugr_from.entrypoint()) + } else { + None + }; - hugr_into + let forest = hugr_into .link_module(hugr_from, &NameLinkingPolicy::default()) .map_err(|err| super::HugrLinkingError::new_err(err.to_string()))?; - + if let Some(new_entrypoint) = replacement_entrypoint { + hugr_into.set_entrypoint(*forest.node_map.get(&new_entrypoint).unwrap()); + } exts_into.extend(exts_from); let mut result = Vec::new(); From 58697fcb068ceed0d9ea5c4de5ad5366cc6ee1b0 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Maximilian=20R=C3=BCsch?= Date: Mon, 16 Mar 2026 12:28:00 +0000 Subject: [PATCH 07/16] Better panic --- hugr-py/rust/linking.rs | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/hugr-py/rust/linking.rs b/hugr-py/rust/linking.rs index b7c585bd5..0c81297f3 100644 --- a/hugr-py/rust/linking.rs +++ b/hugr-py/rust/linking.rs @@ -49,7 +49,10 @@ pub mod linking { .link_module(hugr_from, &NameLinkingPolicy::default()) .map_err(|err| super::HugrLinkingError::new_err(err.to_string()))?; if let Some(new_entrypoint) = replacement_entrypoint { - hugr_into.set_entrypoint(*forest.node_map.get(&new_entrypoint).unwrap()); + let Some(node) = forest.node_map.get(&new_entrypoint) else { + panic!("Entrypoint is to be replaced but was not found after linking"); + }; + hugr_into.set_entrypoint(*node); } exts_into.extend(exts_from); From 2a1d7bc3a80e4cc86d681427ce5d81e13108abf0 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Maximilian=20R=C3=BCsch?= Date: Mon, 16 Mar 2026 12:29:40 +0000 Subject: [PATCH 08/16] Begin linking tests --- hugr-py/src/hugr/_hugr/linking.pyi | 1 - hugr-py/tests/test_linking.py | 56 ++++++++++++++++++++++++++++++ 2 files changed, 56 insertions(+), 1 deletion(-) create mode 100644 hugr-py/tests/test_linking.py diff --git a/hugr-py/src/hugr/_hugr/linking.pyi b/hugr-py/src/hugr/_hugr/linking.pyi index ebf7931cc..e56eb69a3 100644 --- a/hugr-py/src/hugr/_hugr/linking.pyi +++ b/hugr-py/src/hugr/_hugr/linking.pyi @@ -1,4 +1,3 @@ class HugrLinkingError(Exception): ... - def link_modules(module_into: bytes, module_from: bytes) -> bytes: ... diff --git a/hugr-py/tests/test_linking.py b/hugr-py/tests/test_linking.py new file mode 100644 index 000000000..231a6be51 --- /dev/null +++ b/hugr-py/tests/test_linking.py @@ -0,0 +1,56 @@ +import pytest + +from hugr import Hugr, tys +from hugr._hugr.linking import link_modules +from hugr.build import Module +from hugr.ops import FuncDefn + + +def build_module(*, entrypoint: bool) -> Hugr: + builder = Module() + if entrypoint: + main = builder.define_function( + "main", input_types=[tys.Bool], output_types=[tys.Bool], visibility="Public" + ) + main.set_outputs(*main.inputs()) + builder.hugr.entrypoint = main.parent_node + + return builder.hugr + + +def test_no_entrypoints(): + hugr1 = build_module(entrypoint=False) + hugr2 = build_module(entrypoint=False) + + linked = Hugr.from_bytes(link_modules(hugr1.to_bytes(), hugr2.to_bytes())) + assert linked.entrypoint == linked.module_root + + +def test_entrypoint_lhs(): + hugr1 = build_module(entrypoint=True) + hugr2 = build_module(entrypoint=False) + + linked = Hugr.from_bytes(link_modules(hugr1.to_bytes(), hugr2.to_bytes())) + assert linked.entrypoint != linked.module_root + entrypoint = linked.entrypoint_op() + assert isinstance(entrypoint, FuncDefn) + assert entrypoint.f_name == "main" + + +def test_entrypoint_rhs(): + hugr1 = build_module(entrypoint=False) + hugr2 = build_module(entrypoint=True) + + linked = Hugr.from_bytes(link_modules(hugr1.to_bytes(), hugr2.to_bytes())) + assert linked.entrypoint != linked.module_root + entrypoint = linked.entrypoint_op() + assert isinstance(entrypoint, FuncDefn) + assert entrypoint.f_name == "main" + + +def test_multiple_entrypoints(): + hugr1 = build_module(entrypoint=True) + hugr2 = build_module(entrypoint=True) + + with pytest.raises(ValueError, match="Cannot link two executable modules together"): + link_modules(hugr1.to_bytes(), hugr2.to_bytes()) From 76ada5629d013c7555b369520d1ab51165358138 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Maximilian=20R=C3=BCsch?= Date: Mon, 16 Mar 2026 14:39:03 +0000 Subject: [PATCH 09/16] More linking tests and fixes --- hugr-py/src/hugr/package.py | 4 ++- hugr-py/tests/test_linking.py | 46 ++++++++++++++++++++++++++++++++--- 2 files changed, 45 insertions(+), 5 deletions(-) diff --git a/hugr-py/src/hugr/package.py b/hugr-py/src/hugr/package.py index 9ea642e64..c1f31edac 100644 --- a/hugr-py/src/hugr/package.py +++ b/hugr-py/src/hugr/package.py @@ -204,7 +204,9 @@ def link(self, *other: Package): extensions = self.extensions[:] for pkg in other: modules.extend(pkg.modules) - extensions.extend(pkg.extensions) + for new_ext in pkg.extensions: + if new_ext not in extensions: + extensions.append(new_ext) if len(modules) == 0: return Package([], extensions) diff --git a/hugr-py/tests/test_linking.py b/hugr-py/tests/test_linking.py index 231a6be51..525756c99 100644 --- a/hugr-py/tests/test_linking.py +++ b/hugr-py/tests/test_linking.py @@ -4,6 +4,8 @@ from hugr._hugr.linking import link_modules from hugr.build import Module from hugr.ops import FuncDefn +from hugr.package import Package +from hugr.std import float, int, logic, ptr def build_module(*, entrypoint: bool) -> Hugr: @@ -18,7 +20,7 @@ def build_module(*, entrypoint: bool) -> Hugr: return builder.hugr -def test_no_entrypoints(): +def test_link_modules_no_entrypoints(): hugr1 = build_module(entrypoint=False) hugr2 = build_module(entrypoint=False) @@ -26,7 +28,7 @@ def test_no_entrypoints(): assert linked.entrypoint == linked.module_root -def test_entrypoint_lhs(): +def test_link_modules_entrypoint_lhs(): hugr1 = build_module(entrypoint=True) hugr2 = build_module(entrypoint=False) @@ -37,7 +39,7 @@ def test_entrypoint_lhs(): assert entrypoint.f_name == "main" -def test_entrypoint_rhs(): +def test_link_modules_entrypoint_rhs(): hugr1 = build_module(entrypoint=False) hugr2 = build_module(entrypoint=True) @@ -48,9 +50,45 @@ def test_entrypoint_rhs(): assert entrypoint.f_name == "main" -def test_multiple_entrypoints(): +def test_link_modules_multiple_entrypoints(): hugr1 = build_module(entrypoint=True) hugr2 = build_module(entrypoint=True) with pytest.raises(ValueError, match="Cannot link two executable modules together"): link_modules(hugr1.to_bytes(), hugr2.to_bytes()) + + +def test_link_packages_extensions(): + pkg1 = Package( + modules=[build_module(entrypoint=False)], + extensions=[ + int.CONVERSIONS_EXTENSION, + int.INT_TYPES_EXTENSION, + int.INT_OPS_EXTENSION, + # Shared + logic.EXTENSION, + ptr.EXTENSION, + ], + ) + pkg2 = Package( + modules=[build_module(entrypoint=False)], + extensions=[ + float.FLOAT_OPS_EXTENSION, + float.FLOAT_TYPES_EXTENSION, + # Shared + logic.EXTENSION, + ptr.EXTENSION, + ], + ) + + result_pkg = pkg1.link(pkg2) + + assert result_pkg.extensions == [ + int.CONVERSIONS_EXTENSION, + int.INT_TYPES_EXTENSION, + int.INT_OPS_EXTENSION, + logic.EXTENSION, + ptr.EXTENSION, + float.FLOAT_OPS_EXTENSION, + float.FLOAT_TYPES_EXTENSION, + ] From 15574aa50e279f1a6f2d9eabf0a52e890175249a Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Maximilian=20R=C3=BCsch?= Date: Mon, 16 Mar 2026 14:56:43 +0000 Subject: [PATCH 10/16] Empti test --- hugr-py/tests/test_linking.py | 9 +++++++++ 1 file changed, 9 insertions(+) diff --git a/hugr-py/tests/test_linking.py b/hugr-py/tests/test_linking.py index 525756c99..6c5ea4e0b 100644 --- a/hugr-py/tests/test_linking.py +++ b/hugr-py/tests/test_linking.py @@ -58,6 +58,15 @@ def test_link_modules_multiple_entrypoints(): link_modules(hugr1.to_bytes(), hugr2.to_bytes()) +def test_link_packages_no_modules(): + pkg1 = Package(modules=[]) + pkg2 = Package(modules=[]) + + result_pkg = pkg1.link(pkg2) + + assert result_pkg.modules == [] + + def test_link_packages_extensions(): pkg1 = Package( modules=[build_module(entrypoint=False)], From cb6582eb7fbc6e163f770bb10bf1e41d7c386aae Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Maximilian=20R=C3=BCsch?= Date: Tue, 17 Mar 2026 12:08:41 +0000 Subject: [PATCH 11/16] Use owned param --- hugr-py/rust/linking.rs | 13 +++++++------ 1 file changed, 7 insertions(+), 6 deletions(-) diff --git a/hugr-py/rust/linking.rs b/hugr-py/rust/linking.rs index 0c81297f3..7b6ae0dc4 100644 --- a/hugr-py/rust/linking.rs +++ b/hugr-py/rust/linking.rs @@ -25,14 +25,15 @@ pub mod linking { } #[pyfunction] - fn link_modules(module_into: &[u8], module_from: &[u8]) -> PyResult> { - let (mut hugr_into, mut exts_into) = - Hugr::load_with_exts(module_into, None).map_err(|err| { + fn link_modules(module_into: Vec, module_from: Vec) -> PyResult> { + let (mut hugr_into, mut exts_into) = Hugr::load_with_exts(module_into.as_slice(), None) + .map_err(|err| { PyValueError::new_err(format!("Loading of first envelope failed: {}", err)) })?; - let (hugr_from, exts_from) = Hugr::load_with_exts(module_from, None).map_err(|err| { - PyValueError::new_err(format!("Loading of second envelope failed: {}", err)) - })?; + let (hugr_from, exts_from) = + Hugr::load_with_exts(module_from.as_slice(), None).map_err(|err| { + PyValueError::new_err(format!("Loading of second envelope failed: {}", err)) + })?; let into_executable = hugr_into.entrypoint() != hugr_into.module_root(); let from_executable = hugr_from.entrypoint() != hugr_from.module_root(); let replacement_entrypoint = if into_executable && from_executable { From bf371ab8f31c8aa69a9a3054f9de71547557fb7f Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Maximilian=20R=C3=BCsch?= Date: Tue, 17 Mar 2026 12:14:15 +0000 Subject: [PATCH 12/16] Use debug assert --- hugr-py/rust/linking.rs | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/hugr-py/rust/linking.rs b/hugr-py/rust/linking.rs index 7b6ae0dc4..5a02d3de7 100644 --- a/hugr-py/rust/linking.rs +++ b/hugr-py/rust/linking.rs @@ -62,8 +62,7 @@ pub mod linking { .store_with_exts(&mut result, EnvelopeConfig::binary(), &exts_into) .unwrap(); - hugr_core::package::Package::load(&result[..], None) - .map_err(|err| PyValueError::new_err(format!("Roundtrip failed: {:?}", err)))?; + debug_assert!(hugr_core::package::Package::load(&result[..], None).is_ok()); Ok(result) } From 3a2323c293bafc6f061f508ee863cb6b16cb9790 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Maximilian=20R=C3=BCsch?= Date: Tue, 17 Mar 2026 12:17:23 +0000 Subject: [PATCH 13/16] Add comment --- hugr-py/rust/linking.rs | 1 + 1 file changed, 1 insertion(+) diff --git a/hugr-py/rust/linking.rs b/hugr-py/rust/linking.rs index 5a02d3de7..8ba56b2b3 100644 --- a/hugr-py/rust/linking.rs +++ b/hugr-py/rust/linking.rs @@ -62,6 +62,7 @@ pub mod linking { .store_with_exts(&mut result, EnvelopeConfig::binary(), &exts_into) .unwrap(); + // Sanity check roundtrip debug_assert!(hugr_core::package::Package::load(&result[..], None).is_ok()); Ok(result) From 1d200f333d893159fb0516988f3e441c685a68cd Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Maximilian=20R=C3=BCsch?= Date: Tue, 17 Mar 2026 14:35:14 +0000 Subject: [PATCH 14/16] Revert for better performance --- hugr-py/rust/linking.rs | 13 ++++++------- 1 file changed, 6 insertions(+), 7 deletions(-) diff --git a/hugr-py/rust/linking.rs b/hugr-py/rust/linking.rs index 8ba56b2b3..a07e5f55e 100644 --- a/hugr-py/rust/linking.rs +++ b/hugr-py/rust/linking.rs @@ -25,15 +25,14 @@ pub mod linking { } #[pyfunction] - fn link_modules(module_into: Vec, module_from: Vec) -> PyResult> { - let (mut hugr_into, mut exts_into) = Hugr::load_with_exts(module_into.as_slice(), None) - .map_err(|err| { + fn link_modules(module_into: &[u8], module_from: &[u8]) -> PyResult> { + let (mut hugr_into, mut exts_into) = + Hugr::load_with_exts(module_into, None).map_err(|err| { PyValueError::new_err(format!("Loading of first envelope failed: {}", err)) })?; - let (hugr_from, exts_from) = - Hugr::load_with_exts(module_from.as_slice(), None).map_err(|err| { - PyValueError::new_err(format!("Loading of second envelope failed: {}", err)) - })?; + let (hugr_from, exts_from) = Hugr::load_with_exts(module_from, None).map_err(|err| { + PyValueError::new_err(format!("Loading of second envelope failed: {}", err)) + })?; let into_executable = hugr_into.entrypoint() != hugr_into.module_root(); let from_executable = hugr_from.entrypoint() != hugr_from.module_root(); let replacement_entrypoint = if into_executable && from_executable { From 7b7e7908d6f836cc8fb47dc25d0c00e2d261e2b7 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Maximilian=20R=C3=BCsch?= Date: Tue, 17 Mar 2026 15:01:54 +0000 Subject: [PATCH 15/16] Refactor load to use load_with_exts --- hugr-core/src/hugr.rs | 7 ++----- 1 file changed, 2 insertions(+), 5 deletions(-) diff --git a/hugr-core/src/hugr.rs b/hugr-core/src/hugr.rs index fbfea7e1a..5423da243 100644 --- a/hugr-core/src/hugr.rs +++ b/hugr-core/src/hugr.rs @@ -159,11 +159,8 @@ impl Hugr { reader: impl io::BufRead, extensions: Option<&ExtensionRegistry>, ) -> Result { - let pkg = Package::load(reader, extensions)?; - match pkg.modules.into_iter().exactly_one() { - Ok(hugr) => Ok(hugr), - Err(e) => Err(ReadError::ExpectedSingleHugr { count: e.count() }), - } + let (hugr, _) = Self::load_with_exts(reader, extensions)?; + Ok(hugr) } /// Read a HUGR from an Envelope, and return the enclosed extensions. From 705bc20490ec2add05c210e58a7f4191df92a3e9 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Maximilian=20R=C3=BCsch?= Date: Tue, 17 Mar 2026 15:37:30 +0000 Subject: [PATCH 16/16] Add load_with_exts test --- hugr-core/src/hugr.rs | 32 ++++++++++++++++++++++++++++++++ 1 file changed, 32 insertions(+) diff --git a/hugr-core/src/hugr.rs b/hugr-core/src/hugr.rs index 5423da243..dcbeff4f2 100644 --- a/hugr-core/src/hugr.rs +++ b/hugr-core/src/hugr.rs @@ -662,11 +662,16 @@ fn make_module_hugr(root_op: OpType, nodes: usize, ports: usize) -> Option #[cfg(test)] pub(crate) mod test { + use crate::Extension; + use crate::extension::prelude::qb_t; + use crate::extension::prelude::usize_t; use std::{fs::File, io::BufReader}; use super::*; + use crate::builder::test::simple_package; use crate::builder::{Container, Dataflow, DataflowSubContainer, ModuleBuilder}; + use crate::extension::ExtensionId; use crate::extension::prelude::bool_t; use crate::ops::OpaqueOp; use crate::ops::handle::NodeHandle; @@ -854,4 +859,31 @@ pub(crate) mod test { } } } + + #[rstest] + fn load_extensions() { + let my_ext_id = ExtensionId::new("test.ext").unwrap(); + let my_ext = Extension::new_test_arc(my_ext_id, |ext, extension_ref| { + ext.add_op( + "MyOp".into(), + String::new(), + Signature::new(vec![qb_t(), usize_t()], vec![qb_t()]), + extension_ref, + ) + .unwrap(); + }); + + let mut package = simple_package(); + package.extensions.register(my_ext).unwrap(); + let mut hugr_str = Vec::new(); + package + .store(&mut hugr_str, EnvelopeConfig::default()) + .unwrap(); + + let (_, exts) = Hugr::load_with_exts(hugr_str.as_slice(), None).unwrap(); + assert_eq!(exts.len(), 1); + assert_matches!(exts.get("test.ext"), Some(ext) => { + assert!(ext.get_op("MyOp").is_some()); + }); + } }