diff --git a/src/extension/op_def.rs b/src/extension/op_def.rs index 9ef17f73e0..f89da6c188 100644 --- a/src/extension/op_def.rs +++ b/src/extension/op_def.rs @@ -332,3 +332,56 @@ impl Extension { ) } } + +#[cfg(test)] +mod test { + use smol_str::SmolStr; + + use crate::builder::{DFGBuilder, Dataflow, DataflowHugr}; + use crate::extension::prelude::USIZE_T; + use crate::extension::{PRELUDE}; + use crate::ops::custom::ExternalOp; + use crate::ops::LeafOp; + use crate::std_extensions::collections::{EXTENSION, LIST_TYPENAME}; + use crate::types::Type; + use crate::types::{type_param::TypeParam, FunctionType, PolyFuncType, TypeArg, TypeBound}; + use crate::{const_extension_ids, Extension}; + + const_extension_ids! { + const EXT_ID: ExtensionId = "MyExt"; + } + + #[test] + fn op_def_with_type_scheme() -> Result<(), Box> { + let reg1 = [PRELUDE.to_owned(), EXTENSION.to_owned()].into(); + let list_def = EXTENSION.get_type(&LIST_TYPENAME).unwrap(); + let mut e = Extension::new(EXT_ID); + const TP: TypeParam = TypeParam::Type(TypeBound::Any); + let list_of_var = + Type::new_extension(list_def.instantiate(vec![TypeArg::new_var_use(0, TP)])?); + const OP_NAME: SmolStr = SmolStr::new_inline("Reverse"); + let type_scheme = PolyFuncType::new_validated( + vec![TP], + FunctionType::new_linear(vec![list_of_var]), + ®1, + )?; + e.add_op_type_scheme(OP_NAME, "".into(), Default::default(), vec![], type_scheme)?; + + let list_usize = + Type::new_extension(list_def.instantiate(vec![TypeArg::Type { ty: USIZE_T }])?); + let mut dfg = DFGBuilder::new(FunctionType::new_linear(vec![list_usize]))?; + let rev = dfg.add_dataflow_op( + LeafOp::from(ExternalOp::Extension( + e.instantiate_extension_op(&OP_NAME, vec![TypeArg::Type { ty: USIZE_T }], ®1) + .unwrap(), + )), + dfg.input_wires(), + )?; + dfg.finish_hugr_with_outputs( + rev.outputs(), + &[PRELUDE.to_owned(), EXTENSION.to_owned(), e].into(), + )?; + + Ok(()) + } +}