From daa667161fec42b063248c4222a858348d5da21d Mon Sep 17 00:00:00 2001
From: Agustin Borgna <agustin.borgna@quantinuum.com>
Date: Fri, 1 Dec 2023 16:47:49 +0000
Subject: [PATCH] Derive DataflowOpTrait

---
 src/ops/custom.rs | 51 +++++++++++++++++++++++++++++------------------
 src/ops/leaf.rs   |  2 +-
 2 files changed, 33 insertions(+), 20 deletions(-)

diff --git a/src/ops/custom.rs b/src/ops/custom.rs
index 75bde65fc..9d2c0ab00 100644
--- a/src/ops/custom.rs
+++ b/src/ops/custom.rs
@@ -46,23 +46,6 @@ impl ExternalOp {
         qualify_name(res_id, op_name)
     }
 
-    /// A description of the external op.
-    pub fn description(&self) -> &str {
-        match self {
-            Self::Opaque(op) => op.description.as_str(),
-            Self::Extension(ext_op) => DataflowOpTrait::description(ext_op),
-        }
-    }
-
-    /// Note the case of an OpaqueOp without a signature should already
-    /// have been detected in [resolve_extension_ops]
-    pub fn dataflow_signature(&self) -> FunctionType {
-        match self {
-            Self::Opaque(op) => op.signature.clone(),
-            Self::Extension(ext_op) => ext_op.signature(),
-        }
-    }
-
     /// Downgrades this ExternalOp into an OpaqueOp
     pub fn as_opaque(self) -> OpaqueOp {
         match self {
@@ -100,6 +83,24 @@ impl From<ExternalOp> for OpType {
     }
 }
 
+impl DataflowOpTrait for ExternalOp {
+    const TAG: OpTag = OpTag::Leaf;
+
+    fn description(&self) -> &str {
+        match self {
+            Self::Opaque(op) => DataflowOpTrait::description(op),
+            Self::Extension(ext_op) => DataflowOpTrait::description(ext_op),
+        }
+    }
+
+    fn signature(&self) -> FunctionType {
+        match self {
+            Self::Opaque(op) => op.signature.clone(),
+            Self::Extension(ext_op) => ext_op.signature(),
+        }
+    }
+}
+
 /// An operation defined by an [OpDef] from a loaded [Extension].
 /// Note *not* Serializable: container ([ExternalOp]) is serialized as an [OpaqueOp] instead.
 ///
@@ -251,6 +252,18 @@ impl From<OpaqueOp> for OpType {
     }
 }
 
+impl DataflowOpTrait for OpaqueOp {
+    const TAG: OpTag = OpTag::Leaf;
+
+    fn description(&self) -> &str {
+        &self.description
+    }
+
+    fn signature(&self) -> FunctionType {
+        self.signature.clone()
+    }
+}
+
 /// Resolve serialized names of operations into concrete implementation (OpDefs) where possible
 #[allow(dead_code)]
 pub fn resolve_extension_ops(
@@ -352,8 +365,8 @@ mod test {
         );
         let op: ExternalOp = op.into();
         assert_eq!(op.name(), "res.op");
-        assert_eq!(op.description(), "desc");
+        assert_eq!(DataflowOpTrait::description(&op), "desc");
         assert_eq!(op.args(), &[TypeArg::Type { ty: USIZE_T }]);
-        assert_eq!(op.dataflow_signature(), sig);
+        assert_eq!(op.signature(), sig);
     }
 }
diff --git a/src/ops/leaf.rs b/src/ops/leaf.rs
index 6a1e5ac62..437124db7 100644
--- a/src/ops/leaf.rs
+++ b/src/ops/leaf.rs
@@ -170,7 +170,7 @@ impl DataflowOpTrait for LeafOp {
 
         match self {
             LeafOp::Noop { ty: typ } => FunctionType::new(vec![typ.clone()], vec![typ.clone()]),
-            LeafOp::CustomOp(ext) => ext.dataflow_signature(),
+            LeafOp::CustomOp(ext) => ext.signature(),
             LeafOp::MakeTuple { tys: types } => {
                 FunctionType::new(types.clone(), vec![Type::new_tuple(types.clone())])
             }