hugr_core/builder/
handle.rs

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
//! Handles to nodes in HUGR used during the building phase.
use crate::ops::handle::{BasicBlockID, CaseID, DfgID, FuncID, NodeHandle};
use crate::ops::OpTag;
use crate::utils::collect_array;
use crate::{Node, OutgoingPort, Wire};

use std::iter::FusedIterator;

#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash)]
/// Handle to a dataflow node which has a known number of value outputs
pub struct BuildHandle<T> {
    node_handle: T,
    num_value_outputs: usize,
}

impl<T: From<Node>> From<(Node, usize)> for BuildHandle<T> {
    fn from((node, num_value_outputs): (Node, usize)) -> Self {
        Self {
            node_handle: node.into(),
            num_value_outputs,
        }
    }
}

impl<T: NodeHandle> NodeHandle for BuildHandle<T> {
    const TAG: OpTag = T::TAG;

    fn node(&self) -> Node {
        self.node_handle.node()
    }
}

impl<T: NodeHandle> BuildHandle<T> {
    #[inline]
    /// Number of Value kind outputs from this node.
    pub fn num_value_outputs(&self) -> usize {
        self.num_value_outputs
    }

    #[inline]
    /// Return iterator over Value outputs.
    pub fn outputs(&self) -> Outputs {
        Outputs {
            node: self.node(),
            range: (0..self.num_value_outputs()),
        }
    }

    /// Attempt to cast outputs in to array of Wires.
    pub fn outputs_arr<const N: usize>(&self) -> [Wire; N] {
        self.outputs().to_array()
    }

    #[inline]
    /// Retrieve a [`Wire`] corresponding to the given offset.
    /// Does not check whether such a wire is valid for this node.
    pub fn out_wire(&self, offset: usize) -> Wire {
        Wire::new(self.node(), OutgoingPort::from(offset))
    }

    #[inline]
    /// Underlying node handle
    pub fn handle(&self) -> &T {
        &self.node_handle
    }
}

impl From<BuildHandle<DfgID>> for BuildHandle<FuncID<true>> {
    #[inline]
    fn from(value: BuildHandle<DfgID>) -> Self {
        Self {
            node_handle: value.node().into(),
            num_value_outputs: 0,
        }
    }
}

impl From<BuildHandle<DfgID>> for BasicBlockID {
    #[inline]
    fn from(value: BuildHandle<DfgID>) -> Self {
        value.node().into()
    }
}

impl From<BuildHandle<DfgID>> for BuildHandle<CaseID> {
    #[inline]
    fn from(value: BuildHandle<DfgID>) -> Self {
        Self {
            node_handle: value.node().into(),
            num_value_outputs: 0,
        }
    }
}

#[derive(Debug, Clone)]
/// Iterator over output wires of a [`BuildHandle`].
pub struct Outputs {
    node: Node,
    range: std::ops::Range<usize>,
}

impl Outputs {
    #[inline]
    /// Returns the output wires as an array.
    ///
    /// # Panics
    ///
    /// If the length of the slice is not equal to `N`.
    pub fn to_array<const N: usize>(self) -> [Wire; N] {
        collect_array(self)
    }
}

impl Iterator for Outputs {
    type Item = Wire;

    fn next(&mut self) -> Option<Self::Item> {
        self.range
            .next()
            .map(|offset| Wire::new(self.node, OutgoingPort::from(offset)))
    }

    #[inline]
    fn nth(&mut self, n: usize) -> Option<Self::Item> {
        self.range.nth(n).map(|offset| Wire::new(self.node, offset))
    }

    #[inline]
    fn count(self) -> usize {
        self.range.count()
    }

    #[inline]
    fn size_hint(&self) -> (usize, Option<usize>) {
        self.range.size_hint()
    }
}

impl ExactSizeIterator for Outputs {
    #[inline]
    fn len(&self) -> usize {
        self.range.len()
    }
}

impl DoubleEndedIterator for Outputs {
    #[inline]
    fn next_back(&mut self) -> Option<Self::Item> {
        self.range
            .next_back()
            .map(|offset| Wire::new(self.node, offset))
    }
}

impl FusedIterator for Outputs {}