hugr_core/builder/
module.rs

1use super::{
2    BuildError, Container,
3    build_traits::HugrBuilder,
4    dataflow::{DFGBuilder, FunctionBuilder},
5};
6
7use crate::hugr::internal::HugrMutInternals;
8use crate::hugr::views::HugrView;
9use crate::ops;
10use crate::ops::handle::{AliasID, FuncID, NodeHandle};
11use crate::types::{PolyFuncType, Type, TypeBound};
12use crate::{Hugr, Node, Visibility};
13use crate::{hugr::ValidationError, ops::FuncDefn};
14
15use smol_str::SmolStr;
16
17/// Builder for a HUGR module.
18#[derive(Debug, Default, Clone, PartialEq)]
19pub struct ModuleBuilder<T>(pub(super) T);
20
21impl<T: AsMut<Hugr> + AsRef<Hugr>> Container for ModuleBuilder<T> {
22    #[inline]
23    fn container_node(&self) -> Node {
24        self.0.as_ref().module_root()
25    }
26
27    #[inline]
28    fn hugr_mut(&mut self) -> &mut Hugr {
29        self.0.as_mut()
30    }
31
32    fn hugr(&self) -> &Hugr {
33        self.0.as_ref()
34    }
35}
36
37impl ModuleBuilder<Hugr> {
38    /// Begin building a new module.
39    #[must_use]
40    pub fn new() -> Self {
41        Self::default()
42    }
43}
44
45impl HugrBuilder for ModuleBuilder<Hugr> {
46    fn finish_hugr(self) -> Result<Hugr, ValidationError<Node>> {
47        self.0.validate()?;
48        Ok(self.0)
49    }
50}
51
52impl<T: AsMut<Hugr> + AsRef<Hugr>> ModuleBuilder<T> {
53    /// Continue building a module from an existing hugr.
54    #[must_use]
55    pub fn with_hugr(hugr: T) -> Self {
56        ModuleBuilder(hugr)
57    }
58
59    /// Replace a [`ops::FuncDecl`] with [`ops::FuncDefn`] and return a builder for
60    /// the defining graph.
61    ///
62    /// # Errors
63    ///
64    /// This function will return an error if there is an error in adding the
65    /// [`crate::ops::OpType::FuncDefn`] node.
66    pub fn define_declaration(
67        &mut self,
68        f_id: &FuncID<false>,
69    ) -> Result<FunctionBuilder<&mut Hugr>, BuildError> {
70        let f_node = f_id.node();
71        let opty = self.hugr_mut().optype_mut(f_node);
72        let ops::OpType::FuncDecl(decl) = opty else {
73            return Err(BuildError::UnexpectedType {
74                node: f_node,
75                op_desc: "crate::ops::OpType::FuncDecl",
76            });
77        };
78
79        let body = decl.signature().body().clone();
80        *opty = ops::FuncDefn::new_vis(
81            decl.func_name(),
82            decl.signature().clone(),
83            decl.visibility().clone(),
84        )
85        .into();
86
87        let db = DFGBuilder::create_with_io(self.hugr_mut(), f_node, body)?;
88        Ok(FunctionBuilder::from_dfg_builder(db))
89    }
90
91    /// Add a [`ops::FuncDefn`] node of the specified visibility.
92    /// Returns a builder to define the function body graph.
93    ///
94    /// # Errors
95    ///
96    /// This function will return an error if there is an error in adding the
97    /// [`ops::FuncDefn`] node.
98    pub fn define_function_vis(
99        &mut self,
100        name: impl Into<String>,
101        signature: impl Into<PolyFuncType>,
102        visibility: Visibility,
103    ) -> Result<FunctionBuilder<&mut Hugr>, BuildError> {
104        self.define_function_op(FuncDefn::new_vis(name, signature, visibility))
105    }
106
107    fn define_function_op(
108        &mut self,
109        op: FuncDefn,
110    ) -> Result<FunctionBuilder<&mut Hugr>, BuildError> {
111        let body = op.signature().body().clone();
112        let f_node = self.add_child_node(op);
113
114        // Add the extensions used by the function types.
115        self.use_extensions(
116            body.used_extensions().unwrap_or_else(|e| {
117                panic!("Build-time signatures should have valid extensions. {e}")
118            }),
119        );
120
121        let db = DFGBuilder::create_with_io(self.hugr_mut(), f_node, body)?;
122        Ok(FunctionBuilder::from_dfg_builder(db))
123    }
124
125    /// Declare a [Visibility::Public] function with `signature` and return a handle to the declaration.
126    ///
127    /// # Errors
128    ///
129    /// This function will return an error if there is an error in adding the
130    /// [`crate::ops::OpType::FuncDecl`] node.
131    pub fn declare(
132        &mut self,
133        name: impl Into<String>,
134        signature: PolyFuncType,
135    ) -> Result<FuncID<false>, BuildError> {
136        self.declare_vis(name, signature, Visibility::Public)
137    }
138
139    /// Declare a function with the specified `signature` and [Visibility],
140    /// and return a handle to the declaration.
141    ///
142    /// # Errors
143    ///
144    /// This function will return an error if there is an error in adding the
145    /// [`crate::ops::OpType::FuncDecl`] node.
146    pub fn declare_vis(
147        &mut self,
148        name: impl Into<String>,
149        signature: PolyFuncType,
150        visibility: Visibility,
151    ) -> Result<FuncID<false>, BuildError> {
152        let body = signature.body().clone();
153        // TODO add param names to metadata
154        let declare_n = self.add_child_node(ops::FuncDecl::new_vis(name, signature, visibility));
155
156        // Add the extensions used by the function types.
157        self.use_extensions(
158            body.used_extensions().unwrap_or_else(|e| {
159                panic!("Build-time signatures should have valid extensions. {e}")
160            }),
161        );
162
163        Ok(declare_n.into())
164    }
165
166    /// Adds a [`ops::FuncDefn`] node and returns a builder to define the function
167    /// body graph. The function will be private. (See [Self::define_function_vis].)
168    ///
169    /// # Errors
170    ///
171    /// This function will return an error if there is an error in adding the
172    /// [`ops::FuncDefn`] node.
173    pub fn define_function(
174        &mut self,
175        name: impl Into<String>,
176        signature: impl Into<PolyFuncType>,
177    ) -> Result<FunctionBuilder<&mut Hugr>, BuildError> {
178        self.define_function_op(FuncDefn::new(name, signature))
179    }
180
181    /// Add a [`crate::ops::OpType::AliasDefn`] node and return a handle to the Alias.
182    ///
183    /// # Errors
184    ///
185    /// Error in adding [`crate::ops::OpType::AliasDefn`] child node.
186    pub fn add_alias_def(
187        &mut self,
188        name: impl Into<SmolStr>,
189        typ: Type,
190    ) -> Result<AliasID<true>, BuildError> {
191        // TODO: add AliasDefn in other containers
192        // This is currently tricky as they are not connected to anything so do
193        // not appear in topological traversals.
194        // Could be fixed by removing single-entry requirement and sorting from
195        // every 0-input node.
196        let name: SmolStr = name.into();
197        let bound = typ.least_upper_bound();
198        let node = self.add_child_node(ops::AliasDefn {
199            name: name.clone(),
200            definition: typ,
201        });
202
203        Ok(AliasID::new(node, name, bound))
204    }
205
206    /// Add a [`crate::ops::OpType::AliasDecl`] node and return a handle to the Alias.
207    /// # Errors
208    ///
209    /// Error in adding [`crate::ops::OpType::AliasDecl`] child node.
210    pub fn add_alias_declare(
211        &mut self,
212        name: impl Into<SmolStr>,
213        bound: TypeBound,
214    ) -> Result<AliasID<false>, BuildError> {
215        let name: SmolStr = name.into();
216        let node = self.add_child_node(ops::AliasDecl {
217            name: name.clone(),
218            bound,
219        });
220
221        Ok(AliasID::new(node, name, bound))
222    }
223}
224
225#[cfg(test)]
226mod test {
227    use cool_asserts::assert_matches;
228
229    use crate::extension::prelude::usize_t;
230    use crate::{
231        builder::{Dataflow, DataflowSubContainer, test::n_identity},
232        types::Signature,
233    };
234
235    use super::*;
236    #[test]
237    fn basic_recurse() -> Result<(), BuildError> {
238        let build_result = {
239            let mut module_builder = ModuleBuilder::new();
240
241            let f_id = module_builder.declare(
242                "main",
243                Signature::new(vec![usize_t()], vec![usize_t()]).into(),
244            )?;
245
246            let mut f_build = module_builder.define_declaration(&f_id)?;
247            let call = f_build.call(&f_id, &[], f_build.input_wires())?;
248
249            f_build.finish_with_outputs(call.outputs())?;
250            module_builder.finish_hugr()
251        };
252        assert_matches!(build_result, Ok(_));
253        Ok(())
254    }
255
256    #[test]
257    fn simple_alias() -> Result<(), BuildError> {
258        let build_result = {
259            let mut module_builder = ModuleBuilder::new();
260
261            let qubit_state_type =
262                module_builder.add_alias_declare("qubit_state", TypeBound::Linear)?;
263
264            let f_build = module_builder.define_function(
265                "main",
266                Signature::new(
267                    vec![qubit_state_type.get_alias_type()],
268                    vec![qubit_state_type.get_alias_type()],
269                ),
270            )?;
271            n_identity(f_build)?;
272            module_builder.finish_hugr()
273        };
274        assert_matches!(build_result, Ok(_));
275        Ok(())
276    }
277
278    #[test]
279    fn builder_from_existing() -> Result<(), BuildError> {
280        let hugr = Hugr::new();
281
282        let fn_builder = FunctionBuilder::with_hugr(hugr, "main", Signature::new_endo(vec![]))?;
283        let mut hugr = fn_builder.finish_hugr()?;
284
285        let mut module_builder = ModuleBuilder::with_hugr(&mut hugr);
286        module_builder.declare("other", Signature::new_endo(vec![]).into())?;
287
288        hugr.validate()?;
289
290        Ok(())
291    }
292}
hugr_core/builder/module.rs

hugr_core/builder/
module.rs
