hugr_core/builder/
dataflow.rs

1//! Builders for dataflow regions.
2
3use itertools::Itertools;
4
5use super::build_traits::{HugrBuilder, SubContainer};
6use super::handle::BuildHandle;
7use super::{BuildError, Container, Dataflow, DfgID, FuncID};
8
9use std::marker::PhantomData;
10
11use crate::hugr::internal::HugrMutInternals;
12use crate::hugr::{HugrView, ValidationError};
13use crate::ops::handle::NodeHandle;
14use crate::ops::{self, DFG, FuncDefn, NamedOp, OpParent, OpType};
15use crate::types::{PolyFuncType, Signature, Type};
16use crate::{Direction, Hugr, IncomingPort, Node, OutgoingPort, Visibility, Wire, hugr::HugrMut};
17
18/// Builder for a [`ops::DFG`] node.
19#[derive(Debug, Clone, PartialEq)]
20pub struct DFGBuilder<T> {
21    pub(crate) base: T,
22    pub(crate) dfg_node: Node,
23    pub(crate) num_in_wires: usize,
24    pub(crate) num_out_wires: usize,
25}
26
27/// Error returned by [`DFGBuilder::add_input`] and [`DFGBuilder::add_output`].
28#[derive(Debug, Clone, PartialEq, derive_more::Display, derive_more::Error)]
29#[non_exhaustive]
30pub enum DFGAddPortError {
31    /// The parent optype is not a [`FuncDefn`], [`DFG`], or dataflow block so we cannot update its signature.
32    #[display("Adding port to an optype {op} is not supported.")]
33    ParentOpNotSupported {
34        /// The name of the unsupported operation.
35        op: String,
36    },
37}
38
39impl<T: AsMut<Hugr> + AsRef<Hugr>> DFGBuilder<T> {
40    /// Returns a new `DFGBuilder` with the given base and parent node.
41    ///
42    /// Sets up the input and output nodes of the region. If `parent` already has
43    /// input and output nodes, use [`DFGBuilder::create`] instead.
44    pub fn create_with_io(
45        mut base: T,
46        parent: Node,
47        signature: Signature,
48    ) -> Result<Self, BuildError> {
49        debug_assert_eq!(base.as_ref().children(parent).count(), 0);
50
51        let num_in_wires = signature.input_count();
52        let num_out_wires = signature.output_count();
53        let input = ops::Input {
54            types: signature.input().clone(),
55        };
56        let output = ops::Output {
57            types: signature.output().clone(),
58        };
59        base.as_mut().add_node_with_parent(parent, input);
60        base.as_mut().add_node_with_parent(parent, output);
61
62        Ok(Self {
63            base,
64            dfg_node: parent,
65            num_in_wires,
66            num_out_wires,
67        })
68    }
69
70    /// Returns a new `DFGBuilder` with the given base and parent node.
71    ///
72    /// The parent node may be any `DataflowParent` node.
73    ///
74    /// If `parent` doesn't have input and output nodes, use
75    /// [`DFGBuilder::create_with_io`] instead.
76    pub fn create(base: T, parent: Node) -> Result<Self, BuildError> {
77        let sig = base
78            .as_ref()
79            .get_optype(parent)
80            .inner_function_type()
81            .expect("DFG parent must have an inner function signature.");
82        let num_in_wires = sig.input_count();
83        let num_out_wires = sig.output_count();
84
85        Ok(Self {
86            base,
87            dfg_node: parent,
88            num_in_wires,
89            num_out_wires,
90        })
91    }
92
93    /// Add a new input to the dataflow being constructed.
94    ///
95    /// Updates the parent's optype to include the new input type.
96    ///
97    /// # Returns
98    ///
99    /// - The new wire from the input node.
100    ///
101    /// # Errors
102    ///
103    /// - [`DFGAddPortError::ParentOpNotSupported`] if the container optype is not a [`FuncDefn`], [`DFG`], or dataflow block so we cannot add an input.
104    ///   In this case, the Hugr will not be updated.
105    pub fn add_input(&mut self, input_type: Type) -> Result<Wire, DFGAddPortError> {
106        let [inp_node, _] = self.io();
107
108        // Update the parent's root type
109        if !self.update_parent_signature(|mut s| {
110            s.input.to_mut().push(input_type.clone());
111            s
112        }) {
113            return Err(DFGAddPortError::ParentOpNotSupported {
114                op: self
115                    .hugr()
116                    .get_optype(self.container_node())
117                    .name()
118                    .to_string(),
119            });
120        }
121
122        // Update the inner input node
123        let OpType::Input(inp) = self.hugr_mut().optype_mut(inp_node) else {
124            panic!("Input node {inp_node} is not an Input");
125        };
126        inp.types.to_mut().push(input_type);
127
128        let mut new_port = self.hugr_mut().add_ports(inp_node, Direction::Outgoing, 1);
129        let new_port = new_port.next().unwrap();
130
131        // The last port in an input/output node is an order edge port, so we must shift any connections to it.
132        let new_value_port: OutgoingPort = (new_port - 1).into();
133        let new_order_port: OutgoingPort = new_port.into();
134        let order_edge_targets = self
135            .hugr()
136            .linked_inputs(inp_node, new_value_port)
137            .collect_vec();
138        self.hugr_mut().disconnect(inp_node, new_value_port);
139        for (tgt_node, tgt_port) in order_edge_targets {
140            self.hugr_mut()
141                .connect(inp_node, new_order_port, tgt_node, tgt_port);
142        }
143
144        // Update the builder metadata
145        self.num_in_wires += 1;
146
147        Ok(self.input_wires().next_back().unwrap())
148    }
149
150    /// Add a new output to the function being constructed.
151    ///
152    /// Updates the parent's optype to include the new input type, if possible.
153    /// If the container optype is not a [`FuncDefn`], [`DFG`], or dataflow
154    /// block, the optype will not be updated.
155    ///
156    /// # Errors
157    ///
158    /// - [`DFGAddPortError::ParentOpNotSupported`] if the container optype is not a [`FuncDefn`], [`DFG`], or dataflow block so we cannot add an input.
159    ///   In this case, the Hugr will not be updated.
160    pub fn add_output(&mut self, output_type: Type) -> Result<(), DFGAddPortError> {
161        let [_, out_node] = self.io();
162
163        // Update the parent's root type
164        if !self.update_parent_signature(|mut s| {
165            s.output.to_mut().push(output_type.clone());
166            s
167        }) {
168            return Err(DFGAddPortError::ParentOpNotSupported {
169                op: self
170                    .hugr()
171                    .get_optype(self.container_node())
172                    .name()
173                    .to_string(),
174            });
175        }
176
177        // Update the inner input node
178        let OpType::Output(out) = self.hugr_mut().optype_mut(out_node) else {
179            panic!("Output node {out_node} is not an Output");
180        };
181        out.types.to_mut().push(output_type);
182
183        let mut new_port = self.hugr_mut().add_ports(out_node, Direction::Incoming, 1);
184        let new_port = new_port.next().unwrap();
185
186        // The last port in an input/output node is an order edge port, so we must shift any connections to it.
187        let new_value_port: IncomingPort = (new_port - 1).into();
188        let new_order_port: IncomingPort = new_port.into();
189        let order_edge_sources = self
190            .hugr()
191            .linked_outputs(out_node, new_value_port)
192            .collect_vec();
193        self.hugr_mut().disconnect(out_node, new_value_port);
194        for (src_node, src_port) in order_edge_sources {
195            self.hugr_mut()
196                .connect(src_node, src_port, out_node, new_order_port);
197        }
198
199        // Update the builder metadata
200        self.num_out_wires += 1;
201
202        Ok(())
203    }
204
205    /// Update the container's parent signature, if it is a function definition, DFG region, or dataflow block.
206    ///
207    /// Internal function used in [`add_input`] and [`add_output`].
208    ///
209    /// Does not update the input and output nodes.
210    ///
211    /// # Returns
212    ///
213    /// - `true` if the parent signature was updated.
214    /// - `false` if the parent optype is not a [`FuncDefn`], [`DFG`], or dataflow block so we cannot update the signature.
215    fn update_parent_signature(&mut self, f: impl FnOnce(Signature) -> Signature) -> bool {
216        let parent = self.container_node();
217
218        match self.hugr_mut().optype_mut(parent) {
219            ops::OpType::FuncDefn(fd) => {
220                let mut sig = std::mem::take(fd.signature_mut());
221                let body = std::mem::take(sig.body_mut());
222                *sig.body_mut() = f(body);
223                *fd.signature_mut() = sig;
224            }
225            ops::OpType::DFG(dfg) => {
226                let sig = std::mem::take(&mut dfg.signature);
227                dfg.signature = f(sig);
228            }
229            ops::OpType::DataflowBlock(dfb) => {
230                let inp = std::mem::take(&mut dfb.inputs);
231                let other_outputs = std::mem::take(&mut dfb.other_outputs);
232                let sig = f(Signature::new(inp, other_outputs));
233                dfb.inputs = sig.input;
234                dfb.other_outputs = sig.output;
235            }
236            _ => return false,
237        }
238        true
239    }
240}
241
242impl DFGBuilder<Hugr> {
243    /// Begin building a new DFG-rooted HUGR given its inputs, outputs,
244    /// and extension delta.
245    ///
246    /// # Errors
247    ///
248    /// Error in adding DFG child nodes.
249    pub fn new(signature: Signature) -> Result<DFGBuilder<Hugr>, BuildError> {
250        let dfg_op: DFG = ops::DFG {
251            signature: signature.clone(),
252        };
253        let base = Hugr::new_with_entrypoint(dfg_op).expect("DFG entrypoint should be valid");
254        let root = base.entrypoint();
255        DFGBuilder::create_with_io(base, root, signature)
256    }
257}
258
259impl HugrBuilder for DFGBuilder<Hugr> {
260    fn finish_hugr(self) -> Result<Hugr, ValidationError<Node>> {
261        self.base.validate()?;
262        Ok(self.base)
263    }
264}
265
266impl<T: AsMut<Hugr> + AsRef<Hugr>> Container for DFGBuilder<T> {
267    #[inline]
268    fn container_node(&self) -> Node {
269        self.dfg_node
270    }
271
272    #[inline]
273    fn hugr_mut(&mut self) -> &mut Hugr {
274        self.base.as_mut()
275    }
276
277    #[inline]
278    fn hugr(&self) -> &Hugr {
279        self.base.as_ref()
280    }
281}
282
283impl<T: AsMut<Hugr> + AsRef<Hugr>> SubContainer for DFGBuilder<T> {
284    type ContainerHandle = BuildHandle<DfgID>;
285    #[inline]
286    fn finish_sub_container(self) -> Result<Self::ContainerHandle, BuildError> {
287        Ok((self.dfg_node, self.num_out_wires).into())
288    }
289}
290
291impl<T: AsMut<Hugr> + AsRef<Hugr>> Dataflow for DFGBuilder<T> {
292    #[inline]
293    fn num_inputs(&self) -> usize {
294        self.num_in_wires
295    }
296}
297
298/// Wrapper around [`DFGBuilder`] used to build other dataflow regions.
299// Stores option of DFGBuilder so it can be taken out without moving.
300#[derive(Debug, Clone, PartialEq)]
301pub struct DFGWrapper<B, T>(DFGBuilder<B>, PhantomData<T>);
302
303impl<B, T: NodeHandle> DFGWrapper<B, T> {
304    /// Wrap a [`DFGBuilder`] into a [`DFGWrapper`], with the given handle type.
305    ///
306    /// The caller must ensure that the DFGBuilder's parent node has the correct
307    /// handle type given by `T::TAG`.
308    pub fn from_dfg_builder(db: DFGBuilder<B>) -> Self {
309        Self(db, PhantomData)
310    }
311
312    /// Unwrap the [`DFGWrapper`] into a [`DFGBuilder`].
313    pub fn into_dfg_builder(self) -> DFGBuilder<B> {
314        self.0
315    }
316}
317
318impl<B: AsMut<Hugr> + AsRef<Hugr>, T> DFGWrapper<B, T> {
319    /// Add a new input to the dataflow being constructed.
320    ///
321    /// Updates the parent's optype to include the new input type, if possible.
322    /// If the container optype is not a [`FuncDefn`] or a [`DFG`], the optype
323    /// will not be updated.
324    ///
325    /// # Returns
326    ///
327    /// - The new wire from the input node.
328    ///
329    /// # Errors
330    ///
331    /// - [`DFGAddPortError::ParentOpNotSupported`] if the container optype is
332    ///   not a [`FuncDefn`], [`DFG`], or dataflow block so we cannot add an
333    ///   input. In this case, the Hugr will not be updated.
334    pub fn add_input(&mut self, input_type: Type) -> Result<Wire, DFGAddPortError> {
335        self.0.add_input(input_type)
336    }
337
338    /// Add a new output to the dataflow being constructed.
339    ///
340    /// Updates the parent's optype to include the new output type, if possible.
341    /// If the container optype is not a [`FuncDefn`] or a [`DFG`], the optype
342    /// will not be updated.
343    ///
344    /// # Errors
345    ///
346    /// - [`DFGAddPortError::ParentOpNotSupported`] if the container optype is
347    ///   not a [`FuncDefn`], [`DFG`], or dataflow block so we cannot add an
348    ///   input. In this case, the Hugr will not be updated.
349    pub fn add_output(&mut self, output_type: Type) -> Result<(), DFGAddPortError> {
350        self.0.add_output(output_type)
351    }
352}
353
354/// Builder for a [`ops::FuncDefn`] node
355pub type FunctionBuilder<B> = DFGWrapper<B, BuildHandle<FuncID<true>>>;
356
357impl FunctionBuilder<Hugr> {
358    /// Initialize a builder for a [`FuncDefn`]-rooted HUGR; the function will
359    /// be private. (See also [Self::new_vis].)
360    ///
361    /// # Errors
362    ///
363    /// Error in adding DFG child nodes.
364    pub fn new(
365        name: impl Into<String>,
366        signature: impl Into<PolyFuncType>,
367    ) -> Result<Self, BuildError> {
368        Self::new_with_op(FuncDefn::new(name, signature))
369    }
370
371    /// Initialize a builder for a FuncDefn-rooted HUGR, with the specified
372    /// [Visibility].
373    ///
374    /// # Errors
375    ///
376    /// Error in adding DFG child nodes.
377    pub fn new_vis(
378        name: impl Into<String>,
379        signature: impl Into<PolyFuncType>,
380        visibility: Visibility,
381    ) -> Result<Self, BuildError> {
382        Self::new_with_op(FuncDefn::new_vis(name, signature, visibility))
383    }
384
385    fn new_with_op(op: FuncDefn) -> Result<Self, BuildError> {
386        let body = op.signature().body().clone();
387
388        let base = Hugr::new_with_entrypoint(op).expect("FuncDefn entrypoint should be valid");
389        let root = base.entrypoint();
390
391        let db = DFGBuilder::create_with_io(base, root, body)?;
392        Ok(Self::from_dfg_builder(db))
393    }
394}
395
396impl<B: AsMut<Hugr> + AsRef<Hugr>> FunctionBuilder<B> {
397    /// Initialize a new function definition on the root module of an existing HUGR.
398    ///
399    /// The HUGR's entrypoint will **not** be modified.
400    ///
401    /// # Errors
402    ///
403    /// Error in adding DFG child nodes.
404    pub fn with_hugr(
405        mut hugr: B,
406        name: impl Into<String>,
407        signature: impl Into<PolyFuncType>,
408    ) -> Result<Self, BuildError> {
409        let signature: PolyFuncType = signature.into();
410        let body = signature.body().clone();
411        let op = ops::FuncDefn::new(name, signature);
412
413        let module = hugr.as_ref().module_root();
414        let func = hugr.as_mut().add_node_with_parent(module, op);
415
416        let db = DFGBuilder::create_with_io(hugr, func, body)?;
417        Ok(Self::from_dfg_builder(db))
418    }
419}
420
421impl<B: AsMut<Hugr> + AsRef<Hugr>, T> Container for DFGWrapper<B, T> {
422    #[inline]
423    fn container_node(&self) -> Node {
424        self.0.container_node()
425    }
426
427    #[inline]
428    fn hugr_mut(&mut self) -> &mut Hugr {
429        self.0.hugr_mut()
430    }
431
432    #[inline]
433    fn hugr(&self) -> &Hugr {
434        self.0.hugr()
435    }
436}
437
438impl<B: AsMut<Hugr> + AsRef<Hugr>, T> Dataflow for DFGWrapper<B, T> {
439    #[inline]
440    fn num_inputs(&self) -> usize {
441        self.0.num_inputs()
442    }
443}
444
445impl<B: AsMut<Hugr> + AsRef<Hugr>, T: From<BuildHandle<DfgID>>> SubContainer for DFGWrapper<B, T> {
446    type ContainerHandle = T;
447
448    #[inline]
449    fn finish_sub_container(self) -> Result<Self::ContainerHandle, BuildError> {
450        self.0.finish_sub_container().map(Into::into)
451    }
452}
453
454impl<T> HugrBuilder for DFGWrapper<Hugr, T> {
455    fn finish_hugr(self) -> Result<Hugr, ValidationError<Node>> {
456        self.0.finish_hugr()
457    }
458}
459
460#[cfg(test)]
461pub(crate) mod test {
462    use cool_asserts::assert_matches;
463    use rstest::rstest;
464    use serde_json::json;
465    use std::collections::HashMap;
466
467    use crate::builder::build_traits::DataflowHugr;
468    use crate::builder::test::dfg_calling_defn_decl;
469    use crate::builder::{
470        BuilderWiringError, CFGBuilder, DataflowSubContainer, ModuleBuilder, TailLoopBuilder,
471        endo_sig, inout_sig,
472    };
473    use crate::extension::SignatureError;
474    use crate::extension::prelude::Noop;
475    use crate::extension::prelude::{bool_t, qb_t, usize_t};
476    use crate::hugr::linking::NodeLinkingDirective;
477    use crate::hugr::validate::InterGraphEdgeError;
478    use crate::ops::{FuncDecl, FuncDefn, OpParent, OpTag, OpTrait, Value, handle::NodeHandle};
479
480    use crate::std_extensions::logic::test::and_op;
481    use crate::types::type_param::TypeParam;
482    use crate::types::{EdgeKind, FuncValueType, RowVariable, Signature, Type, TypeBound, TypeRV};
483    use crate::utils::test_quantum_extension::h_gate;
484    use crate::{Wire, builder::test::n_identity, type_row};
485
486    use super::super::test::simple_dfg_hugr;
487    use super::*;
488    #[test]
489    fn nested_identity() -> Result<(), BuildError> {
490        let build_result = {
491            let mut outer_builder = DFGBuilder::new(endo_sig(vec![usize_t(), qb_t()]))?;
492
493            let [int, qb] = outer_builder.input_wires_arr();
494
495            let q_out = outer_builder.add_dataflow_op(h_gate(), vec![qb])?;
496
497            let inner_builder = outer_builder.dfg_builder_endo([(usize_t(), int)])?;
498            let inner_id = n_identity(inner_builder)?;
499
500            outer_builder.finish_hugr_with_outputs(inner_id.outputs().chain(q_out.outputs()))
501        };
502
503        assert_eq!(build_result.err(), None);
504
505        Ok(())
506    }
507
508    // Scaffolding for copy insertion tests
509    fn copy_scaffold<F>(f: F, msg: &'static str) -> Result<(), BuildError>
510    where
511        F: FnOnce(&mut DFGBuilder<Hugr>) -> Result<(), BuildError>,
512    {
513        let build_result = {
514            let mut builder = DFGBuilder::new(inout_sig(bool_t(), vec![bool_t(), bool_t()]))?;
515
516            f(&mut builder)?;
517
518            builder.finish_hugr()
519        };
520        assert_matches!(build_result, Ok(_), "Failed on example: {}", msg);
521
522        Ok(())
523    }
524    #[test]
525    fn copy_insertion() -> Result<(), BuildError> {
526        copy_scaffold(
527            |f_build| {
528                let [b1] = f_build.input_wires_arr();
529                f_build.set_outputs([b1, b1])
530            },
531            "Copy input and output",
532        )?;
533
534        copy_scaffold(
535            |f_build| {
536                let [b1] = f_build.input_wires_arr();
537                let xor = f_build.add_dataflow_op(and_op(), [b1, b1])?;
538                f_build.set_outputs([xor.out_wire(0), b1])
539            },
540            "Copy input and use with binary function",
541        )?;
542
543        copy_scaffold(
544            |f_build| {
545                let [b1] = f_build.input_wires_arr();
546                let xor1 = f_build.add_dataflow_op(and_op(), [b1, b1])?;
547                let xor2 = f_build.add_dataflow_op(and_op(), [b1, xor1.out_wire(0)])?;
548                f_build.set_outputs([xor2.out_wire(0), b1])
549            },
550            "Copy multiple times",
551        )?;
552
553        Ok(())
554    }
555
556    #[test]
557    fn copy_insertion_qubit() {
558        let builder = || {
559            let mut module_builder = ModuleBuilder::new();
560
561            let f_build = module_builder
562                .define_function("main", Signature::new(vec![qb_t()], vec![qb_t(), qb_t()]))?;
563
564            let [q1] = f_build.input_wires_arr();
565            f_build.finish_with_outputs([q1, q1])?;
566
567            Ok(module_builder.finish_hugr()?)
568        };
569
570        assert_matches!(
571            builder(),
572            Err(BuildError::OutputWiring {
573                error: BuilderWiringError::NoCopyLinear { typ, .. },
574                ..
575            })
576            if *typ == qb_t()
577        );
578    }
579
580    #[test]
581    fn simple_inter_graph_edge() {
582        let builder = || -> Result<Hugr, BuildError> {
583            let mut f_build =
584                FunctionBuilder::new("main", Signature::new(vec![bool_t()], vec![bool_t()]))?;
585
586            let [i1] = f_build.input_wires_arr();
587            let noop = f_build.add_dataflow_op(Noop(bool_t()), [i1])?;
588            let i1 = noop.out_wire(0);
589
590            let mut nested =
591                f_build.dfg_builder(Signature::new(type_row![], vec![bool_t()]), [])?;
592
593            let id = nested.add_dataflow_op(Noop(bool_t()), [i1])?;
594
595            let nested = nested.finish_with_outputs([id.out_wire(0)])?;
596
597            f_build.finish_hugr_with_outputs([nested.out_wire(0)])
598        };
599
600        assert_matches!(builder(), Ok(_));
601    }
602
603    #[rstest]
604    #[case::function(|sig| FunctionBuilder::new("main", sig).unwrap().into_dfg_builder())]
605    #[case::dfg(|sig| DFGBuilder::new(sig).unwrap())]
606    #[case::df_block(|sig: Signature| {
607        let outs = sig.output.clone();
608        let mut h = CFGBuilder::new(sig).unwrap();
609        let entry_node = h.entry_builder([], outs).unwrap().finish_sub_container().unwrap().node();
610        let hugr = std::mem::take(h.hugr_mut());
611        DFGBuilder::create(hugr, entry_node).unwrap()
612    })]
613    fn add_inputs_outputs(#[case] builder: impl FnOnce(Signature) -> DFGBuilder<Hugr>) {
614        let builder = || -> Result<(Hugr, Node, Signature), BuildError> {
615            let mut dfg = builder(Signature::new(vec![bool_t()], vec![bool_t()]));
616            let dfg_node = dfg.container_node();
617
618            // Initial inner signature, before any inputs or outputs are added
619            let initial_sig = dfg
620                .hugr()
621                .get_optype(dfg_node)
622                .inner_function_type()
623                .unwrap()
624                .into_owned();
625
626            let [i0] = dfg.input_wires_arr();
627            let noop0 = dfg.add_dataflow_op(Noop(bool_t()), [i0])?;
628
629            // Some some order edges
630            dfg.set_order(&dfg.io()[0], &noop0.node());
631            dfg.set_order(&noop0.node(), &dfg.io()[1]);
632
633            // Add a new input and output, and connect them with a noop in between
634            dfg.add_output(qb_t()).unwrap();
635            let i1 = dfg.add_input(qb_t()).unwrap();
636            let noop1 = dfg.add_dataflow_op(Noop(qb_t()), [i1])?;
637
638            // Do not validate the final hugr, as it may have disconnected inputs.
639            dfg.set_outputs([noop0.out_wire(0), noop1.out_wire(0)])?;
640            let hugr = std::mem::take(dfg.hugr_mut());
641            Ok((hugr, dfg_node, initial_sig))
642        };
643
644        let (hugr, dfg_node, initial_sig) = builder().unwrap_or_else(|e| panic!("{e}"));
645
646        let container_sig = hugr.get_optype(dfg_node).inner_function_type().unwrap();
647        let mut expected_sig = initial_sig;
648        expected_sig.input.to_mut().push(qb_t());
649        expected_sig.output.to_mut().push(qb_t());
650        assert_eq!(
651            container_sig.io(),
652            (&expected_sig.input, &expected_sig.output),
653            "Got signature: {container_sig}, expected: {expected_sig}",
654        );
655    }
656
657    #[rstest]
658    #[case::tail_loop(|sig: Signature| TailLoopBuilder::new(sig.input, vec![], sig.output).unwrap().into_dfg_builder())]
659    fn add_inputs_outputs_unsupported(#[case] builder: impl FnOnce(Signature) -> DFGBuilder<Hugr>) {
660        let mut dfg = builder(Signature::new(vec![bool_t()], vec![bool_t()]));
661
662        // Add a new input and output, and connect them with a noop in between
663        assert!(dfg.add_output(qb_t()).is_err());
664        assert!(dfg.add_input(qb_t()).is_err());
665    }
666
667    #[test]
668    fn error_on_linear_inter_graph_edge() -> Result<(), BuildError> {
669        let mut f_build = FunctionBuilder::new("main", Signature::new(vec![qb_t()], vec![qb_t()]))?;
670
671        let [i1] = f_build.input_wires_arr();
672        let noop = f_build.add_dataflow_op(Noop(qb_t()), [i1])?;
673        let i1 = noop.out_wire(0);
674
675        let mut nested = f_build.dfg_builder(Signature::new(type_row![], vec![qb_t()]), [])?;
676
677        let id_res = nested.add_dataflow_op(Noop(qb_t()), [i1]);
678
679        // The error would anyway be caught in validation when we finish the Hugr,
680        // but the builder catches it earlier
681        assert_matches!(
682            id_res.map(|bh| bh.handle().node()), // Transform into something that impl's Debug
683            Err(BuildError::OperationWiring {
684                error: BuilderWiringError::NonCopyableIntergraph { .. },
685                ..
686            })
687        );
688
689        Ok(())
690    }
691
692    #[rstest]
693    fn dfg_hugr(simple_dfg_hugr: Hugr) {
694        assert_eq!(simple_dfg_hugr.num_nodes(), 7);
695        assert_eq!(simple_dfg_hugr.entry_descendants().count(), 3);
696        assert_matches!(simple_dfg_hugr.entrypoint_optype().tag(), OpTag::Dfg);
697    }
698
699    #[test]
700    fn add_hugr() -> Result<(), BuildError> {
701        // Create a simple DFG
702        let mut dfg_builder = DFGBuilder::new(Signature::new(vec![bool_t()], vec![bool_t()]))?;
703        let [i1] = dfg_builder.input_wires_arr();
704        dfg_builder.set_metadata("x", 42);
705        let dfg_hugr = dfg_builder.finish_hugr_with_outputs([i1])?;
706
707        // Create a module, and insert the DFG into it
708        let mut module_builder = ModuleBuilder::new();
709
710        let (dfg_node, f_node) = {
711            let mut f_build =
712                module_builder.define_function("main", Signature::new_endo(bool_t()))?;
713
714            let [i1] = f_build.input_wires_arr();
715            let dfg = f_build.add_hugr_with_wires(dfg_hugr, [i1])?;
716            let f = f_build.finish_with_outputs([dfg.out_wire(0)])?;
717            module_builder.set_child_metadata(f.node(), "x", "hi");
718            (dfg.node(), f.node())
719        };
720
721        let hugr = module_builder.finish_hugr()?;
722        assert_eq!(hugr.entry_descendants().count(), 7);
723
724        assert_eq!(hugr.get_metadata(hugr.entrypoint(), "x"), None);
725        assert_eq!(hugr.get_metadata(dfg_node, "x").cloned(), Some(json!(42)));
726        assert_eq!(hugr.get_metadata(f_node, "x").cloned(), Some(json!("hi")));
727
728        Ok(())
729    }
730
731    #[rstest]
732    fn add_hugr_link_nodes(
733        #[values(false, true)] replace: bool,
734        #[values(true, false)] view: bool,
735    ) {
736        let mut fb = FunctionBuilder::new("main", Signature::new_endo(bool_t())).unwrap();
737        let my_decl = fb
738            .module_root_builder()
739            .declare("func1", Signature::new_endo(bool_t()).into())
740            .unwrap();
741        let (insert, ins_defn, ins_decl) = dfg_calling_defn_decl();
742        let ins_defn_name = insert
743            .get_optype(ins_defn.node())
744            .as_func_defn()
745            .unwrap()
746            .func_name()
747            .clone();
748        let ins_decl_name = insert
749            .get_optype(ins_decl.node())
750            .as_func_decl()
751            .unwrap()
752            .func_name()
753            .clone();
754        let decl_mode = if replace {
755            NodeLinkingDirective::UseExisting(my_decl.node())
756        } else {
757            NodeLinkingDirective::add()
758        };
759        let link_spec = HashMap::from([
760            (ins_defn.node(), NodeLinkingDirective::add()),
761            (ins_decl.node(), decl_mode),
762        ]);
763        let inserted = if view {
764            fb.add_link_view_by_node_with_wires(&insert, [], link_spec)
765                .unwrap()
766        } else {
767            fb.add_link_hugr_by_node_with_wires(insert, [], link_spec)
768                .unwrap()
769        };
770        let h = fb.finish_hugr_with_outputs(inserted.outputs()).unwrap();
771        let defn_names = h
772            .nodes()
773            .filter_map(|n| h.get_optype(n).as_func_defn().map(FuncDefn::func_name))
774            .collect_vec();
775        assert_eq!(defn_names, [&"main".to_string(), &ins_defn_name]);
776        let decl_names = h
777            .nodes()
778            .filter_map(|n| h.get_optype(n).as_func_decl().map(FuncDecl::func_name))
779            .cloned()
780            .collect_vec();
781        let mut expected_decl_names = vec!["func1".to_string()];
782        if !replace {
783            expected_decl_names.push(ins_decl_name)
784        }
785        assert_eq!(decl_names, expected_decl_names);
786    }
787
788    #[test]
789    fn barrier_node() -> Result<(), BuildError> {
790        let mut parent = DFGBuilder::new(endo_sig(bool_t()))?;
791
792        let [w] = parent.input_wires_arr();
793
794        let mut dfg_b = parent.dfg_builder(endo_sig(bool_t()), [w])?;
795        let [w] = dfg_b.input_wires_arr();
796
797        let barr0 = dfg_b.add_barrier([w])?;
798        let [w] = barr0.outputs_arr();
799
800        let barr1 = dfg_b.add_barrier([w])?;
801        let [w] = barr1.outputs_arr();
802
803        let dfg = dfg_b.finish_with_outputs([w])?;
804        let [w] = dfg.outputs_arr();
805
806        let mut dfg2_b = parent.dfg_builder(endo_sig(vec![bool_t(), bool_t()]), [w, w])?;
807        let [w1, w2] = dfg2_b.input_wires_arr();
808        let barr2 = dfg2_b.add_barrier([w1, w2])?;
809        let wires: Vec<Wire> = barr2.outputs().collect();
810
811        let dfg2 = dfg2_b.finish_with_outputs(wires)?;
812        let [w, _] = dfg2.outputs_arr();
813        parent.finish_hugr_with_outputs([w])?;
814
815        Ok(())
816    }
817
818    #[test]
819    fn non_cfg_ancestor() -> Result<(), BuildError> {
820        let unit_sig = Signature::new(type_row![Type::UNIT], type_row![Type::UNIT]);
821        let mut b = DFGBuilder::new(unit_sig.clone())?;
822        let b_child = b.dfg_builder(unit_sig.clone(), [b.input().out_wire(0)])?;
823        let b_child_in_wire = b_child.input().out_wire(0);
824        b_child.finish_with_outputs([])?;
825        let b_child_2 = b.dfg_builder(unit_sig.clone(), [])?;
826
827        // DFG block has edge coming a sibling block, which is only valid for
828        // CFGs
829        let b_child_2_handle = b_child_2.finish_with_outputs([b_child_in_wire])?;
830
831        let res = b.finish_hugr_with_outputs([b_child_2_handle.out_wire(0)]);
832
833        assert_matches!(
834            res,
835            Err(BuildError::InvalidHUGR(
836                ValidationError::InterGraphEdgeError(InterGraphEdgeError::NonCFGAncestor { .. })
837            ))
838        );
839        Ok(())
840    }
841
842    #[test]
843    fn no_relation_edge() -> Result<(), BuildError> {
844        let unit_sig = Signature::new(type_row![Type::UNIT], type_row![Type::UNIT]);
845        let mut b = DFGBuilder::new(unit_sig.clone())?;
846        let mut b_child = b.dfg_builder(unit_sig.clone(), [b.input().out_wire(0)])?;
847        let b_child_child = b_child.dfg_builder(unit_sig.clone(), [b_child.input().out_wire(0)])?;
848        let b_child_child_in_wire = b_child_child.input().out_wire(0);
849
850        b_child_child.finish_with_outputs([])?;
851        b_child.finish_with_outputs([])?;
852
853        let mut b_child_2 = b.dfg_builder(unit_sig.clone(), [])?;
854        let b_child_2_child =
855            b_child_2.dfg_builder(unit_sig.clone(), [b_child_2.input().out_wire(0)])?;
856
857        let res = b_child_2_child.finish_with_outputs([b_child_child_in_wire]);
858
859        assert_matches!(
860            res.map(|h| h.handle().node()), // map to something that implements Debug
861            Err(BuildError::OutputWiring {
862                error: BuilderWiringError::NoRelationIntergraph { .. },
863                ..
864            })
865        );
866        Ok(())
867    }
868
869    #[test]
870    fn no_outer_row_variables() -> Result<(), BuildError> {
871        let e = crate::hugr::validate::test::extension_with_eval_parallel();
872        let tv = TypeRV::new_row_var_use(0, TypeBound::Copyable);
873        // Can *declare* a function that takes a function-value of unknown #args
874        FunctionBuilder::new(
875            "bad_eval",
876            PolyFuncType::new(
877                [TypeParam::new_list_type(TypeBound::Copyable)],
878                Signature::new(
879                    Type::new_function(FuncValueType::new(usize_t(), tv.clone())),
880                    vec![],
881                ),
882            ),
883        )?;
884
885        // But cannot eval it...
886        let ev = e.instantiate_extension_op(
887            "eval",
888            [vec![usize_t().into()].into(), vec![tv.into()].into()],
889        );
890        assert_eq!(
891            ev,
892            Err(SignatureError::RowVarWhereTypeExpected {
893                var: RowVariable(0, TypeBound::Copyable)
894            })
895        );
896        Ok(())
897    }
898
899    #[test]
900    fn order_edges() {
901        let (mut hugr, load_constant, call) = {
902            let mut builder = ModuleBuilder::new();
903            let func = builder
904                .declare("func", Signature::new_endo(bool_t()).into())
905                .unwrap();
906            let (load_constant, call) = {
907                let mut builder = builder
908                    .define_function("main", Signature::new(Type::EMPTY_TYPEROW, bool_t()))
909                    .unwrap();
910                let load_constant = builder.add_load_value(Value::true_val());
911                let [r] = builder
912                    .call(&func, &[], [load_constant])
913                    .unwrap()
914                    .outputs_arr();
915                builder.finish_with_outputs([r]).unwrap();
916                (load_constant.node(), r.node())
917            };
918            (builder.finish_hugr().unwrap(), load_constant, call)
919        };
920
921        let lc_optype = hugr.get_optype(load_constant);
922        let call_optype = hugr.get_optype(call);
923        assert_eq!(EdgeKind::StateOrder, lc_optype.other_input().unwrap());
924        assert_eq!(EdgeKind::StateOrder, lc_optype.other_output().unwrap());
925        assert_eq!(EdgeKind::StateOrder, call_optype.other_input().unwrap());
926        assert_eq!(EdgeKind::StateOrder, call_optype.other_output().unwrap());
927
928        hugr.connect(
929            load_constant,
930            lc_optype.other_output_port().unwrap(),
931            call,
932            call_optype.other_input_port().unwrap(),
933        );
934
935        hugr.validate().unwrap();
936    }
937}
hugr_core/builder/dataflow.rs

hugr_core/builder/
dataflow.rs
