Skip to content

Commit

Permalink
fix: OpDef serialisation
Browse files Browse the repository at this point in the history
  • Loading branch information
@doug-q
doug-q committed May 8, 2024
1 parent dbacfe0 commit 62ce904
Show file tree
Hide file tree
Showing 6 changed files with 209 additions and 21 deletions.
20 changes: 11 additions & 9 deletions hugr-py/src/hugr/serialization/ops.py
View file
Original file line number Diff line number Diff line change
@@ -1,7 +1,7 @@
import inspect
import sys
from abc import ABC
from typing import Any, Literal
from typing import Any, Literal, Optional

from pydantic import Field, RootModel

Expand Down Expand Up @@ -506,18 +506,20 @@ class Config:
# --------------------------------------


class OpDef(BaseOp, populate_by_name=True):
class FixedHugr(ConfiguredBaseModel):
extensions: ExtensionSet
hugr: Any


class OpDef(ConfiguredBaseModel, populate_by_name=True):
"""Serializable definition for dynamically loaded operations."""

extension: ExtensionId
name: str # Unique identifier of the operation.
description: str # Human readable description of the operation.
inputs: list[tuple[str | None, Type]]
outputs: list[tuple[str | None, Type]]
misc: dict[str, Any] # Miscellaneous data associated with the operation.
def_: str | None = Field(
..., alias="def"
) # (YAML?)-encoded definition of the operation.
extension_reqs: ExtensionSet # Resources required to execute this operation.
misc: Optional[dict[str, Any]] = None
signature: Optional[PolyFuncType] = None
lower_funcs: list[FixedHugr]


# Now that all classes are defined, we need to update the ForwardRefs in all type
Expand Down
3 changes: 2 additions & 1 deletion hugr-py/src/hugr/serialization/testing_hugr.py
View file
Original file line number Diff line number Diff line change
@@ -1,7 +1,7 @@
from pydantic import ConfigDict
from typing import Literal
from .tys import Type, SumType, PolyFuncType, ConfiguredBaseModel, model_rebuild
from .ops import Value, OpType, classes as ops_classes
from .ops import Value, OpType, OpDef, classes as ops_classes


class TestingHugr(ConfiguredBaseModel):
Expand All @@ -14,6 +14,7 @@ class TestingHugr(ConfiguredBaseModel):
poly_func_type: PolyFuncType | None = None
value: Value | None = None
optype: OpType | None = None
op_def: OpDef | None = None

@classmethod
def get_version(cls) -> str:
Expand Down
2 changes: 1 addition & 1 deletion hugr/src/extension.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -28,7 +28,7 @@ pub use infer::{ExtensionSolution, InferExtensionError};

mod op_def;
pub use op_def::{
CustomSignatureFunc, CustomValidator, OpDef, SignatureFromArgs, SignatureFunc,
CustomSignatureFunc, CustomValidator, LowerFunc, OpDef, SignatureFromArgs, SignatureFunc,
ValidateJustArgs, ValidateTypeArgs,
};
mod type_def;
Expand Down
29 changes: 19 additions & 10 deletions hugr/src/extension/op_def.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -122,7 +122,7 @@ pub struct CustomValidator {
#[serde(flatten)]
poly_func: PolyFuncType,
#[serde(skip)]
validate: Box<dyn ValidateTypeArgs>,
pub(crate) validate: Box<dyn ValidateTypeArgs>,
}

impl CustomValidator {
Expand Down Expand Up @@ -265,21 +265,28 @@ 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.
#[derive(serde::Deserialize, serde::Serialize)]
#[serde(untagged)]
pub enum LowerFunc {
/// Lowering to a fixed Hugr. Since this cannot depend upon the [TypeArg]s,
/// this will generally only be applicable if the [OpDef] has no [TypeParam]s.
#[serde(rename = "hugr")]
FixedHugr(ExtensionSet, Hugr),
FixedHugr {
/// The extensions required by the [`Hugr`]
extensions: ExtensionSet,
/// The [`Hugr`] to be used to replace [`CustomOp`]s matching the parent
/// [`OpDef`]
hugr: Hugr,
},
/// Custom binary function that can (fallibly) compute a Hugr
/// for the particular instance and set of available extensions.
/// [`CustomOp`]: crate::op::CustomOp
#[serde(skip)]
CustomFunc(Box<dyn CustomLowerFunc>),
}

impl Debug for LowerFunc {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
match self {
Self::FixedHugr(_, _) => write!(f, "FixedHugr"),
Self::FixedHugr { .. } => write!(f, "FixedHugr"),
Self::CustomFunc(_) => write!(f, "<custom lower>"),
}
}
Expand All @@ -305,8 +312,7 @@ 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(flatten)]
lower_funcs: Vec<LowerFunc>,
pub(crate) lower_funcs: Vec<LowerFunc>,

