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::{OpTag, 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    /// Gets a mutable reference to the optype.
314    ///
315    /// Changing this may invalidate the ports, which may need to be resized to
316    /// match the OpType signature.
317    ///
318    /// Will panic for the root node unless [`Self::RootHandle`](RootTagged::RootHandle)
319    /// is [OpTag::Any], as mutation could invalidate the bound.
320    fn optype_mut(&mut self, node: Node) -> &mut OpType {
321        if Self::RootHandle::TAG.is_superset(OpTag::Any) {
322            panic_invalid_node(self, node);
323        } else {
324            panic_invalid_non_root(self, node);
325        }
326        self.hugr_mut().op_types.get_mut(node.pg_index())
327    }
328}
329
330/// Impl for non-wrapped Hugrs. Overwrites the recursive default-impls to directly use the hugr.
331impl<T: RootTagged<RootHandle = Node, Node = Node> + AsMut<Hugr>> HugrMutInternals for T {
332    fn hugr_mut(&mut self) -> &mut Hugr {
333        self.as_mut()
334    }
335
336    #[inline]
337    fn set_num_ports(&mut self, node: Node, incoming: usize, outgoing: usize) {
338        self.hugr_mut()
339            .graph
340            .set_num_ports(node.pg_index(), incoming, outgoing, |_, _| {})
341    }
342
343    fn insert_ports(
344        &mut self,
345        node: Node,
346        direction: Direction,
347        index: usize,
348        amount: usize,
349    ) -> Range<usize> {
350        let old_num_ports = self.base_hugr().graph.num_ports(node.pg_index(), direction);
351
352        self.add_ports(node, direction, amount as isize);
353
354        for swap_from_port in (index..old_num_ports).rev() {
355            let swap_to_port = swap_from_port + amount;
356            let [from_port_index, to_port_index] = [swap_from_port, swap_to_port].map(|p| {
357                self.base_hugr()
358                    .graph
359                    .port_index(node.pg_index(), PortOffset::new(direction, p))
360                    .unwrap()
361            });
362            let linked_ports = self
363                .base_hugr()
364                .graph
365                .port_links(from_port_index)
366                .map(|(_, to_subport)| to_subport.port())
367                .collect_vec();
368            self.hugr_mut().graph.unlink_port(from_port_index);
369            for linked_port_index in linked_ports {
370                let _ = self
371                    .hugr_mut()
372                    .graph
373                    .link_ports(to_port_index, linked_port_index)
374                    .expect("Ports exist");
375            }
376        }
377        index..index + amount
378    }
379
380    fn add_ports(&mut self, node: Node, direction: Direction, amount: isize) -> Range<usize> {
381        let mut incoming = self.hugr_mut().graph.num_inputs(node.pg_index());
382        let mut outgoing = self.hugr_mut().graph.num_outputs(node.pg_index());
383        let increment = |num: &mut usize| {
384            let new = num.saturating_add_signed(amount);
385            let range = *num..new;
386            *num = new;
387            range
388        };
389        let range = match direction {
390            Direction::Incoming => increment(&mut incoming),
391            Direction::Outgoing => increment(&mut outgoing),
392        };
393        self.hugr_mut()
394            .graph
395            .set_num_ports(node.pg_index(), incoming, outgoing, |_, _| {});
396        range
397    }
398
399    fn set_parent(&mut self, node: Node, parent: Node) {
400        self.hugr_mut().hierarchy.detach(node.pg_index());
401        self.hugr_mut()
402            .hierarchy
403            .push_child(node.pg_index(), parent.pg_index())
404            .expect("Inserting a newly-created node into the hierarchy should never fail.");
405    }
406
407    fn move_after_sibling(&mut self, node: Node, after: Node) {
408        self.hugr_mut().hierarchy.detach(node.pg_index());
409        self.hugr_mut()
410            .hierarchy
411            .insert_after(node.pg_index(), after.pg_index())
412            .expect("Inserting a newly-created node into the hierarchy should never fail.");
413    }
414
415    fn move_before_sibling(&mut self, node: Node, before: Node) {
416        self.hugr_mut().hierarchy.detach(node.pg_index());
417        self.hugr_mut()
418            .hierarchy
419            .insert_before(node.pg_index(), before.pg_index())
420            .expect("Inserting a newly-created node into the hierarchy should never fail.");
421    }
422
423    fn replace_op(&mut self, node: Node, op: impl Into<OpType>) -> Result<OpType, HugrError> {
424        // We know RootHandle=Node here so no need to check
425        Ok(std::mem::replace(self.optype_mut(node), op.into()))
426    }
427}
428
429#[cfg(test)]
430mod test {
431    use crate::{
432        builder::{Container, DFGBuilder, Dataflow, DataflowHugr},
433        extension::prelude::Noop,
434        hugr::internal::HugrMutInternals as _,
435        ops::handle::NodeHandle,
436        types::{Signature, Type},
437        Direction, HugrView as _,
438    };
439
440    #[test]
441    fn insert_ports() {
442        let (nop, mut hugr) = {
443            let mut builder =
444                DFGBuilder::new(Signature::new_endo(Type::UNIT).with_prelude()).unwrap();
445            let [nop_in] = builder.input_wires_arr();
446            let nop = builder
447                .add_dataflow_op(Noop::new(Type::UNIT), [nop_in])
448                .unwrap();
449            builder.add_other_wire(nop.node(), builder.output().node());
450            let [nop_out] = nop.outputs_arr();
451            (
452                nop.node(),
453                builder.finish_hugr_with_outputs([nop_out]).unwrap(),
454            )
455        };
456        let [i, o] = hugr.get_io(hugr.root()).unwrap();
457        assert_eq!(0..2, hugr.insert_ports(nop, Direction::Incoming, 0, 2));
458        assert_eq!(1..3, hugr.insert_ports(nop, Direction::Outgoing, 1, 2));
459
460        assert_eq!(hugr.single_linked_input(i, 0), Some((nop, 2.into())));
461        assert_eq!(hugr.single_linked_output(o, 0), Some((nop, 0.into())));
462        assert_eq!(hugr.single_linked_output(o, 1), Some((nop, 3.into())));
463    }
464}
hugr_core/hugr/internal.rs

hugr_core/hugr/
internal.rs
