hugr_core/hugr/
internal.rs

1//! Internal traits, not exposed in the public `hugr` API.
2
3use std::borrow::Cow;
4use std::ops::Range;
5use std::rc::Rc;
6use std::sync::Arc;
7
8use delegate::delegate;
9use itertools::Itertools;
10use portgraph::{LinkMut, LinkView, MultiPortGraph, PortMut, PortOffset, PortView};
11
12use crate::ops::handle::NodeHandle;
13use crate::ops::OpTrait;
14use crate::{Direction, Hugr, Node};
15
16use super::hugrmut::{panic_invalid_node, panic_invalid_non_root};
17use super::{HugrError, OpType, RootTagged};
18
19/// Trait for accessing the internals of a Hugr(View).
20///
21/// Specifically, this trait provides access to the underlying portgraph
22/// view.
23pub trait HugrInternals {
24    /// The underlying portgraph view type.
25    type Portgraph<'p>: LinkView + Clone + 'p
26    where
27        Self: 'p;
28
29    /// The type of nodes in the Hugr.
30    type Node: Copy + Ord + std::fmt::Debug + std::fmt::Display + std::hash::Hash;
31
32    /// Returns a reference to the underlying portgraph.
33    fn portgraph(&self) -> Self::Portgraph<'_>;
34
35    /// Returns the Hugr at the base of a chain of views.
36    fn base_hugr(&self) -> &Hugr;
37
38    /// Return the root node of this view.
39    fn root_node(&self) -> Self::Node;
40
41    /// Convert a node to a portgraph node index.
42    fn get_pg_index(&self, node: Self::Node) -> portgraph::NodeIndex;
43
44    /// Convert a portgraph node index to a node.
45    fn get_node(&self, index: portgraph::NodeIndex) -> Self::Node;
46}
47
48impl HugrInternals for Hugr {
49    type Portgraph<'p>
50        = &'p MultiPortGraph
51    where
52        Self: 'p;
53
54    type Node = Node;
55
56    #[inline]
57    fn portgraph(&self) -> Self::Portgraph<'_> {
58        &self.graph
59    }
60
61    #[inline]
62    fn base_hugr(&self) -> &Hugr {
63        self
64    }
65
66    #[inline]
67    fn root_node(&self) -> Self::Node {
68        self.root.into()
69    }
70
71    fn get_pg_index(&self, node: Self::Node) -> portgraph::NodeIndex {
72        node.pg_index()
73    }
74
75    fn get_node(&self, index: portgraph::NodeIndex) -> Self::Node {
76        index.into()
77    }
78}
79
80impl<T: HugrInternals> HugrInternals for &T {
81    type Portgraph<'p>
82        = T::Portgraph<'p>
83    where
84        Self: 'p;
85    type Node = T::Node;
86
87    delegate! {
88        to (**self) {
89            fn portgraph(&self) -> Self::Portgraph<'_>;
90            fn base_hugr(&self) -> &Hugr;
91            fn root_node(&self) -> Self::Node;
92            fn get_pg_index(&self, node: Self::Node) -> portgraph::NodeIndex;
93            fn get_node(&self, index: portgraph::NodeIndex) -> Self::Node;
94        }
95    }
96}
97
98impl<T: HugrInternals> HugrInternals for &mut T {
99    type Portgraph<'p>
100        = T::Portgraph<'p>
101    where
102        Self: 'p;
103    type Node = T::Node;
104
105    delegate! {
106        to (**self) {
107            fn portgraph(&self) -> Self::Portgraph<'_>;
108            fn base_hugr(&self) -> &Hugr;
109            fn root_node(&self) -> Self::Node;
110            fn get_pg_index(&self, node: Self::Node) -> portgraph::NodeIndex;
111            fn get_node(&self, index: portgraph::NodeIndex) -> Self::Node;
112        }
113    }
114}
115
116impl<T: HugrInternals> HugrInternals for Rc<T> {
117    type Portgraph<'p>
118        = T::Portgraph<'p>
119    where
120        Self: 'p;
121    type Node = T::Node;
122
123    delegate! {
124        to (**self) {
125            fn portgraph(&self) -> Self::Portgraph<'_>;
126            fn base_hugr(&self) -> &Hugr;
127            fn root_node(&self) -> Self::Node;
128            fn get_pg_index(&self, node: Self::Node) -> portgraph::NodeIndex;
129            fn get_node(&self, index: portgraph::NodeIndex) -> Self::Node;
130        }
131    }
132}
133
134impl<T: HugrInternals> HugrInternals for Arc<T> {
135    type Portgraph<'p>
136        = T::Portgraph<'p>
137    where
138        Self: 'p;
139    type Node = T::Node;
140
141    delegate! {
142        to (**self) {
143            fn portgraph(&self) -> Self::Portgraph<'_>;
144            fn base_hugr(&self) -> &Hugr;
145            fn root_node(&self) -> Self::Node;
146            fn get_pg_index(&self, node: Self::Node) -> portgraph::NodeIndex;
147            fn get_node(&self, index: portgraph::NodeIndex) -> Self::Node;
148        }
149    }
150}
151
152impl<T: HugrInternals> HugrInternals for Box<T> {
153    type Portgraph<'p>
154        = T::Portgraph<'p>
155    where
156        Self: 'p;
157    type Node = T::Node;
158
159    delegate! {
160        to (**self) {
161            fn portgraph(&self) -> Self::Portgraph<'_>;
162            fn base_hugr(&self) -> &Hugr;
163            fn root_node(&self) -> Self::Node;
164            fn get_pg_index(&self, node: Self::Node) -> portgraph::NodeIndex;
165            fn get_node(&self, index: portgraph::NodeIndex) -> Self::Node;
166        }
167    }
168}
169
170impl<T: HugrInternals + ToOwned> HugrInternals for Cow<'_, T> {
171    type Portgraph<'p>
172        = T::Portgraph<'p>
173    where
174        Self: 'p;
175    type Node = T::Node;
176
177    delegate! {
178        to self.as_ref() {
179            fn portgraph(&self) -> Self::Portgraph<'_>;
180            fn base_hugr(&self) -> &Hugr;
181            fn root_node(&self) -> Self::Node;
182            fn get_pg_index(&self, node: Self::Node) -> portgraph::NodeIndex;
183            fn get_node(&self, index: portgraph::NodeIndex) -> Self::Node;
184        }
185    }
186}
187/// Trait for accessing the mutable internals of a Hugr(Mut).
188///
189/// Specifically, this trait lets you apply arbitrary modifications that may
190/// invalidate the HUGR.
191pub trait HugrMutInternals: RootTagged<Node = Node> {
192    /// Returns the Hugr at the base of a chain of views.
193    fn hugr_mut(&mut self) -> &mut Hugr;
194
195    /// Set the number of ports on a node. This may invalidate the node's `PortIndex`.
196    ///
197    /// # Panics
198    ///
199    /// If the node is not in the graph.
200    fn set_num_ports(&mut self, node: Node, incoming: usize, outgoing: usize) {
201        panic_invalid_node(self, node);
202        self.hugr_mut().set_num_ports(node, incoming, outgoing)
203    }
204
205    /// Alter the number of ports on a node and returns a range with the new
206    /// port offsets, if any. This may invalidate the node's `PortIndex`.
207    ///
208    /// The `direction` parameter specifies whether to add ports to the incoming
209    /// or outgoing list.
210    ///
211    /// Returns the range of newly created ports.
212    ///
213    /// # Panics
214    ///
215    /// If the node is not in the graph.
216    fn add_ports(&mut self, node: Node, direction: Direction, amount: isize) -> Range<usize> {
217        panic_invalid_node(self, node);
218        self.hugr_mut().add_ports(node, direction, amount)
219    }
220
221    /// Insert `amount` new ports for a node, starting at `index`.  The
222    /// `direction` parameter specifies whether to add ports to the incoming or
223    /// outgoing list. Links from this node are preserved, even when ports are
224    /// renumbered by the insertion.
225    ///
226    /// Returns the range of newly created ports.
227    /// # Panics
228    ///
229    /// If the node is not in the graph.
230    fn insert_ports(
231        &mut self,
232        node: Node,
233        direction: Direction,
234        index: usize,
235        amount: usize,
236    ) -> Range<usize> {
237        panic_invalid_node(self, node);
238        self.hugr_mut().insert_ports(node, direction, index, amount)
239    }
240
241    /// Sets the parent of a node.
242    ///
243    /// The node becomes the parent's last child.
244    ///
245    /// # Panics
246    ///
247    /// If either the node or the parent is not in the graph.
248    fn set_parent(&mut self, node: Node, parent: Node) {
249        panic_invalid_node(self, parent);
250        panic_invalid_non_root(self, node);
251        self.hugr_mut().set_parent(node, parent);
252    }
253
254    /// Move a node in the hierarchy to be the subsequent sibling of another
255    /// node.
256    ///
257    /// The sibling node's parent becomes the new node's parent.
258    ///
259    /// The node becomes the parent's last child.
260    ///
261    /// # Panics
262    ///
263    /// If either node is not in the graph, or if it is a root.
264    fn move_after_sibling(&mut self, node: Node, after: Node) {
265        panic_invalid_non_root(self, node);
266        panic_invalid_non_root(self, after);
267        self.hugr_mut().move_after_sibling(node, after);
268    }
269
270    /// Move a node in the hierarchy to be the prior sibling of another node.
271    ///
272    /// The sibling node's parent becomes the new node's parent.
273    ///
274    /// The node becomes the parent's last child.
275    ///
276    /// # Panics
277    ///
278    /// If either node is not in the graph, or if it is a root.
279    fn move_before_sibling(&mut self, node: Node, before: Node) {
280        panic_invalid_non_root(self, node);
281        panic_invalid_non_root(self, before);
282        self.hugr_mut().move_before_sibling(node, before)
283    }
284
285    /// Replace the OpType at node and return the old OpType.
286    /// In general this invalidates the ports, which may need to be resized to
287    /// match the OpType signature.
288    ///
289    /// Returns the old OpType.
290    ///
291    /// TODO: Add a version which ignores input extensions
292    ///
293    /// # Errors
294    ///
295    /// Returns a [`HugrError::InvalidTag`] if this would break the bound
296    /// (`Self::RootHandle`) on the root node's OpTag.
297    ///
298    /// # Panics
299    ///
300    /// If the node is not in the graph.
301    fn replace_op(&mut self, node: Node, op: impl Into<OpType>) -> Result<OpType, HugrError> {
302        panic_invalid_node(self, node);
303        let op = op.into();
304        if node == self.root() && !Self::RootHandle::TAG.is_superset(op.tag()) {
305            return Err(HugrError::InvalidTag {
306                required: Self::RootHandle::TAG,
307                actual: op.tag(),
308            });
309        }
310        self.hugr_mut().replace_op(node, op)
311    }
312}
313
314/// Impl for non-wrapped Hugrs. Overwrites the recursive default-impls to directly use the hugr.
315impl<T: RootTagged<RootHandle = Node, Node = Node> + AsMut<Hugr>> HugrMutInternals for T {
316    fn hugr_mut(&mut self) -> &mut Hugr {
317        self.as_mut()
318    }
319
320    #[inline]
321    fn set_num_ports(&mut self, node: Node, incoming: usize, outgoing: usize) {
322        self.hugr_mut()
323            .graph
324            .set_num_ports(node.pg_index(), incoming, outgoing, |_, _| {})
325    }
326
327    fn insert_ports(
328        &mut self,
329        node: Node,
330        direction: Direction,
331        index: usize,
332        amount: usize,
333    ) -> Range<usize> {
334        let old_num_ports = self.base_hugr().graph.num_ports(node.pg_index(), direction);
335
336        self.add_ports(node, direction, amount as isize);
337
338        for swap_from_port in (index..old_num_ports).rev() {
339            let swap_to_port = swap_from_port + amount;
340            let [from_port_index, to_port_index] = [swap_from_port, swap_to_port].map(|p| {
341                self.base_hugr()
342                    .graph
343                    .port_index(node.pg_index(), PortOffset::new(direction, p))
344                    .unwrap()
345            });
346            let linked_ports = self
347                .base_hugr()
348                .graph
349                .port_links(from_port_index)
350                .map(|(_, to_subport)| to_subport.port())
351                .collect_vec();
352            self.hugr_mut().graph.unlink_port(from_port_index);
353            for linked_port_index in linked_ports {
354                let _ = self
355                    .hugr_mut()
356                    .graph
357                    .link_ports(to_port_index, linked_port_index)
358                    .expect("Ports exist");
359            }
360        }
361        index..index + amount
362    }
363
364    fn add_ports(&mut self, node: Node, direction: Direction, amount: isize) -> Range<usize> {
365        let mut incoming = self.hugr_mut().graph.num_inputs(node.pg_index());
366        let mut outgoing = self.hugr_mut().graph.num_outputs(node.pg_index());
367        let increment = |num: &mut usize| {
368            let new = num.saturating_add_signed(amount);
369            let range = *num..new;
370            *num = new;
371            range
372        };
373        let range = match direction {
374            Direction::Incoming => increment(&mut incoming),
375            Direction::Outgoing => increment(&mut outgoing),
376        };
377        self.hugr_mut()
378            .graph
379            .set_num_ports(node.pg_index(), incoming, outgoing, |_, _| {});
380        range
381    }
382
383    fn set_parent(&mut self, node: Node, parent: Node) {
384        self.hugr_mut().hierarchy.detach(node.pg_index());
385        self.hugr_mut()
386            .hierarchy
387            .push_child(node.pg_index(), parent.pg_index())
388            .expect("Inserting a newly-created node into the hierarchy should never fail.");
389    }
390
391    fn move_after_sibling(&mut self, node: Node, after: Node) {
392        self.hugr_mut().hierarchy.detach(node.pg_index());
393        self.hugr_mut()
394            .hierarchy
395            .insert_after(node.pg_index(), after.pg_index())
396            .expect("Inserting a newly-created node into the hierarchy should never fail.");
397    }
398
399    fn move_before_sibling(&mut self, node: Node, before: Node) {
400        self.hugr_mut().hierarchy.detach(node.pg_index());
401        self.hugr_mut()
402            .hierarchy
403            .insert_before(node.pg_index(), before.pg_index())
404            .expect("Inserting a newly-created node into the hierarchy should never fail.");
405    }
406
407    fn replace_op(&mut self, node: Node, op: impl Into<OpType>) -> Result<OpType, HugrError> {
408        // We know RootHandle=Node here so no need to check
409        let cur = self.hugr_mut().op_types.get_mut(node.pg_index());
410        Ok(std::mem::replace(cur, op.into()))
411    }
412}
413
414#[cfg(test)]
415mod test {
416    use crate::{
417        builder::{Container, DFGBuilder, Dataflow, DataflowHugr},
418        extension::prelude::Noop,
419        hugr::internal::HugrMutInternals as _,
420        ops::handle::NodeHandle,
421        types::{Signature, Type},
422        Direction, HugrView as _,
423    };
424
425    #[test]
426    fn insert_ports() {
427        let (nop, mut hugr) = {
428            let mut builder =
429                DFGBuilder::new(Signature::new_endo(Type::UNIT).with_prelude()).unwrap();
430            let [nop_in] = builder.input_wires_arr();
431            let nop = builder
432                .add_dataflow_op(Noop::new(Type::UNIT), [nop_in])
433                .unwrap();
434            builder.add_other_wire(nop.node(), builder.output().node());
435            let [nop_out] = nop.outputs_arr();
436            (
437                nop.node(),
438                builder.finish_hugr_with_outputs([nop_out]).unwrap(),
439            )
440        };
441        let [i, o] = hugr.get_io(hugr.root()).unwrap();
442        assert_eq!(0..2, hugr.insert_ports(nop, Direction::Incoming, 0, 2));
443        assert_eq!(1..3, hugr.insert_ports(nop, Direction::Outgoing, 1, 2));
444
445        assert_eq!(hugr.single_linked_input(i, 0), Some((nop, 2.into())));
446        assert_eq!(hugr.single_linked_output(o, 0), Some((nop, 0.into())));
447        assert_eq!(hugr.single_linked_output(o, 1), Some((nop, 3.into())));
448    }
449}
hugr_core/hugr/internal.rs

hugr_core/hugr/
internal.rs
