diff --git a/src/extension/op_def.rs b/src/extension/op_def.rs index b50db4c38..e194c169f 100644 --- a/src/extension/op_def.rs +++ b/src/extension/op_def.rs @@ -274,14 +274,9 @@ 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( @@ -289,16 +284,7 @@ impl<'e> OpDefMut<'e> { k: impl ToString, v: serde_yaml::Value, ) -> Option { - 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, - ) -> 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, - ) -> 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();