diff --git a/hugr-core/src/builder/dataflow.rs b/hugr-core/src/builder/dataflow.rs index cc6b1536fc..a5faa0aef9 100644 --- a/hugr-core/src/builder/dataflow.rs +++ b/hugr-core/src/builder/dataflow.rs @@ -259,6 +259,31 @@ impl FunctionBuilder { } } +impl + AsRef> FunctionBuilder { + /// Initialize a new function definition on the root module of an existing HUGR. + /// + /// The HUGR's entrypoint will **not** be modified. + /// + /// # Errors + /// + /// Error in adding DFG child nodes. + pub fn with_hugr( + mut hugr: B, + name: impl Into, + signature: impl Into, + ) -> Result { + let signature: PolyFuncType = signature.into(); + let body = signature.body().clone(); + let op = ops::FuncDefn::new(name, signature); + + let module = hugr.as_ref().module_root(); + let func = hugr.as_mut().add_node_with_parent(module, op); + + let db = DFGBuilder::create_with_io(hugr, func, body)?; + Ok(Self::from_dfg_builder(db)) + } +} + impl + AsRef, T> Container for DFGWrapper { #[inline] fn container_node(&self) -> Node { diff --git a/hugr-core/src/builder/module.rs b/hugr-core/src/builder/module.rs index b57d25c144..9e045bc2a4 100644 --- a/hugr-core/src/builder/module.rs +++ b/hugr-core/src/builder/module.rs @@ -51,6 +51,12 @@ impl HugrBuilder for ModuleBuilder { } impl + AsRef> ModuleBuilder { + /// Continue building a module from an existing hugr. + #[must_use] + pub fn with_hugr(hugr: T) -> Self { + ModuleBuilder(hugr) + } + /// Replace a [`ops::FuncDecl`] with [`ops::FuncDefn`] and return a builder for /// the defining graph. /// @@ -229,4 +235,19 @@ mod test { assert_matches!(build_result, Ok(_)); Ok(()) } + + #[test] + fn builder_from_existing() -> Result<(), BuildError> { + let hugr = Hugr::new(); + + let fn_builder = FunctionBuilder::with_hugr(hugr, "main", Signature::new_endo(vec![]))?; + let mut hugr = fn_builder.finish_hugr()?; + + let mut module_builder = ModuleBuilder::with_hugr(&mut hugr); + module_builder.declare("other", Signature::new_endo(vec![]).into())?; + + hugr.validate()?; + + Ok(()) + } }