diff --git a/hugr-py/src/hugr/serialization/ops.py b/hugr-py/src/hugr/serialization/ops.py index 9ce52e5bd..0c62bd108 100644 --- a/hugr-py/src/hugr/serialization/ops.py +++ b/hugr-py/src/hugr/serialization/ops.py @@ -500,18 +500,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: dict[str, Any] | None = None + signature: PolyFuncType | None = 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 bd10146b3..32bf2b95f 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 b504487ed..263a2d14e 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..91a49cfa5 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,11 +265,19 @@ 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] + /// + /// [CustomOp]: crate::ops::CustomOp + hugr: Hugr, + }, /// Custom binary function that can (fallibly) compute a Hugr /// for the particular instance and set of available extensions. #[serde(skip)] @@ -279,7 +287,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 +313,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 +367,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 } @@ -477,7 +484,6 @@ mod test { use crate::std_extensions::collections::{EXTENSION, LIST_TYPENAME}; use crate::types::Type; use crate::types::{type_param::TypeParam, FunctionType, PolyFuncType, TypeArg, TypeBound}; - use crate::Hugr; use crate::{const_extension_ids, Extension}; const_extension_ids! { @@ -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: crate::builder::test::simple_dfg_hugr(), // this is nonsense, but we are not testing the actual lowering here + }); def.add_misc("key", Default::default()); assert_eq!(def.description(), "desc"); assert_eq!(def.lower_funcs.len(), 1); diff --git a/hugr/src/hugr/serialize/test.rs b/hugr/src/hugr/serialize/test.rs index 57ae945ed..c9ff8eae0 100644 --- a/hugr/src/hugr/serialize/test.rs +++ b/hugr/src/hugr/serialize/test.rs @@ -377,7 +377,7 @@ fn roundtrip_sumtype(#[case] sum_type: SumType) { // #[case(Value::extension(ConstF64::new(std::f64::NAN)))] // #[case(Value::extension(ConstF64::new(std::f64::INFINITY)))] // #[case(Value::extension(ConstF64::new(std::f64::NEG_INFINITY)))] -#[case(Value::extension(ConstF64::new(std::f64::MIN_POSITIVE)))] +#[case(Value::extension(ConstF64::new(f64::MIN_POSITIVE)))] #[case(Value::sum(1,[Value::extension(ConstInt::new_u(2,1).unwrap())], SumType::new([vec![], vec![INT_TYPES[2].clone()]])).unwrap())] #[case(Value::tuple([Value::false_val(), Value::extension(ConstInt::new_s(2,1).unwrap())]))] #[case(Value::function(crate::builder::test::simple_dfg_hugr()).unwrap())] diff --git a/specification/schema/testing_hugr_schema_strict_v1.json b/specification/schema/testing_hugr_schema_strict_v1.json index 776715a99..39687363c 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": "TestingHugr", diff --git a/specification/schema/testing_hugr_schema_v1.json b/specification/schema/testing_hugr_schema_v1.json index ca1d49c8f..7b87683af 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": "TestingHugr",