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

hugr_core/hugr/
internal.rs
