feat!: Use binary envelopes for operation lower_func encoding (#2447)

~~Blocked by #2445 due to errors with the hugr-py model encoding.~~

Changes the format used to encode the `LowerFunc::FixedHugr`s in an
extension's operation from string-encoded envelopes to base64-encoded
binary envelopes.

Adds an `envelope::serde_with::AsBinaryEnvelope` wrapper similar to
`AsStringEnvelope`, that lets us include base64-encoded envelopes in a
serde definition.

- hugr-py does not yet support loading binary envelopes, so we ignore
lowering functions when loading extensions.
- The deserializer can still load older extensions. We check for the
envelope magic numbers on the encoded string and skip the base64
deserialization if found.
- Older hugr versions will **not** be able to read extensions encoded
this way, and will fail with an "invalid magic" error.
- The schema does not change, since the hugr field is still a `string`.

drive-by: Adds a custom deserializing function for `LowerFunc` so errors
found while decoding the fixed hugr get communicated back up. The
default `serde` derive produced an opaque "no variants match" error due
to using `#[serde(untagged)]`.

BREAKING CHANGE: Lowering functions in extension operations are now
encoded as binary envelopes. Older hugr versions will error out when
trying to load them.

Author: aborgna-q

hugr-core/src/envelope/serde_with.rs

@@ -36,7 +36,7 @@ mod type_def;
 pub use const_fold::{ConstFold, ConstFoldResult, Folder, fold_out_row};
 pub use op_def::{
     CustomSignatureFunc, CustomValidator, LowerFunc, OpDef, SignatureFromArgs, SignatureFunc,
-    ValidateJustArgs, ValidateTypeArgs,
+    ValidateJustArgs, ValidateTypeArgs, deserialize_lower_funcs,
 };
 pub use prelude::{PRELUDE, PRELUDE_REGISTRY};
 pub use type_def::{TypeDef, TypeDefBound};

hugr-core/src/extension.rs

@@ -12,7 +12,7 @@ use super::{
 };
 
 use crate::Hugr;
-use crate::envelope::serde_with::AsStringEnvelope;
+use crate::envelope::serde_with::AsBinaryEnvelope;
 use crate::ops::{OpName, OpNameRef};
 use crate::types::type_param::{TypeArg, TypeParam, check_term_types};
 use crate::types::{FuncValueType, PolyFuncType, PolyFuncTypeRV, Signature};
@@ -268,8 +268,12 @@ impl Debug for SignatureFunc {
 
 /// Different ways that an [OpDef] can lower operation nodes i.e. provide a Hugr
 /// that implements the operation using a set of other extensions.
+///
+/// Does not implement [`serde::Deserialize`] directly since the serde error for
+/// untagged enums is unhelpful. Use [`deserialize_lower_funcs`] with
+/// [`serde(deserialize_with = "deserialize_lower_funcs")] instead.
 #[serde_as]
-#[derive(serde::Deserialize, serde::Serialize)]
+#[derive(serde::Serialize)]
 #[serde(untagged)]
 pub enum LowerFunc {
     /// Lowering to a fixed Hugr. Since this cannot depend upon the [TypeArg]s,
@@ -281,7 +285,7 @@ pub enum LowerFunc {
         /// [OpDef]
         ///
         /// [ExtensionOp]: crate::ops::ExtensionOp
-        #[serde_as(as = "Box<AsStringEnvelope>")]
+        #[serde_as(as = "Box<AsBinaryEnvelope>")]
         hugr: Box<Hugr>,
     },
     /// Custom binary function that can (fallibly) compute a Hugr
@@ -290,6 +294,34 @@ pub enum LowerFunc {
     CustomFunc(Box<dyn CustomLowerFunc>),
 }
 
+/// A function for deserializing sequences of [`LowerFunc::FixedHugr`].
+///
+/// We could let serde deserialize [`LowerFunc`] as-is, but if the LowerFunc
+/// deserialization fails it just returns an opaque "data did not match any
+/// variant of untagged enum LowerFunc" error. This function will return the
+/// internal errors instead.
+pub fn deserialize_lower_funcs<'de, D>(deserializer: D) -> Result<Vec<LowerFunc>, D::Error>
+where
+    D: serde::Deserializer<'de>,
+{
+    #[serde_as]
+    #[derive(serde::Deserialize)]
+    struct FixedHugrDeserializer {
+        pub extensions: ExtensionSet,
+        #[serde_as(as = "Box<AsBinaryEnvelope>")]
+        pub hugr: Box<Hugr>,
+    }
+
+    let funcs: Vec<FixedHugrDeserializer> = serde::Deserialize::deserialize(deserializer)?;
+    Ok(funcs
+        .into_iter()
+        .map(|f| LowerFunc::FixedHugr {
+            extensions: f.extensions,
+            hugr: f.hugr,
+        })
+        .collect())
+}
+
 impl Debug for LowerFunc {
     fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
         match self {
@@ -322,7 +354,11 @@ pub struct OpDef {
     signature_func: SignatureFunc,
     // Some operations cannot lower themselves and tools that do not understand them
     // can only treat them as opaque/black-box ops.
-    #[serde(default, skip_serializing_if = "Vec::is_empty")]
+    #[serde(
+        default,
+        skip_serializing_if = "Vec::is_empty",
+        deserialize_with = "deserialize_lower_funcs"
+    )]
     pub(crate) lower_funcs: Vec<LowerFunc>,
 
     /// Operations can optionally implement [`ConstFold`] to implement constant folding.

hugr-core/src/extension/op_def.rs

@@ -66,12 +66,26 @@ def deserialize(self, extension: ext.Extension) -> ext.TypeDef:
 
 
 class FixedHugr(ConfiguredBaseModel):
+    """Fixed HUGR used to define the lowering of an operation.
+
+    Args:
+        extensions: Extensions used in the HUGR.
+        hugr: Base64-encoded HUGR envelope.
+    """
+
     extensions: ExtensionSet
     hugr: str
 
     def deserialize(self) -> ext.FixedHugr:
-        hugr = Hugr.from_str(self.hugr)
-        return ext.FixedHugr(extensions=self.extensions, hugr=hugr)
+        # Loading fixed HUGRs requires reading hugr-model envelopes,
+        # which is not currently supported in Python.
+        # TODO: Add support for loading fixed HUGRs in Python.
+        # https://github.com/CQCL/hugr/issues/2287
+        msg = (
+            "Loading extensions with operation lowering functions is not "
+            + "supported in Python"
+        )
+        raise NotImplementedError(msg)
 
 
 class OpDef(ConfiguredBaseModel, populate_by_name=True):
@@ -91,13 +105,21 @@ def deserialize(self, extension: ext.Extension) -> ext.OpDef:
             self.binary,
         )
 
+        # Loading fixed HUGRs requires reading hugr-model envelopes,
+        # which is not currently supported in Python.
+        # We currently ignore any lower functions instead of raising an error.
+        #
+        # TODO: Add support for loading fixed HUGRs in Python.
+        # https://github.com/CQCL/hugr/issues/2287
+        lower_funcs: list[ext.FixedHugr] = []
+
         return extension.add_op_def(
             ext.OpDef(
                 name=self.name,
                 description=self.description,
                 misc=self.misc or {},
                 signature=signature,
-                lower_funcs=[f.deserialize() for f in self.lower_funcs],
+                lower_funcs=lower_funcs,
             )
         )
 

hugr-py/src/hugr/_serialization/extension.py

@@ -2,6 +2,7 @@
 
 from __future__ import annotations
 
+import base64
 from dataclasses import dataclass, field
 from typing import TYPE_CHECKING, Any, TypeVar
 
@@ -154,7 +155,8 @@ class FixedHugr:
     hugr: Hugr
 
     def _to_serial(self) -> ext_s.FixedHugr:
-        return ext_s.FixedHugr(extensions=self.extensions, hugr=self.hugr.to_str())
+        hugr_64: str = base64.b64encode(self.hugr.to_bytes()).decode()
+        return ext_s.FixedHugr(extensions=self.extensions, hugr=hugr_64)
 
 
 @dataclass

hugr-py/src/hugr/ext.py

@@ -226,6 +226,12 @@ def validate(
                         h1_hash == h2_hash
                     ), f"HUGRs are not the same for {write_fmt} -> {load_fmt}"
 
+                # Lowering functions are currently ignored in Python,
+                # because we don't support loading -model envelopes yet.
+                for ext in loaded.extensions:
+                    for op in ext.operations.values():
+                        assert op.lower_funcs == []
+
 
 @dataclass(frozen=True, order=True)
 class _NodeHash:

hugr-py/tests/conftest.py

@@ -581,6 +581,7 @@
             "type": "object"
         },
         "FixedHugr": {
+            "description": "Fixed HUGR used to define the lowering of an operation.\n\nArgs:\n    extensions: Extensions used in the HUGR.\n    hugr: Base64-encoded HUGR envelope.",
             "properties": {
                 "extensions": {
                     "items": {

specification/schema/hugr_schema_live.json

@@ -581,6 +581,7 @@
             "type": "object"
         },
         "FixedHugr": {
+            "description": "Fixed HUGR used to define the lowering of an operation.\n\nArgs:\n    extensions: Extensions used in the HUGR.\n    hugr: Base64-encoded HUGR envelope.",
             "properties": {
                 "extensions": {
                     "items": {

specification/schema/hugr_schema_strict_live.json

@@ -581,6 +581,7 @@
             "type": "object"
         },
         "FixedHugr": {
+            "description": "Fixed HUGR used to define the lowering of an operation.\n\nArgs:\n    extensions: Extensions used in the HUGR.\n    hugr: Base64-encoded HUGR envelope.",
             "properties": {
                 "extensions": {
                     "items": {

specification/schema/testing_hugr_schema_live.json

@@ -581,6 +581,7 @@
             "type": "object"
         },
         "FixedHugr": {
+            "description": "Fixed HUGR used to define the lowering of an operation.\n\nArgs:\n    extensions: Extensions used in the HUGR.\n    hugr: Base64-encoded HUGR envelope.",
             "properties": {
                 "extensions": {
                     "items": {