diff --git a/hugr-py/src/hugr/serialization/ops.py b/hugr-py/src/hugr/serialization/ops.py index 6f075f2cf..a66d0ed2e 100644 --- a/hugr-py/src/hugr/serialization/ops.py +++ b/hugr-py/src/hugr/serialization/ops.py @@ -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 @@ -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 diff --git a/hugr-py/src/hugr/serialization/testing_hugr.py b/hugr-py/src/hugr/serialization/testing_hugr.py index 59db4b80d..98904aaf6 100644 --- a/hugr-py/src/hugr/serialization/testing_hugr.py +++ b/hugr-py/src/hugr/serialization/testing_hugr.py @@ -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): @@ -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: diff --git a/hugr/src/extension.rs b/hugr/src/extension.rs index 13794406a..8ab82e7c3 100644 --- a/hugr/src/extension.rs +++ b/hugr/src/extension.rs @@ -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; diff --git a/hugr/src/extension/op_def.rs b/hugr/src/extension/op_def.rs index 1e68400ea..fb75d2e87 100644 --- a/hugr/src/extension/op_def.rs +++ b/hugr/src/extension/op_def.rs @@ -122,7 +122,7 @@ pub struct CustomValidator { #[serde(flatten)] poly_func: PolyFuncType, #[serde(skip)] - validate: Box, + pub(crate) validate: Box, } impl CustomValidator { @@ -265,13 +265,20 @@ 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), } @@ -279,7 +286,7 @@ pub enum LowerFunc { 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, ""), } } @@ -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, + pub(crate) lower_funcs: Vec, /// Operations can optionally implement [`ConstFold`] to implement constant folding. #[serde(skip)] @@ -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 } @@ -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); diff --git a/specification/schema/testing_hugr_schema_strict_v1.json b/specification/schema/testing_hugr_schema_strict_v1.json index 850d2674d..0202bc282 100644 --- a/specification/schema/testing_hugr_schema_strict_v1.json +++ b/specification/schema/testing_hugr_schema_strict_v1.json @@ -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.", @@ -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": { @@ -2325,6 +2402,17 @@ } ], "default": null + }, + "op_def": { + "anyOf": [ + { + "$ref": "#/$defs/OpDef" + }, + { + "type": "null" + } + ], + "default": null } }, "title": "HugrTesting", diff --git a/specification/schema/testing_hugr_schema_v1.json b/specification/schema/testing_hugr_schema_v1.json index 40e1b894c..40e57d4fe 100644 --- a/specification/schema/testing_hugr_schema_v1.json +++ b/specification/schema/testing_hugr_schema_v1.json @@ -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.", @@ -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": { @@ -2325,6 +2402,17 @@ } ], "default": null + }, + "op_def": { + "anyOf": [ + { + "$ref": "#/$defs/OpDef" + }, + { + "type": "null" + } + ], + "default": null } }, "title": "HugrTesting",