hugr_core/builder/
handle.rs

1//! Handles to nodes in HUGR used during the building phase.
2use crate::ops::OpTag;
3use crate::ops::handle::{BasicBlockID, CaseID, DfgID, FuncID, NodeHandle};
4use crate::utils::collect_array;
5use crate::{Node, OutgoingPort, Wire};
6
7use std::iter::FusedIterator;
8
9#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash)]
10/// Handle to a dataflow node which has a known number of value outputs
11pub struct BuildHandle<T> {
12    node_handle: T,
13    num_value_outputs: usize,
14}
15
16impl<T: From<Node>> From<(Node, usize)> for BuildHandle<T> {
17    fn from((node, num_value_outputs): (Node, usize)) -> Self {
18        Self {
19            node_handle: node.into(),
20            num_value_outputs,
21        }
22    }
23}
24
25impl<T: NodeHandle> NodeHandle for BuildHandle<T> {
26    const TAG: OpTag = T::TAG;
27
28    fn node(&self) -> Node {
29        self.node_handle.node()
30    }
31}
32
33impl<T: NodeHandle> BuildHandle<T> {
34    #[inline]
35    /// Number of Value kind outputs from this node.
36    pub fn num_value_outputs(&self) -> usize {
37        self.num_value_outputs
38    }
39
40    #[inline]
41    /// Return iterator over Value outputs.
42    pub fn outputs(&self) -> Outputs {
43        Outputs {
44            node: self.node(),
45            range: (0..self.num_value_outputs()),
46        }
47    }
48
49    /// Attempt to cast outputs in to array of Wires.
50    pub fn outputs_arr<const N: usize>(&self) -> [Wire; N] {
51        self.outputs().to_array()
52    }
53
54    #[inline]
55    /// Retrieve a [`Wire`] corresponding to the given offset.
56    /// Does not check whether such a wire is valid for this node.
57    pub fn out_wire(&self, offset: usize) -> Wire {
58        Wire::new(self.node(), OutgoingPort::from(offset))
59    }
60
61    #[inline]
62    /// Underlying node handle
63    pub fn handle(&self) -> &T {
64        &self.node_handle
65    }
66}
67
68impl From<BuildHandle<DfgID>> for BuildHandle<FuncID<true>> {
69    #[inline]
70    fn from(value: BuildHandle<DfgID>) -> Self {
71        Self {
72            node_handle: value.node().into(),
73            num_value_outputs: 0,
74        }
75    }
76}
77
78impl From<BuildHandle<DfgID>> for BasicBlockID {
79    #[inline]
80    fn from(value: BuildHandle<DfgID>) -> Self {
81        value.node().into()
82    }
83}
84
85impl From<BuildHandle<DfgID>> for BuildHandle<CaseID> {
86    #[inline]
87    fn from(value: BuildHandle<DfgID>) -> Self {
88        Self {
89            node_handle: value.node().into(),
90            num_value_outputs: 0,
91        }
92    }
93}
94
95#[derive(Debug, Clone)]
96/// Iterator over output wires of a [`BuildHandle`].
97pub struct Outputs {
98    node: Node,
99    range: std::ops::Range<usize>,
100}
101
102impl Outputs {
103    #[inline]
104    /// Returns the output wires as an array.
105    ///
106    /// # Panics
107    ///
108    /// If the length of the slice is not equal to `N`.
109    #[must_use]
110    pub fn to_array<const N: usize>(self) -> [Wire; N] {
111        collect_array(self)
112    }
113}
114
115impl Iterator for Outputs {
116    type Item = Wire;
117
118    fn next(&mut self) -> Option<Self::Item> {
119        self.range
120            .next()
121            .map(|offset| Wire::new(self.node, OutgoingPort::from(offset)))
122    }
123
124    #[inline]
125    fn nth(&mut self, n: usize) -> Option<Self::Item> {
126        self.range.nth(n).map(|offset| Wire::new(self.node, offset))
127    }
128
129    #[inline]
130    fn count(self) -> usize {
131        self.range.count()
132    }
133
134    #[inline]
135    fn size_hint(&self) -> (usize, Option<usize>) {
136        self.range.size_hint()
137    }
138}
139
140impl ExactSizeIterator for Outputs {
141    #[inline]
142    fn len(&self) -> usize {
143        self.range.len()
144    }
145}
146
147impl DoubleEndedIterator for Outputs {
148    #[inline]
149    fn next_back(&mut self) -> Option<Self::Item> {
150        self.range
151            .next_back()
152            .map(|offset| Wire::new(self.node, offset))
153    }
154}
155
156impl FusedIterator for Outputs {}