Test creating and instantiating an OpDef with a PolyFuncType scheme

CQCL · Nov 7, 2023 · 846f516 · 846f516
1 parent ebfec42
commit 846f516
Showing 1 changed file with 53 additions and 0 deletions.
diff --git 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<dyn std::error::Error>> {
+        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]),
+            &reg1,
+        )?;
+        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 }], &reg1)
+                    .unwrap(),
+            )),
+            dfg.input_wires(),
+        )?;
+        dfg.finish_hugr_with_outputs(
+            rev.outputs(),
+            &[PRELUDE.to_owned(), EXTENSION.to_owned(), e].into(),
+        )?;
+
+        Ok(())
+    }
+}