refactor!: Merge CustomOp and ExternalOp. (#923)

Followup of #922 .

Drops `ExternalOp` and makes `CustomOp` an enum with two boxed variants
for `OpaqueOp` and `ExtensionOp`.
The schema remains the same.

---------

Co-authored-by: Seyon Sivarajah <seyon.sivarajah@quantinuum.com>

hugr/src/algorithm/const_fold.rs

@@ -65,7 +65,7 @@ pub fn fold_leaf_op(op: &OpType, consts: &[(IncomingPort, Const)]) -> ConstFoldR
         )
         .unwrap()]),
         OpType::CustomOp(op) => {
-            let ext_op = op.as_ref().as_extension_op()?;
+            let ext_op = op.as_extension_op()?;
             ext_op.constant_fold(consts)
         }
         _ => None,

hugr/src/builder/circuit.rs

@@ -278,13 +278,13 @@ mod test {
 
     #[test]
     fn with_nonlinear_and_outputs() {
-        let my_custom_op = CustomOp::new(crate::ops::custom::ExternalOp::Opaque(OpaqueOp::new(
+        let my_custom_op = CustomOp::new_opaque(OpaqueOp::new(
             "MissingRsrc".try_into().unwrap(),
             "MyOp",
             "unknown op".to_string(),
             vec![],
             FunctionType::new(vec![QB, NAT], vec![QB]),
-        )));
+        ));
         let build_res = build_main(
             FunctionType::new(type_row![QB, QB, NAT], type_row![QB, QB, BOOL_T]).into(),
             |mut f_build| {

hugr/src/extension/infer/test.rs

@@ -9,7 +9,7 @@ use crate::extension::ExtensionId;
 use crate::extension::{prelude::PRELUDE_REGISTRY, ExtensionSet};
 use crate::hugr::{Hugr, HugrMut, HugrView, NodeType};
 use crate::macros::const_extension_ids;
-use crate::ops::custom::{ExternalOp, OpaqueOp};
+use crate::ops::custom::OpaqueOp;
 use crate::ops::{self, dataflow::IOTrait};
 use crate::ops::{CustomOp, Lift, OpType};
 #[cfg(feature = "extension_inference")]
@@ -439,8 +439,7 @@ fn extension_adding_sequence() -> Result<(), Box<dyn Error>> {
 }
 
 fn make_opaque(extension: impl Into<ExtensionId>, signature: FunctionType) -> CustomOp {
-    let opaque = ops::custom::OpaqueOp::new(extension.into(), "", "".into(), vec![], signature);
-    ops::custom::ExternalOp::from(opaque).into()
+    ops::custom::OpaqueOp::new(extension.into(), "", "".into(), vec![], signature).into()
 }
 
 fn make_block(
@@ -930,21 +929,21 @@ fn plus_on_self() -> Result<(), Box<dyn std::error::Error>> {
     // While https://github.com/CQCL/hugr/issues/388 is unsolved,
     // most operations have empty extension_reqs (not including their own extension).
     // Define some that do.
-    let binop = CustomOp::new(ExternalOp::Opaque(OpaqueOp::new(
+    let binop = CustomOp::new_opaque(OpaqueOp::new(
         ext.clone(),
         "2qb_op",
         String::new(),
         vec![],
         ft,
-    )));
+    ));
     let unary_sig = FunctionType::new_endo(type_row![QB_T]).with_extension_delta(ext.clone());
-    let unop = CustomOp::new(ExternalOp::Opaque(OpaqueOp::new(
+    let unop = CustomOp::new_opaque(OpaqueOp::new(
         ext,
         "1qb_op",
         String::new(),
         vec![],
         unary_sig,
-    )));
+    ));
     // Constrain q1,q2 as PLUS(ext1, inputs):
     let [q1, q2] = dfg
         .add_dataflow_op(binop.clone(), dfg.input_wires())?

hugr/src/extension/op_def.rs

@@ -476,7 +476,6 @@ mod test {
     use crate::extension::prelude::USIZE_T;
     use crate::extension::{ExtensionRegistry, ExtensionSet, PRELUDE};
     use crate::extension::{SignatureError, EMPTY_REG, PRELUDE_REGISTRY};
-    use crate::ops::custom::ExternalOp;
     use crate::ops::CustomOp;
     use crate::std_extensions::collections::{EXTENSION, LIST_TYPENAME};
     use crate::types::Type;
@@ -513,10 +512,10 @@ mod test {
             Type::new_extension(list_def.instantiate(vec![TypeArg::Type { ty: USIZE_T }])?);
         let mut dfg = DFGBuilder::new(FunctionType::new_endo(vec![list_usize]))?;
         let rev = dfg.add_dataflow_op(
-            CustomOp::new(ExternalOp::Extension(
+            CustomOp::new_extension(
                 e.instantiate_extension_op(&OP_NAME, vec![TypeArg::Type { ty: USIZE_T }], &reg)
                     .unwrap(),
-            )),
+            ),
             dfg.input_wires(),
         )?;
         dfg.finish_hugr_with_outputs(rev.outputs(), &reg)?;

hugr/src/extension/simple_op.rs

@@ -94,7 +94,7 @@ pub trait MakeExtensionOp: OpName {
     where
         Self: Sized,
     {
-        let ext: &ExtensionOp = op.as_custom_op()?.as_ref().as_extension_op()?;
+        let ext: &ExtensionOp = op.as_custom_op()?.as_extension_op()?;
         Self::from_extension_op(ext).ok()
     }
 

hugr/src/hugr.rs

@@ -203,8 +203,6 @@ impl Hugr {
 
     /// Infer extension requirements and add new information to `op_types` field
     /// (if the "extension_inference" feature is on; otherwise, do nothing)
-    ///
-    /// See [`infer_extensions`] for details on the "closure" value
     pub fn infer_extensions(&mut self) -> Result<(), InferExtensionError> {
         #[cfg(feature = "extension_inference")]
         {

hugr/src/hugr/rewrite/replace.rs

@@ -455,10 +455,10 @@ mod test {
     use crate::hugr::hugrmut::sealed::HugrMutInternals;
     use crate::hugr::rewrite::replace::WhichHugr;
     use crate::hugr::{HugrMut, NodeType, Rewrite};
-    use crate::ops::custom::{ExternalOp, OpaqueOp};
+    use crate::ops::custom::{CustomOp, OpaqueOp};
     use crate::ops::dataflow::DataflowOpTrait;
     use crate::ops::handle::{BasicBlockID, ConstID, NodeHandle};
-    use crate::ops::{self, Case, CustomOp, DataflowBlock, OpTag, OpType, DFG};
+    use crate::ops::{self, Case, DataflowBlock, OpTag, OpType, DFG};
     use crate::std_extensions::collections;
     use crate::types::{FunctionType, Type, TypeArg, TypeRow};
     use crate::{type_row, Direction, Hugr, HugrView, OutgoingPort};
@@ -650,13 +650,13 @@ mod test {
     fn test_invalid() -> Result<(), Box<dyn std::error::Error>> {
         let utou = FunctionType::new_endo(vec![USIZE_T]);
         let mk_op = |s| {
-            CustomOp::new(ExternalOp::Opaque(OpaqueOp::new(
+            CustomOp::new_opaque(OpaqueOp::new(
                 ExtensionId::new("unknown_ext").unwrap(),
                 s,
                 String::new(),
                 vec![],
                 utou.clone(),
-            )))
+            ))
         };
         let mut h = DFGBuilder::new(FunctionType::new(
             type_row![USIZE_T, BOOL_T],

hugr/src/hugr/serialize.rs

@@ -347,18 +347,7 @@ pub mod test {
         for node in new_hugr.nodes() {
             let new_op = new_hugr.get_optype(node);
             let old_op = h_canon.get_optype(node);
-            if let OpType::CustomOp(new_op) = new_op {
-                if let OpType::CustomOp(old_op) = old_op {
-                    assert_eq!(
-                        new_op.as_ref().clone().into_opaque(),
-                        old_op.as_ref().clone().into_opaque()
-                    );
-                } else {
-                    panic!("Expected old_op to be a custom op");
-                }
-            } else {
-                assert_eq!(new_op, old_op);
-            }
+            assert_eq!(new_op, old_op);
         }
 
         // Check that the graphs are equivalent up to port renumbering.

hugr/src/hugr/validate.rs

@@ -16,7 +16,7 @@ use crate::extension::{
 };
 
 use crate::ops::custom::CustomOpError;
-use crate::ops::custom::{resolve_opaque_op, ExternalOp};
+use crate::ops::custom::{resolve_opaque_op, CustomOp};
 use crate::ops::validate::{ChildrenEdgeData, ChildrenValidationError, EdgeValidationError};
 use crate::ops::{FuncDefn, OpTag, OpTrait, OpType, ValidateOp};
 use crate::types::type_param::TypeParam;
@@ -525,33 +525,32 @@ impl<'a, 'b> ValidationContext<'a, 'b> {
         // The op_type must be defined only in terms of type variables defined outside the node
         // TODO consider turning this match into a trait method?
         match op_type {
-            OpType::CustomOp(b) => {
+            OpType::CustomOp(op) => {
                 // Try to resolve serialized names to actual OpDefs in Extensions.
-                let temp: ExternalOp;
-                let external = b.as_ref();
-                let resolved = match external {
-                    ExternalOp::Opaque(op) => {
+                let temp: CustomOp;
+                let resolved = match op {
+                    CustomOp::Opaque(opaque) => {
                         // If resolve_extension_ops has been called first, this would always return Ok(None)
-                        match resolve_opaque_op(node, op, self.extension_registry)? {
+                        match resolve_opaque_op(node, opaque, self.extension_registry)? {
                             Some(exten) => {
-                                temp = ExternalOp::Extension(exten);
+                                temp = CustomOp::new_extension(exten);
                                 &temp
                             }
-                            None => external,
+                            None => op,
                         }
                     }
-                    ExternalOp::Extension(_) => external,
+                    CustomOp::Extension(_) => op,
                 };
                 // Check TypeArgs are valid, and if we can, fit the declared TypeParams
                 match resolved {
-                    ExternalOp::Extension(exten) => exten
+                    CustomOp::Extension(exten) => exten
                         .def()
                         .validate_args(exten.args(), self.extension_registry, var_decls)
                         .map_err(|cause| ValidationError::SignatureError { node, cause })?,
-                    ExternalOp::Opaque(opaq) => {
+                    CustomOp::Opaque(opaque) => {
                         // Best effort. Just check TypeArgs are valid in themselves, allowing any of them
                         // to contain type vars (we don't know how many are binary params, so accept if in doubt)
-                        for arg in opaq.args() {
+                        for arg in opaque.args() {
                             arg.validate(self.extension_registry, var_decls)
                                 .map_err(|cause| ValidationError::SignatureError { node, cause })?;
                         }
@@ -699,7 +698,7 @@ pub enum ValidationError {
     /// Error in a [CustomOp] serialized as an [Opaque]
     ///
     /// [CustomOp]: crate::ops::CustomOp
-    /// [Opaque]: crate::ops::custom::ExternalOp::Opaque
+    /// [Opaque]: crate::ops::CustomOp::Opaque
     #[error(transparent)]
     CustomOpError(#[from] CustomOpError),
 }

hugr/src/ops.rs

@@ -22,8 +22,9 @@ use enum_dispatch::enum_dispatch;
 
 pub use constant::Const;
 pub use controlflow::{BasicBlock, Case, Conditional, DataflowBlock, ExitBlock, TailLoop, CFG};
+pub use custom::CustomOp;
 pub use dataflow::{Call, CallIndirect, DataflowParent, Input, LoadConstant, Output, DFG};
-pub use leaf::{CustomOp, Lift, MakeTuple, Noop, Tag, UnpackTuple};
+pub use leaf::{Lift, MakeTuple, Noop, Tag, UnpackTuple};
 pub use module::{AliasDecl, AliasDefn, FuncDecl, FuncDefn, Module};
 pub use tag::OpTag;