get rid of OpDefMut

CQCL · Nov 20, 2023 · c89135c · c89135c
1 parent 9b07d79
commit c89135c
Showing 1 changed file with 19 additions and 44 deletions.
diff --git a/src/extension/op_def.rs b/src/extension/op_def.rs
@@ -274,31 +274,17 @@ impl OpDef {
         }
         Ok(())
     }
-}
-
-/// Wrapper to allow editing an OpDef before it is added to an extension.
-pub struct OpDefMut<'e>(&'e mut Extension, OpDef);
 
-impl<'e> OpDefMut<'e> {
     pub fn add_lower_func(&mut self, lower: LowerFunc) {
-        self.1.lower_funcs.push(lower);
+        self.lower_funcs.push(lower);
     }
 
     pub fn add_misc(
         &mut self,
         k: impl ToString,
         v: serde_yaml::Value,
     ) -> Option<serde_yaml::Value> {
-        self.1.misc.insert(k.to_string(), v)
-    }
-
-    pub fn finish(self) -> Result<&'e OpDef, ExtensionBuildError> {
-        let Self(e, op) = self;
-
-        match e.operations.entry(op.name.clone()) {
-            Entry::Occupied(_) => Err(ExtensionBuildError::OpDefExists(op.name)),
-            Entry::Vacant(ve) => Ok(ve.insert(Arc::new(op))),
-        }
+        self.misc.insert(k.to_string(), v)
     }
 }
 
@@ -309,28 +295,20 @@ impl Extension {
         name: SmolStr,
         description: String,
         signature_func: impl Into<SignatureFunc>,
-    ) -> Result<&OpDef, ExtensionBuildError> {
-        self.start_op_mut(name, description, signature_func)
-            .finish()
-    }
-
-    pub fn start_op_mut(
-        &mut self,
-        name: SmolStr,
-        description: String,
-        signature_func: impl Into<SignatureFunc>,
-    ) -> OpDefMut {
-        OpDefMut(
-            self,
-            OpDef {
-                extension: self.name.clone(),
-                name,
-                description,
-                signature_func: signature_func.into(),
-                misc: HashMap::new(),
-                lower_funcs: vec![],
-            },
-        )
+    ) -> Result<&mut OpDef, ExtensionBuildError> {
+        let op = Arc::new(OpDef {
+            extension: self.name.clone(),
+            name,
+            description,
+            signature_func: signature_func.into(),
+            misc: HashMap::new(),
+            lower_funcs: vec![],
+        });
+
+        match self.operations.entry(op.name.clone()) {
+            Entry::Occupied(_) => Err(ExtensionBuildError::OpDefExists(op.name.clone())),
+            Entry::Vacant(ve) => Ok(Arc::get_mut(ve.insert(op)).unwrap()),
+        }
     }
 }
 
@@ -379,12 +357,9 @@ mod test {
             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(vec![TP], FunctionType::new_endo(vec![list_of_var]));
-        let _def = {
-            let mut def_mut = e.start_op_mut(OP_NAME, "".into(), type_scheme);
-            def_mut.add_lower_func(LowerFunc::FixedHugr(ExtensionSet::new(), Hugr::default()));
-            def_mut.add_misc("key", Default::default());
-            def_mut.finish()
-        }?;
+        let def_mut = e.add_op(OP_NAME, "".into(), type_scheme)?;
+        def_mut.add_lower_func(LowerFunc::FixedHugr(ExtensionSet::new(), Hugr::default()));
+        def_mut.add_misc("key", Default::default());
         let reg =
             ExtensionRegistry::try_new([PRELUDE.to_owned(), EXTENSION.to_owned(), e]).unwrap();
         let e = reg.get(&EXT_ID).unwrap();