refactor: Direct import of model representation to Python (#2683)

Closes #2287 .

Testing:

* All existing hugr-py tests that call `validate()` perform a
round-trip, checking that the hashes of the start and end Hugrs
(computed using the `_NodeHash` method defined in `conftest.py`) agree.
* I exported all the Hugrs generated by tests in the guppylang repo
(using `just export-integration-tests`) and tried importing them (1)
using the old method via json and (2) using the new method, and checked
that the hashes of the two imported Hugrs agree (using the same
`_NodeHash` method).

There are a few TODO comments remaining in `load.py` which I could not
see a way to resolve using the existing model. I will investigate these
further and raise issues if necessary.

---------

Co-authored-by: Lukas Heidemann <[email protected]>

Author: cqc-alec

hugr-core/src/import.rs

@@ -1140,12 +1140,12 @@ impl<'a> Context<'a> {
 
     /// Import a node with a custom operation.
     ///
-    /// A custom operation in `hugr-model` referred to by a symbol application
-    /// term. The name of the symbol specifies the name of the custom operation,
-    /// and the arguments supplied to the symbol are the arguments to be passed
-    /// to the custom operation. This method imports the custom operations as
-    /// [`OpaqueOp`]s. The [`OpaqueOp`]s are then resolved later against the
-    /// [`ExtensionRegistry`].
+    /// A custom operation in `hugr-model` is referred to by a symbol
+    /// application term. The name of the symbol specifies the name of the
+    /// custom operation, and the arguments supplied to the symbol are the
+    /// arguments to be passed to the custom operation. This method imports the
+    /// custom operations as [`OpaqueOp`]s. The [`OpaqueOp`]s are then resolved
+    /// later against the [`ExtensionRegistry`].
     ///
     /// Some operations that needed to be builtins in `hugr-core` are custom
     /// operations in `hugr-model`. This method detects these and converts them

hugr-model/src/v0/binary/mod.rs

@@ -10,5 +10,5 @@
 mod read;
 mod write;
 
-pub use read::{ReadError, read_from_reader, read_from_slice};
+pub use read::{ReadError, read_from_reader, read_from_slice, read_from_slice_with_suffix};
 pub use write::{WriteError, write_to_vec, write_to_writer};

hugr-model/src/v0/binary/read.rs

@@ -3,7 +3,7 @@ use crate::v0::table;
 use crate::{CURRENT_VERSION, v0 as model};
 use bumpalo::Bump;
 use bumpalo::collections::Vec as BumpVec;
-use std::io::BufRead;
+use std::io::{BufRead, BufReader, Read};
 
 /// An error encountered while deserialising a model.
 #[derive(Debug, derive_more::From, derive_more::Display, derive_more::Error)]
@@ -27,10 +27,27 @@ pub enum ReadError {
 type ReadResult<T> = Result<T, ReadError>;
 
 /// Read a hugr package from a byte slice.
+///
+/// If the slice contains bytes beyond the encoded package, they are ignored.
 pub fn read_from_slice<'a>(slice: &[u8], bump: &'a Bump) -> ReadResult<table::Package<'a>> {
     read_from_reader(slice, bump)
 }
 
+/// Read a hugr package from a byte slice.
+///
+/// If the slice contains bytes beyond the encoded package, they are returned
+/// along with the decoded package.
+pub fn read_from_slice_with_suffix<'a>(
+    slice: &[u8],
+    bump: &'a Bump,
+) -> ReadResult<(table::Package<'a>, Vec<u8>)> {
+    let mut buffer = BufReader::new(slice);
+    let package = read_from_reader(&mut buffer, bump)?;
+    let mut suffix: Vec<u8> = vec![];
+    buffer.read_to_end(&mut suffix)?;
+    Ok((package, suffix))
+}
+
 /// Read a hugr package from an impl of [`BufRead`].
 pub fn read_from_reader(reader: impl BufRead, bump: &Bump) -> ReadResult<table::Package<'_>> {
     let reader =

hugr-py/rust/lib.rs

@@ -74,14 +74,14 @@ fn package_to_bytes(package: ast::Package) -> PyResult<Vec<u8>> {
 }
 
 #[pyfunction]
