From 5e9d7e14ae8bdedc5cfe018785539d82d5b3320e Mon Sep 17 00:00:00 2001
From: Seyon Sivarajah <seyon.sivarajah@quantinuum.com>
Date: Mon, 2 Sep 2024 16:57:37 +0100
Subject: [PATCH] feat: `is_ext` utility for checking equality against
 extension ops

Open to bikeshedding on name of `is_ext`
---
 hugr-core/src/extension/prelude/array.rs      | 3 +--
 hugr-core/src/hugr/rewrite/insert_identity.rs | 4 +---
 hugr-core/src/ops.rs                          | 8 ++++++++
 3 files changed, 10 insertions(+), 5 deletions(-)

diff --git a/hugr-core/src/extension/prelude/array.rs b/hugr-core/src/extension/prelude/array.rs
index a15bf23cc..abfbf7088 100644
--- a/hugr-core/src/extension/prelude/array.rs
+++ b/hugr-core/src/extension/prelude/array.rs
@@ -335,8 +335,7 @@ mod tests {
             };
             let op = def.to_concrete(ty, size);
             let optype: OpType = op.clone().into();
-            let new_op: ArrayOp = optype.cast().unwrap();
-            assert_eq!(new_op, op);
+            assert!(optype.is_ext(op));
         }
     }
 
diff --git a/hugr-core/src/hugr/rewrite/insert_identity.rs b/hugr-core/src/hugr/rewrite/insert_identity.rs
index d3a63bf97..afd763982 100644
--- a/hugr-core/src/hugr/rewrite/insert_identity.rs
+++ b/hugr-core/src/hugr/rewrite/insert_identity.rs
@@ -125,9 +125,7 @@ mod tests {
 
         assert_eq!(h.node_count(), 7);
 
-        let noop: Noop = h.get_optype(noop_node).cast().unwrap();
-
-        assert_eq!(noop, Noop(QB_T));
+        assert_eq!(h.get_optype(noop_node).cast(), Some(Noop(QB_T)));
 
         h.update_validate(&PRELUDE_REGISTRY).unwrap();
     }
diff --git a/hugr-core/src/ops.rs b/hugr-core/src/ops.rs
index 14305020d..a4b20cb7d 100644
--- a/hugr-core/src/ops.rs
+++ b/hugr-core/src/ops.rs
@@ -294,6 +294,14 @@ impl OpType {
         self.as_extension_op()
             .and_then(|o| T::from_extension_op(o).ok())
     }
+
+    /// Check if the operation is an extension operation by casting to it.
+    pub fn is_ext<T>(&self, op: T) -> bool
+    where
+        T: MakeExtensionOp + PartialEq,
+    {
+        self.cast::<T>() == Some(op)
+    }
 }
 
 /// Macro used by operations that want their