/// Operations can optionally implement [`ConstFold`] to implement constant folding.
#[serde(skip)]
Expand Down Expand Up @@ -360,9 +366,9 @@ impl OpDef {
self.lower_funcs
.iter()
.flat_map(|f| match f {
LowerFunc::FixedHugr(req_res, h) => {
if available_extensions.is_superset(req_res) {
Some(h.clone())
LowerFunc::FixedHugr { extensions, hugr } => {
if available_extensions.is_superset(extensions) {
Some(hugr.clone())
} else {
None
}
Expand Down Expand Up @@ -495,7 +501,10 @@ mod test {
let type_scheme = PolyFuncType::new(vec![TP], FunctionType::new_endo(vec![list_of_var]));

let def = e.add_op(OP_NAME, "desc".into(), type_scheme)?;
def.add_lower_func(LowerFunc::FixedHugr(ExtensionSet::new(), Hugr::default()));
def.add_lower_func(LowerFunc::FixedHugr {
extensions: ExtensionSet::new(),
hugr: Hugr::default(),
});
def.add_misc("key", Default::default());
assert_eq!(def.description(), "desc");
assert_eq!(def.lower_funcs.len(), 1);
Expand Down
88 changes: 88 additions & 0 deletions specification/schema/testing_hugr_schema_strict_v1.json
View file
Original file line number Diff line number Diff line change
Expand Up @@ -817,6 +817,27 @@
"title": "ExtensionsParam",
"type": "object"
},
"FixedHugr": {
"additionalProperties": false,
"properties": {
"extensions": {
"items": {
"type": "string"
},
"title": "Extensions",
"type": "array"
},
"hugr": {
"title": "Hugr"
}
},
"required": [
"extensions",
"hugr"
],
"title": "FixedHugr",
"type": "object"
},
"FuncDecl": {
"additionalProperties": false,
"description": "External function declaration, linked at runtime.",
Expand Down Expand Up @@ -1363,6 +1384,62 @@
"title": "Noop",
"type": "object"
},
"OpDef": {
"additionalProperties": false,
"description": "Serializable definition for dynamically loaded operations.",
"properties": {
"extension": {
"title": "Extension",
"type": "string"
},
"name": {
"title": "Name",
"type": "string"
},
"description": {
"title": "Description",
"type": "string"
},
"misc": {
"anyOf": [
{
"type": "object"
},
{
"type": "null"
}
],
"default": null,
"title": "Misc"
},
"signature": {
"anyOf": [
{
"$ref": "#/$defs/PolyFuncType"
},
{
"type": "null"
}
],
"default": null
},
"lower_funcs": {
"items": {
"$ref": "#/$defs/FixedHugr"
},
"title": "Lower Funcs",
"type": "array"
}
},
"required": [
"extension",
"name",
"description",
"lower_funcs"
],
"title": "OpDef",
"type": "object"
},
"OpType": {
"description": "A constant operation.",
"discriminator": {
Expand Down Expand Up @@ -2325,6 +2402,17 @@
}
],
"default": null
},
"op_def": {
"anyOf": [
{
"$ref": "#/$defs/OpDef"
},
{
"type": "null"
}
],
"default": null
}
},
"title": "HugrTesting",
Expand Down
88 changes: 88 additions & 0 deletions specification/schema/testing_hugr_schema_v1.json
View file
Original file line number Diff line number Diff line change
Expand Up @@ -817,6 +817,27 @@
"title": "ExtensionsParam",
"type": "object"
},
"FixedHugr": {
"additionalProperties": true,
"properties": {
"extensions": {
"items": {
"type": "string"
},
"title": "Extensions",
"type": "array"
},
"hugr": {
"title": "Hugr"
}
},
"required": [
"extensions",
"hugr"
],
"title": "FixedHugr",
"type": "object"
},
"FuncDecl": {
"additionalProperties": true,
"description": "External function declaration, linked at runtime.",
Expand Down Expand Up @@ -1363,6 +1384,62 @@
"title": "Noop",
"type": "object"
},
"OpDef": {
"additionalProperties": true,
"description": "Serializable definition for dynamically loaded operations.",
"properties": {
"extension": {
"title": "Extension",
"type": "string"
},
"name": {
"title": "Name",
"type": "string"
},
"description": {
"title": "Description",
"type": "string"
},
"misc": {
"anyOf": [
{
"type": "object"
},
{
"type": "null"
}
],
"default": null,
"title": "Misc"
},
"signature": {
"anyOf": [
{
"$ref": "#/$defs/PolyFuncType"
},
{
"type": "null"
}
],
"default": null
},
"lower_funcs": {
"items": {
"$ref": "#/$defs/FixedHugr"
},
"title": "Lower Funcs",
"type": "array"
}
},
"required": [
"extension",
"name",
"description",
"lower_funcs"
],
"title": "OpDef",
"type": "object"
},
"OpType": {
"description": "A constant operation.",
"discriminator": {
Expand Down Expand Up @@ -2325,6 +2402,17 @@
}
],
"default": null
},
"op_def": {
"anyOf": [
{
"$ref": "#/$defs/OpDef"
},
{
"type": "null"
}
],
"default": null
}
},
"title": "HugrTesting",
Expand Down

0 comments on commit 62ce904

Please sign in to comment.