From 64351e21e99eabc1dc28b7b250103344be5001f2 Mon Sep 17 00:00:00 2001
From: Seyon Sivarajah <seyon.sivarajah@quantinuum.com>
Date: Mon, 27 Nov 2023 12:44:00 +0000
Subject: [PATCH] move op adding to extension

---
 src/extension/op_def.rs    | 15 +++++++++++++++
 src/extension/simple_op.rs | 34 +++++++++-------------------------
 2 files changed, 24 insertions(+), 25 deletions(-)

diff --git a/src/extension/op_def.rs b/src/extension/op_def.rs
index 246be5af2..155089a1b 100644
--- a/src/extension/op_def.rs
+++ b/src/extension/op_def.rs
@@ -435,6 +435,21 @@ impl Extension {
             Entry::Vacant(ve) => Ok(Arc::get_mut(ve.insert(Arc::new(op))).unwrap()),
         }
     }
+
+    pub fn add_op_enum(
+        &mut self,
+        op: &impl super::simple_op::OpEnum,
+    ) -> Result<&mut OpDef, ExtensionBuildError> {
+        let def = self.add_op(
+            op.name().into(),
+            op.description().to_string(),
+            op.def_signature(),
+        )?;
+
+        op.post_opdef(def);
+
+        Ok(def)
+    }
 }
 
 #[cfg(test)]
diff --git a/src/extension/simple_op.rs b/src/extension/simple_op.rs
index 72649def2..5c7a99488 100644
--- a/src/extension/simple_op.rs
+++ b/src/extension/simple_op.rs
@@ -76,36 +76,12 @@ pub trait OpEnum: FromStr + IntoEnumIterator + IntoStaticSt {
     /// Try to load one of the operations of this set from an [OpDef].
     fn from_op_def(op_def: &OpDef, args: &[TypeArg]) -> Result<Self, Self::LoadError>;
 
-    /// Add an operation to an extension.
-    fn add_to_extension<'e>(
-        &self,
-        ext: &'e mut Extension,
-    ) -> Result<&'e OpDef, ExtensionBuildError> {
-        let def = ext.add_op(
-            self.name().into(),
-            self.description().to_string(),
-            self.def_signature(),
-        )?;
-
-        self.post_opdef(def);
-
-        Ok(def)
-    }
-
     /// Iterator over all operations in the set. Non-trivial variants will have
     /// default values used for the members.
     fn all_variants() -> <Self as IntoEnumIterator>::Iterator {
         <Self as IntoEnumIterator>::iter()
     }
 
-    /// load all variants of a `SimpleOpEnum` in to an extension as op defs.
-    fn load_all_ops(extension: &mut Extension) -> Result<(), ExtensionBuildError> {
-        for op in Self::all_variants() {
-            op.add_to_extension(extension)?;
-        }
-        Ok(())
-    }
-
     /// Try to instantiate a variant from an [OpType]. Default behaviour assumes
     /// an [ExtensionOp] and loads from the name.
     fn from_optype(op: &OpType) -> Option<Self> {
@@ -124,6 +100,14 @@ pub trait OpEnum: FromStr + IntoEnumIterator + IntoStaticSt {
             op_enum: self,
         }
     }
+
+    /// load all variants of a [OpEnum] in to an extension as op defs.
+    fn load_all_ops(extension: &mut Extension) -> Result<(), ExtensionBuildError> {
+        for op in Self::all_variants() {
+            extension.add_op_enum(&op)?;
+        }
+        Ok(())
+    }
 }
 
 pub struct RegisteredEnum<'r, T> {
@@ -208,7 +192,7 @@ mod test {
         let ext_name = ExtensionId::new("dummy").unwrap();
         let mut e = Extension::new(ext_name.clone());
 
-        o.add_to_extension(&mut e).unwrap();
+        e.add_op_enum(&o).unwrap();
 
         assert_eq!(
             DummyEnum::from_op_def(e.get_op(o.name()).unwrap(), &[]).unwrap(),