hugr_core/builder/
handle.rs

1//! Handles to nodes in HUGR used during the building phase.
2use crate::ops::handle::{BasicBlockID, CaseID, DfgID, FuncID, NodeHandle};
3use crate::ops::OpTag;
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    pub fn to_array<const N: usize>(self) -> [Wire; N] {
110        collect_array(self)
111    }
112}
113
114impl Iterator for Outputs {
115    type Item = Wire;
116
117    fn next(&mut self) -> Option<Self::Item> {
118        self.range
119            .next()
120            .map(|offset| Wire::new(self.node, OutgoingPort::from(offset)))
121    }
122
123    #[inline]
124    fn nth(&mut self, n: usize) -> Option<Self::Item> {
125        self.range.nth(n).map(|offset| Wire::new(self.node, offset))
126    }
127
128    #[inline]
129    fn count(self) -> usize {
130        self.range.count()
131    }
132
133    #[inline]
134    fn size_hint(&self) -> (usize, Option<usize>) {
135        self.range.size_hint()
136    }
137}
138
139impl ExactSizeIterator for Outputs {
140    #[inline]
141    fn len(&self) -> usize {
142        self.range.len()
143    }
144}
145
146impl DoubleEndedIterator for Outputs {
147    #[inline]
148    fn next_back(&mut self) -> Option<Self::Item> {
149        self.range
150            .next_back()
151            .map(|offset| Wire::new(self.node, offset))
152    }
153}
154
155impl FusedIterator for Outputs {}