hugr_core/hugr/
internal.rs

1//! Internal traits, not exposed in the public `hugr` API.
2
3use std::ops::Range;
4use std::sync::OnceLock;
5
6use itertools::Itertools;
7use portgraph::{LinkMut, LinkView, MultiPortGraph, PortMut, PortOffset, PortView};
8
9use crate::core::HugrNode;
10use crate::extension::ExtensionRegistry;
11use crate::{Direction, Hugr, Node};
12
13use super::HugrView;
14use super::views::{panic_invalid_node, panic_invalid_non_entrypoint};
15use super::{NodeMetadataMap, OpType};
16use crate::ops::handle::NodeHandle;
17
18/// Trait for accessing the internals of a Hugr(View).
19///
20/// Specifically, this trait provides access to the underlying portgraph
21/// view.
22pub trait HugrInternals {
23    /// The portgraph graph structure used internally by [`scheduling_graph`].
24    ///
25    /// [`scheduling_graph`]: HugrView::scheduling_graph
26    type RegionPortgraph<'p>: LinkView<LinkEndpoint: Eq, NodeIndexBase = u32, PortIndexBase = u32, PortOffsetBase = u32>
27        + Clone
28        + 'p
29    where
30        Self: 'p;
31
32    /// The type of nodes in the Hugr.
33    type Node: Copy + Ord + std::fmt::Debug + std::fmt::Display + std::hash::Hash;
34
35    /// A mapping between HUGR nodes and portgraph nodes in the graph returned by
36    /// [`scheduling_graph`].
37    ///
38    /// [`scheduling_graph`]: HugrView::scheduling_graph
39    type RegionPortgraphNodes: PortgraphNodeMap<Self::Node>;
40
41    /// Returns a metadata entry associated with a node.
42    ///
43    /// # Panics
44    ///
45    /// If the node is not in the graph.
46    fn node_metadata_map(&self, node: Self::Node) -> &NodeMetadataMap;
47}
48
49/// A map between hugr nodes and portgraph nodes in a [HugrInternals::RegionPortgraph]
50/// or [`scheduling_graph`].
51///
52/// [`scheduling_graph`]: HugrView::scheduling_graph
53pub trait PortgraphNodeMap<N>: Clone + Sized + std::fmt::Debug {
54    /// Returns the portgraph index of a HUGR node in the associated region
55    /// graph.
56    ///
57    /// If the node is not in the region, the result is undefined.
58    fn to_portgraph(&self, node: N) -> portgraph::NodeIndex;
59
60    /// Returns the HUGR node for a portgraph node in the associated region
61    /// graph.
62    ///
63    /// If the node is not in the region, the result is undefined.
64    #[allow(clippy::wrong_self_convention)]
65    fn from_portgraph(&self, node: portgraph::NodeIndex) -> N;
66}
67
68/// An identity map between HUGR nodes and portgraph nodes.
69#[derive(
70    Copy, Clone, Debug, Default, Eq, PartialEq, Hash, PartialOrd, Ord, derive_more::Display,
71)]
72pub struct DefaultPGNodeMap;
73
74impl PortgraphNodeMap<Node> for DefaultPGNodeMap {
75    #[inline]
76    fn to_portgraph(&self, node: Node) -> portgraph::NodeIndex {
77        node.into_portgraph()
78    }
79
80    #[inline]
81    fn from_portgraph(&self, node: portgraph::NodeIndex) -> Node {
82        node.into()
83    }
84}
85
86impl<N: HugrNode> PortgraphNodeMap<N> for std::collections::HashMap<N, Node> {
87    #[inline]
88    fn to_portgraph(&self, node: N) -> portgraph::NodeIndex {
89        self[&node].into_portgraph()
90    }
91
92    #[inline]
93    fn from_portgraph(&self, node: portgraph::NodeIndex) -> N {
94        let node = node.into();
95        self.iter()
96            .find_map(|(&k, &v)| (v == node).then_some(k))
97            .expect("Portgraph node not found in map")
98    }
99}
100
101impl HugrInternals for Hugr {
102    type RegionPortgraph<'p>
103        = &'p MultiPortGraph<u32, u32, u32>
104    where
105        Self: 'p;
106
107    type Node = Node;
108
109    type RegionPortgraphNodes = DefaultPGNodeMap;
110
111    #[inline]
112    fn node_metadata_map(&self, node: Self::Node) -> &NodeMetadataMap {
113        static EMPTY: OnceLock<NodeMetadataMap> = OnceLock::new();
114        panic_invalid_node(self, node);
115        let map = self.metadata.get(node.into_portgraph()).as_ref();
116        map.unwrap_or(EMPTY.get_or_init(Default::default))
117    }
118}
119
120/// Trait for accessing the mutable internals of a Hugr(Mut).
121///
122/// Specifically, this trait lets you apply arbitrary modifications that may
123/// invalidate the HUGR.
124pub trait HugrMutInternals: HugrView {
125    /// Set the node at the root of the HUGR hierarchy.
126    ///
127    /// Any node not reachable from this root should be manually removed from
128    /// the HUGR.
129    ///
130    /// To set the working entrypoint of the HUGR, use
131    /// [`HugrMut::set_entrypoint`][crate::hugr::HugrMut::set_entrypoint]
132    /// instead.
133    ///
134    /// # Panics
135    ///
136    /// If the node is not in the graph.
137    fn set_module_root(&mut self, root: Self::Node);
138
139    /// Set the number of ports on a node. This may invalidate the node's `PortIndex`.
140    ///
141    /// # Panics
142    ///
143    /// If the node is not in the graph.
144    fn set_num_ports(&mut self, node: Self::Node, incoming: usize, outgoing: usize);
145
146    /// Alter the number of ports on a node and returns a range with the new
147    /// port offsets, if any. This may invalidate the node's `PortIndex`.
148    ///
149    /// The `direction` parameter specifies whether to add ports to the incoming
150    /// or outgoing list.
151    ///
152    /// Returns the range of newly created ports.
153    ///
154    /// # Panics
155    ///
156    /// If the node is not in the graph.
157    fn add_ports(&mut self, node: Self::Node, direction: Direction, amount: isize) -> Range<usize>;
158
159    /// Insert `amount` new ports for a node, starting at `index`.  The
160    /// `direction` parameter specifies whether to add ports to the incoming or
161    /// outgoing list. Links from this node are preserved, even when ports are
162    /// renumbered by the insertion.
163    ///
164    /// Returns the range of newly created ports.
165    /// # Panics
166    ///
167    /// If the node is not in the graph.
168    fn insert_ports(
169        &mut self,
170        node: Self::Node,
171        direction: Direction,
172        index: usize,
173        amount: usize,
174    ) -> Range<usize>;
175
176    /// Sets the parent of a node.
177    ///
178    /// The node becomes the parent's last child.
179    ///
180    /// # Panics
181    ///
182    /// If either the node or the parent is not in the graph.
183    fn set_parent(&mut self, node: Self::Node, parent: Self::Node);
184
185    /// Move a node in the hierarchy to be the subsequent sibling of another
186    /// node.
187    ///
188    /// The sibling node's parent becomes the new node's parent.
189    ///
190    /// The node becomes the parent's last child.
191    ///
192    /// # Panics
193    ///
194    /// If either node is not in the graph, or if it is a root.
195    fn move_after_sibling(&mut self, node: Self::Node, after: Self::Node);
196
197    /// Move a node in the hierarchy to be the prior sibling of another node.
198    ///
199    /// The sibling node's parent becomes the new node's parent.
200    ///
201    /// The node becomes the parent's last child.
202    ///
203    /// # Panics
204    ///
205    /// If either node is not in the graph, or if it is a root.
206    fn move_before_sibling(&mut self, node: Self::Node, before: Self::Node);
207
208    /// Replace the `OpType` at node and return the old `OpType`.
209    /// In general this invalidates the ports, which may need to be resized to
210    /// match the `OpType` signature.
211    ///
212    /// Returns the old `OpType`.
213    ///
214    /// If the module root is set to a non-module operation the hugr will
215    /// become invalid.
216    ///
217    /// # Panics
218    ///
219    /// If the node is not in the graph.
220    fn replace_op(&mut self, node: Self::Node, op: impl Into<OpType>) -> OpType;
221
222    /// Gets a mutable reference to the optype.
223    ///
224    /// Changing this may invalidate the ports, which may need to be resized to
225    /// match the `OpType` signature.
226    ///
227    /// Mutating the root node operation may invalidate the root tag.
228    ///
229    /// Mutating the module root into a non-module operation will invalidate the hugr.
230    ///
231    /// # Panics
232    ///
233    /// If the node is not in the graph.
234    fn optype_mut(&mut self, node: Self::Node) -> &mut OpType;
235
236    /// Returns a metadata entry associated with a node.
237    ///
238    /// # Panics
239    ///
240    /// If the node is not in the graph.
241    fn node_metadata_map_mut(&mut self, node: Self::Node) -> &mut NodeMetadataMap;
242
243    /// Returns a mutable reference to the extension registry for this HUGR.
244    ///
245    /// This set contains all extensions required to define the operations and
246    /// types in the HUGR.
247    fn extensions_mut(&mut self) -> &mut ExtensionRegistry;
248}
249
250/// Impl for non-wrapped Hugrs. Overwrites the recursive default-impls to directly use the hugr.
251impl HugrMutInternals for Hugr {
252    fn set_module_root(&mut self, root: Node) {
253        panic_invalid_node(self, root.node());
254        let root = root.into_portgraph();
255        self.hierarchy.detach(root);
256        self.module_root = root;
257    }
258
259    #[inline]
260    fn set_num_ports(&mut self, node: Node, incoming: usize, outgoing: usize) {
261        panic_invalid_node(self, node);
262        self.graph
263            .set_num_ports(node.into_portgraph(), incoming, outgoing, |_, _| {});
264    }
265
266    fn add_ports(&mut self, node: Node, direction: Direction, amount: isize) -> Range<usize> {
267        panic_invalid_node(self, node);
268        let mut incoming = self.graph.num_inputs(node.into_portgraph());
269        let mut outgoing = self.graph.num_outputs(node.into_portgraph());
270        let increment = |num: &mut usize| {
271            let new = num.saturating_add_signed(amount);
272            let range = *num..new;
273            *num = new;
274            range
275        };
276        let range = match direction {
277            Direction::Incoming => increment(&mut incoming),
278            Direction::Outgoing => increment(&mut outgoing),
279        };
280        self.graph
281            .set_num_ports(node.into_portgraph(), incoming, outgoing, |_, _| {});
282        range
283    }
284
285    fn insert_ports(
286        &mut self,
287        node: Node,
288        direction: Direction,
289        index: usize,
290        amount: usize,
291    ) -> Range<usize> {
292        panic_invalid_node(self, node);
293        let old_num_ports = self.graph.num_ports(node.into_portgraph(), direction);
294
295        self.add_ports(node, direction, amount as isize);
296
297        for swap_from_port in (index..old_num_ports).rev() {
298            let swap_to_port = swap_from_port + amount;
299            let [from_port_index, to_port_index] = [swap_from_port, swap_to_port].map(|p| {
300                self.graph
301                    .port_index(node.into_portgraph(), PortOffset::new(direction, p))
302                    .unwrap()
303            });
304            let linked_ports = self
305                .graph
306                .port_links(from_port_index)
307                .map(|(_, to_subport)| to_subport.port())
308                .collect_vec();
309            self.graph.unlink_port(from_port_index);
310            for linked_port_index in linked_ports {
311                let _ = self
312                    .graph
313                    .link_ports(to_port_index, linked_port_index)
314                    .expect("Ports exist");
315            }
316        }
317        index..index + amount
318    }
319
320    fn set_parent(&mut self, node: Node, parent: Node) {
321        panic_invalid_node(self, parent);
322        panic_invalid_node(self, node);
323        self.hierarchy.detach(node.into_portgraph());
324        self.hierarchy
325            .push_child(node.into_portgraph(), parent.into_portgraph())
326            .expect("Inserting a newly-created node into the hierarchy should never fail.");
327    }
328
329    fn move_after_sibling(&mut self, node: Node, after: Node) {
330        panic_invalid_non_entrypoint(self, node);
331        panic_invalid_non_entrypoint(self, after);
332        self.hierarchy.detach(node.into_portgraph());
333        self.hierarchy
334            .insert_after(node.into_portgraph(), after.into_portgraph())
335            .expect("Inserting a newly-created node into the hierarchy should never fail.");
336    }
337
338    fn move_before_sibling(&mut self, node: Node, before: Node) {
339        panic_invalid_non_entrypoint(self, node);
340        panic_invalid_non_entrypoint(self, before);
341        self.hierarchy.detach(node.into_portgraph());
342        self.hierarchy
343            .insert_before(node.into_portgraph(), before.into_portgraph())
344            .expect("Inserting a newly-created node into the hierarchy should never fail.");
345    }
346
347    fn replace_op(&mut self, node: Node, op: impl Into<OpType>) -> OpType {
348        panic_invalid_node(self, node);
349        std::mem::replace(self.optype_mut(node), op.into())
350    }
351
352    fn optype_mut(&mut self, node: Node) -> &mut OpType {
353        panic_invalid_node(self, node);
354        let node = node.into_portgraph();
355        self.op_types.get_mut(node)
356    }
357
358    fn node_metadata_map_mut(&mut self, node: Self::Node) -> &mut NodeMetadataMap {
359        panic_invalid_node(self, node);
360        self.metadata
361            .get_mut(node.into_portgraph())
362            .get_or_insert_with(Default::default)
363    }
364
365    fn extensions_mut(&mut self) -> &mut ExtensionRegistry {
366        &mut self.extensions
367    }
368}
369
370#[cfg(test)]
371mod test {
372    use crate::{
373        Direction, HugrView as _,
374        builder::{Container, DFGBuilder, Dataflow, DataflowHugr},
375        extension::prelude::Noop,
376        hugr::internal::HugrMutInternals as _,
377        ops::handle::NodeHandle,
378        types::{Signature, Type},
379    };
380
381    #[test]
382    fn insert_ports() {
383        let (nop, mut hugr) = {
384            let mut builder = DFGBuilder::new(Signature::new_endo([Type::UNIT])).unwrap();
385            let [nop_in] = builder.input_wires_arr();
386            let nop = builder
387                .add_dataflow_op(Noop::new(Type::UNIT), [nop_in])
388                .unwrap();
389            builder.add_other_wire(nop.node(), builder.output().node());
390            let [nop_out] = nop.outputs_arr();
391            (
392                nop.node(),
393                builder.finish_hugr_with_outputs([nop_out]).unwrap(),
394            )
395        };
396        let [i, o] = hugr.get_io(hugr.entrypoint()).unwrap();
397        assert_eq!(0..2, hugr.insert_ports(nop, Direction::Incoming, 0, 2));
398        assert_eq!(1..3, hugr.insert_ports(nop, Direction::Outgoing, 1, 2));
399
400        assert_eq!(hugr.single_linked_input(i, 0), Some((nop, 2.into())));
401        assert_eq!(hugr.single_linked_output(o, 0), Some((nop, 0.into())));
402        assert_eq!(hugr.single_linked_output(o, 1), Some((nop, 3.into())));
403    }
404}
hugr_core/hugr/internal.rs

hugr_core/hugr/
internal.rs
