Skip to content

Commit 2c16a77

Browse files
authored
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.
1 parent 4bc7f65 commit 2c16a77

File tree

10 files changed

+460
-10
lines changed

10 files changed

+460
-10
lines changed

hugr-core/src/envelope/serde_with.rs

Lines changed: 381 additions & 1 deletion
Large diffs are not rendered by default.

hugr-core/src/extension.rs

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -36,7 +36,7 @@ mod type_def;
3636
pub use const_fold::{ConstFold, ConstFoldResult, Folder, fold_out_row};
3737
pub use op_def::{
3838
CustomSignatureFunc, CustomValidator, LowerFunc, OpDef, SignatureFromArgs, SignatureFunc,
39-
ValidateJustArgs, ValidateTypeArgs,
39+
ValidateJustArgs, ValidateTypeArgs, deserialize_lower_funcs,
4040
};
4141
pub use prelude::{PRELUDE, PRELUDE_REGISTRY};
4242
pub use type_def::{TypeDef, TypeDefBound};

hugr-core/src/extension/op_def.rs

Lines changed: 40 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -12,7 +12,7 @@ use super::{
1212
};
1313

1414
use crate::Hugr;
15-
use crate::envelope::serde_with::AsStringEnvelope;
15+
use crate::envelope::serde_with::AsBinaryEnvelope;
1616
use crate::ops::{OpName, OpNameRef};
1717
use crate::types::type_param::{TypeArg, TypeParam, check_term_types};
1818
use crate::types::{FuncValueType, PolyFuncType, PolyFuncTypeRV, Signature};
@@ -268,8 +268,12 @@ impl Debug for SignatureFunc {
268268

269269
/// Different ways that an [OpDef] can lower operation nodes i.e. provide a Hugr
270270
/// that implements the operation using a set of other extensions.
271+
///
272+
/// Does not implement [`serde::Deserialize`] directly since the serde error for
273+
/// untagged enums is unhelpful. Use [`deserialize_lower_funcs`] with
274+
/// [`serde(deserialize_with = "deserialize_lower_funcs")] instead.
271275
#[serde_as]
272-
#[derive(serde::Deserialize, serde::Serialize)]
276+
#[derive(serde::Serialize)]
273277
#[serde(untagged)]
274278
pub enum LowerFunc {
275279
/// Lowering to a fixed Hugr. Since this cannot depend upon the [TypeArg]s,
@@ -281,7 +285,7 @@ pub enum LowerFunc {
281285
/// [OpDef]
282286
///
283287
/// [ExtensionOp]: crate::ops::ExtensionOp
284-
#[serde_as(as = "Box<AsStringEnvelope>")]
288+
#[serde_as(as = "Box<AsBinaryEnvelope>")]
285289
hugr: Box<Hugr>,
286290
},
287291
/// Custom binary function that can (fallibly) compute a Hugr
@@ -290,6 +294,34 @@ pub enum LowerFunc {
290294
CustomFunc(Box<dyn CustomLowerFunc>),
291295
}
292296

297+
/// A function for deserializing sequences of [`LowerFunc::FixedHugr`].
298+
///
299+
/// We could let serde deserialize [`LowerFunc`] as-is, but if the LowerFunc
300+
/// deserialization fails it just returns an opaque "data did not match any
301+
/// variant of untagged enum LowerFunc" error. This function will return the
302+
/// internal errors instead.
303+
pub fn deserialize_lower_funcs<'de, D>(deserializer: D) -> Result<Vec<LowerFunc>, D::Error>
304+
where
305+
D: serde::Deserializer<'de>,
306+
{
307+
#[serde_as]
308+
#[derive(serde::Deserialize)]
309+
struct FixedHugrDeserializer {
310+
pub extensions: ExtensionSet,
311+
#[serde_as(as = "Box<AsBinaryEnvelope>")]
312+
pub hugr: Box<Hugr>,
313+
}
314+
315+
let funcs: Vec<FixedHugrDeserializer> = serde::Deserialize::deserialize(deserializer)?;
316+
Ok(funcs
317+
.into_iter()
318+
.map(|f| LowerFunc::FixedHugr {
319+
extensions: f.extensions,
320+
hugr: f.hugr,
321+
})
322+
.collect())
323+
}
324+
293325
impl Debug for LowerFunc {
294326
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
295327
match self {
@@ -322,7 +354,11 @@ pub struct OpDef {
322354
signature_func: SignatureFunc,
323355
// Some operations cannot lower themselves and tools that do not understand them
324356
// can only treat them as opaque/black-box ops.
325-
#[serde(default, skip_serializing_if = "Vec::is_empty")]
357+
#[serde(
358+
default,
359+
skip_serializing_if = "Vec::is_empty",
360+
deserialize_with = "deserialize_lower_funcs"
361+
)]
326362
pub(crate) lower_funcs: Vec<LowerFunc>,
327363

328364
/// Operations can optionally implement [`ConstFold`] to implement constant folding.

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

Lines changed: 25 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -66,12 +66,26 @@ def deserialize(self, extension: ext.Extension) -> ext.TypeDef:
6666

6767

6868
class FixedHugr(ConfiguredBaseModel):
69+
"""Fixed HUGR used to define the lowering of an operation.
70+
71+
Args:
72+
extensions: Extensions used in the HUGR.
73+
hugr: Base64-encoded HUGR envelope.
74+
"""
75+
6976
extensions: ExtensionSet
7077
hugr: str
7178

7279
def deserialize(self) -> ext.FixedHugr:
73-
hugr = Hugr.from_str(self.hugr)
74-
return ext.FixedHugr(extensions=self.extensions, hugr=hugr)
80+
# Loading fixed HUGRs requires reading hugr-model envelopes,
81+
# which is not currently supported in Python.
82+
# TODO: Add support for loading fixed HUGRs in Python.
83+
# https://github.com/CQCL/hugr/issues/2287
84+
msg = (
85+
"Loading extensions with operation lowering functions is not "
86+
+ "supported in Python"
87+
)
88+
raise NotImplementedError(msg)
7589

7690

7791
class OpDef(ConfiguredBaseModel, populate_by_name=True):
@@ -91,13 +105,21 @@ def deserialize(self, extension: ext.Extension) -> ext.OpDef:
91105
self.binary,
92106
)
93107

108+
# Loading fixed HUGRs requires reading hugr-model envelopes,
109+
# which is not currently supported in Python.
110+
# We currently ignore any lower functions instead of raising an error.
111+
#
112+
# TODO: Add support for loading fixed HUGRs in Python.
113+
# https://github.com/CQCL/hugr/issues/2287
114+
lower_funcs: list[ext.FixedHugr] = []
115+
94116
return extension.add_op_def(
95117
ext.OpDef(
96118
name=self.name,
97119
description=self.description,
98120
misc=self.misc or {},
99121
signature=signature,
100-
lower_funcs=[f.deserialize() for f in self.lower_funcs],
122+
lower_funcs=lower_funcs,
101123
)
102124
)
103125

hugr-py/src/hugr/ext.py

Lines changed: 3 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -2,6 +2,7 @@
22

33
from __future__ import annotations
44

5+
import base64
56
from dataclasses import dataclass, field
67
from typing import TYPE_CHECKING, Any, TypeVar
78

@@ -154,7 +155,8 @@ class FixedHugr:
154155
hugr: Hugr
155156

156157
def _to_serial(self) -> ext_s.FixedHugr:
157-
return ext_s.FixedHugr(extensions=self.extensions, hugr=self.hugr.to_str())
158+
hugr_64: str = base64.b64encode(self.hugr.to_bytes()).decode()
159+
return ext_s.FixedHugr(extensions=self.extensions, hugr=hugr_64)
158160

159161

160162
@dataclass

hugr-py/tests/conftest.py

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -226,6 +226,12 @@ def validate(
226226
h1_hash == h2_hash
227227
), f"HUGRs are not the same for {write_fmt} -> {load_fmt}"
228228

229+
# Lowering functions are currently ignored in Python,
230+
# because we don't support loading -model envelopes yet.
231+
for ext in loaded.extensions:
232+
for op in ext.operations.values():
233+
assert op.lower_funcs == []
234+
229235

230236
@dataclass(frozen=True, order=True)
231237
class _NodeHash:

specification/schema/hugr_schema_live.json

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -581,6 +581,7 @@
581581
"type": "object"
582582
},
583583
"FixedHugr": {
584+
"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.",
584585
"properties": {
585586
"extensions": {
586587
"items": {

specification/schema/hugr_schema_strict_live.json

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -581,6 +581,7 @@
581581
"type": "object"
582582
},
583583
"FixedHugr": {
584+
"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.",
584585
"properties": {
585586
"extensions": {
586587
"items": {

specification/schema/testing_hugr_schema_live.json

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -581,6 +581,7 @@
581581
"type": "object"
582582
},
583583
"FixedHugr": {
584+
"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.",
584585
"properties": {
585586
"extensions": {
586587
"items": {

specification/schema/testing_hugr_schema_strict_live.json

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -581,6 +581,7 @@
581581
"type": "object"
582582
},
583583
"FixedHugr": {
584+
"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.",
584585
"properties": {
585586
"extensions": {
586587
"items": {

0 commit comments

Comments
 (0)