-fn bytes_to_package(bytes: &[u8]) -> PyResult<ast::Package> {
+fn bytes_to_package(bytes: &[u8]) -> PyResult<(ast::Package, Vec<u8>)> {
     let bump = bumpalo::Bump::new();
-    let table = hugr_model::v0::binary::read_from_slice(bytes, &bump)
+    let (table, suffix) = hugr_model::v0::binary::read_from_slice_with_suffix(bytes, &bump)
         .map_err(|err| PyValueError::new_err(err.to_string()))?;
     let package = table
         .as_ast()
         .ok_or_else(|| PyValueError::new_err("Malformed package"))?;
-    Ok(package)
+    Ok((package, suffix))
 }
 
 /// Returns the current version of the HUGR model format as a tuple of (major, minor, patch).

hugr-py/src/hugr/_hugr/__init__.pyi

@@ -20,7 +20,7 @@ def bytes_to_module(binary: bytes) -> hugr.model.Module: ...
 def package_to_string(package: hugr.model.Package) -> str: ...
 def string_to_package(string: str) -> hugr.model.Package: ...
 def package_to_bytes(package: hugr.model.Package) -> bytes: ...
-def bytes_to_package(binary: bytes) -> hugr.model.Package: ...
+def bytes_to_package(binary: bytes) -> tuple[hugr.model.Package, bytes]: ...
 def current_model_version() -> tuple[int, int, int]: ...
 def to_json_envelope(binary: bytes) -> bytes: ...
 def run_cli() -> None: ...

hugr-py/src/hugr/envelope.py

@@ -39,7 +39,7 @@
 
 import pyzstd
 
-from hugr import cli
+import hugr._hugr as rust
 
 if TYPE_CHECKING:
     from hugr.hugr.base import Hugr
@@ -64,7 +64,6 @@ def make_envelope(package: Package | Hugr, config: EnvelopeConfig) -> bytes:
     if not isinstance(package, Package):
         package = Package(modules=[package], extensions=[])
 
-    # Currently only uncompressed JSON is supported.
     payload: bytes
     match config.format:
         case EnvelopeFormat.JSON:
@@ -103,6 +102,7 @@ def make_envelope_str(package: Package | Hugr, config: EnvelopeConfig) -> str:
 def read_envelope(envelope: bytes) -> Package:
     """Decode a HUGR package from an envelope."""
     import hugr._serialization.extension as ext_s
+    from hugr.package import Package
 
     header = EnvelopeHeader.from_bytes(envelope)
     payload = envelope[10:]
@@ -113,12 +113,23 @@ def read_envelope(envelope: bytes) -> Package:
     match header.format:
         case EnvelopeFormat.JSON:
             return ext_s.Package.model_validate_json(payload).deserialize()
-        case EnvelopeFormat.MODEL | EnvelopeFormat.MODEL_WITH_EXTS:
-            # TODO Going via JSON is a temporary solution, until we get model import to
-            # python properly implemented.
-            # https://github.com/CQCL/hugr/issues/2287
-            json_data = cli.convert(envelope, format="json")
-            return read_envelope(json_data)
+        case EnvelopeFormat.MODEL:
+            model_package, suffix = rust.bytes_to_package(payload)
+            if suffix:
+                msg = f"Excess bytes in envelope with format {EnvelopeFormat.MODEL}."
+                raise ValueError(msg)
+            return Package.from_model(model_package)
+        case EnvelopeFormat.MODEL_WITH_EXTS:
+            from hugr.ext import Extension
+
+            model_package, suffix = rust.bytes_to_package(payload)
+            return Package(
+                modules=Package.from_model(model_package).modules,
+                extensions=[
+                    Extension.from_json(json.dumps(extension))
+                    for extension in json.loads(suffix)
+                ],
+            )
 
 
 def read_envelope_hugr(envelope: bytes) -> Hugr:

hugr-py/src/hugr/hugr/base.py

@@ -1114,6 +1114,19 @@ def from_str(envelope: str) -> Hugr:
         """
         return read_envelope_hugr_str(envelope)
 
+    @staticmethod
+    def from_model(module: model.Module) -> Hugr:
+        """Import from the hugr model format."""
+        from hugr.model.load import ModelImport
+
+        loader = ModelImport(module=module)
+        for i, node in enumerate(module.root.children):
+            loader.import_node_in_module(node, i)
+        loader.link_ports()
+        loader.link_static_ports()
+        loader.add_module_metadata()
+        return loader.hugr
+
     def to_bytes(self, config: EnvelopeConfig | None = None) -> bytes:
         """Serialize the HUGR into an envelope byte string.
 

hugr-py/src/hugr/model/__init__.py

@@ -1,9 +1,9 @@
 """HUGR model data structures."""
 
-from collections.abc import Sequence
+import warnings
+from collections.abc import Generator, Sequence
 from dataclasses import dataclass, field
 from enum import Enum
-from typing import Protocol
 
 from semver import Version
 
@@ -21,7 +21,7 @@ def _current_version() -> Version:
 CURRENT_VERSION: Version = _current_version()
 
 
-class Term(Protocol):
+class Term:
     """A model term for static data such as types, constants and metadata."""
 
     def __str__(self) -> str:
@@ -33,6 +33,26 @@ def from_str(s: str) -> "Term":
         """Read the term from its string representation."""
         return rust.string_to_term(s)
 
+    def to_list_parts(self) -> Generator["SeqPart"]:
+        if isinstance(self, List):
+            for part in self.parts:
+                if isinstance(part, Splice):
+                    yield from part.seq.to_list_parts()
+                else:
+                    yield part
+        else:
+            yield Splice(self)
+
+    def to_tuple_parts(self) -> Generator["SeqPart"]:
+        if isinstance(self, Tuple):
+            for part in self.parts:
+                if isinstance(part, Splice):
+                    yield from part.seq.to_tuple_parts()
+                else:
+                    yield part
+        else:
+            yield Splice(self)
+
 
 @dataclass(frozen=True)
 class Wildcard(Term):
@@ -129,9 +149,13 @@ def from_str(s: str) -> "Symbol":
         return rust.string_to_symbol(s)
 
 
-class Op(Protocol):
+class Op:
     """The operation of a node."""
 
+    def symbol_name(self) -> str | None:
+        """Returns name of the symbol introduced by this node, if any."""
+        return None
+
 
 @dataclass(frozen=True)
 class InvalidOp(Op):
@@ -159,13 +183,19 @@ class DefineFunc(Op):
 
     symbol: Symbol
 
+    def symbol_name(self) -> str | None:
+        return self.symbol.name
+
 
 @dataclass(frozen=True)
 class DeclareFunc(Op):
     """Function declaration."""
 
     symbol: Symbol
 
+    def symbol_name(self) -> str | None:
+        return self.symbol.name
+
 
 @dataclass(frozen=True)
 class CustomOp(Op):
@@ -181,13 +211,19 @@ class DefineAlias(Op):
     symbol: Symbol
     value: Term
 
+    def symbol_name(self) -> str | None:
+        return self.symbol.name
+
 
 @dataclass(frozen=True)
 class DeclareAlias(Op):
     """Alias declaration."""
 
     symbol: Symbol
 
+    def symbol_name(self) -> str | None:
+        return self.symbol.name
+
 
 @dataclass(frozen=True)
 class TailLoop(Op):
@@ -205,20 +241,29 @@ class DeclareConstructor(Op):
 
     symbol: Symbol
 
+    def symbol_name(self) -> str | None:
+        return self.symbol.name
+
 
 @dataclass(frozen=True)
 class DeclareOperation(Op):
     """Operation declaration."""
 
     symbol: Symbol
 
+    def symbol_name(self) -> str | None:
+        return self.symbol.name
+
 
 @dataclass(frozen=True)
 class Import(Op):
     """Import operation."""
 
     name: str
 
+    def symbol_name(self) -> str | None:
+        return self.name
+
 
 @dataclass
 class Node:
@@ -302,7 +347,13 @@ def from_str(s: str) -> "Package":
     @staticmethod
     def from_bytes(b: bytes) -> "Package":
         """Read a package from its binary representation."""
-        return rust.bytes_to_package(b)
+        package, suffix = rust.bytes_to_package(b)
+        if suffix:
+            warnings.warn(
+                "Binary encoding of model Package contains extra bytes: ignoring them.",
+                stacklevel=2,
+            )
+        return package
 
     @property
     def version(self) -> Version: