hugr_core/hugr/views/
rerooted.rs

1//! A HUGR wrapper with a modified entrypoint node, returned by
2//! [`HugrView::with_entrypoint`] and [`HugrMut::with_entrypoint_mut`].
3
4use crate::hugr::HugrMut;
5use crate::hugr::internal::{HugrInternals, HugrMutInternals};
6
7use super::{HugrView, panic_invalid_node};
8
9/// A HUGR wrapper with a modified entrypoint node.
10///
11/// All nodes from the original are still present, but the main entrypoint used
12/// for traversals and optimizations is altered.
13#[derive(Clone)]
14pub struct Rerooted<H: HugrView> {
15    hugr: H,
16    entrypoint: H::Node,
17}
18
19impl<H: HugrView> Rerooted<H> {
20    /// Create a hierarchical view of a whole HUGR
21    ///
22    /// # Panics
23    ///
24    /// If the new entrypoint is not in the HUGR.
25    ///
26    /// [`OpTag`]: crate::ops::OpTag
27    pub fn new(hugr: H, entrypoint: H::Node) -> Self {
28        panic_invalid_node(&hugr, entrypoint);
29        Self { hugr, entrypoint }
30    }
31
32    /// Returns the HUGR wrapped in this view.
33    pub fn into_unwrapped(self) -> H {
34        self.hugr
35    }
36}
37
38impl<H: HugrView> HugrInternals for Rerooted<H> {
39    type RegionPortgraph<'p>
40        = H::RegionPortgraph<'p>
41    where
42        Self: 'p;
43
44    type Node = H::Node;
45
46    type RegionPortgraphNodes = H::RegionPortgraphNodes;
47
48    super::impls::hugr_internal_methods! {this, &this.hugr}
49}
50
51impl<H: HugrView> HugrView for Rerooted<H> {
52    #[inline]
53    fn entrypoint(&self) -> Self::Node {
54        self.entrypoint
55    }
56
57    #[inline]
58    fn entrypoint_optype(&self) -> &crate::ops::OpType {
59        self.hugr.get_optype(self.entrypoint)
60    }
61
62    fn mermaid_string_with_formatter(
63        &self,
64        formatter: crate::hugr::views::render::MermaidFormatter<Self>,
65    ) -> String {
66        self.hugr
67            .mermaid_string_with_formatter(formatter.with_hugr(&self.hugr))
68    }
69
70    delegate::delegate! {
71        to (&self.hugr) {
72                fn module_root(&self) -> Self::Node;
73                fn contains_node(&self, node: Self::Node) -> bool;
74                fn get_parent(&self, node: Self::Node) -> Option<Self::Node>;
75                fn get_metadata(&self, node: Self::Node, key: impl AsRef<str>) -> Option<&crate::hugr::NodeMetadata>;
76                fn get_optype(&self, node: Self::Node) -> &crate::ops::OpType;
77                fn num_nodes(&self) -> usize;
78                fn num_edges(&self) -> usize;
79                fn num_ports(&self, node: Self::Node, dir: crate::Direction) -> usize;
80                fn num_inputs(&self, node: Self::Node) -> usize;
81                fn num_outputs(&self, node: Self::Node) -> usize;
82                fn nodes(&self) -> impl Iterator<Item = Self::Node> + Clone;
83                fn node_ports(&self, node: Self::Node, dir: crate::Direction) -> impl Iterator<Item = crate::Port> + Clone;
84                fn node_outputs(&self, node: Self::Node) -> impl Iterator<Item = crate::OutgoingPort> + Clone;
85                fn node_inputs(&self, node: Self::Node) -> impl Iterator<Item = crate::IncomingPort> + Clone;
86                fn all_node_ports(&self, node: Self::Node) -> impl Iterator<Item = crate::Port> + Clone;
87                fn linked_ports(&self, node: Self::Node, port: impl Into<crate::Port>) -> impl Iterator<Item = (Self::Node, crate::Port)> + Clone;
88                fn all_linked_ports(&self, node: Self::Node, dir: crate::Direction) -> itertools::Either<impl Iterator<Item = (Self::Node, crate::OutgoingPort)>, impl Iterator<Item = (Self::Node, crate::IncomingPort)>>;
89                fn all_linked_outputs(&self, node: Self::Node) -> impl Iterator<Item = (Self::Node, crate::OutgoingPort)>;
90                fn all_linked_inputs(&self, node: Self::Node) -> impl Iterator<Item = (Self::Node, crate::IncomingPort)>;
91                fn single_linked_port(&self, node: Self::Node, port: impl Into<crate::Port>) -> Option<(Self::Node, crate::Port)>;
92                fn single_linked_output(&self, node: Self::Node, port: impl Into<crate::IncomingPort>) -> Option<(Self::Node, crate::OutgoingPort)>;
93                fn single_linked_input(&self, node: Self::Node, port: impl Into<crate::OutgoingPort>) -> Option<(Self::Node, crate::IncomingPort)>;
94                fn linked_outputs(&self, node: Self::Node, port: impl Into<crate::IncomingPort>) -> impl Iterator<Item = (Self::Node, crate::OutgoingPort)>;
95                fn linked_inputs(&self, node: Self::Node, port: impl Into<crate::OutgoingPort>) -> impl Iterator<Item = (Self::Node, crate::IncomingPort)>;
96                fn node_connections(&self, node: Self::Node, other: Self::Node) -> impl Iterator<Item = [crate::Port; 2]> + Clone;
97                fn is_linked(&self, node: Self::Node, port: impl Into<crate::Port>) -> bool;
98                fn children(&self, node: Self::Node) -> impl DoubleEndedIterator<Item = Self::Node> + Clone;
99                fn descendants(&self, node: Self::Node) -> impl Iterator<Item = Self::Node> + Clone;
100                fn first_child(&self, node: Self::Node) -> Option<Self::Node>;
101                fn neighbours(&self, node: Self::Node, dir: crate::Direction) -> impl Iterator<Item = Self::Node> + Clone;
102                fn all_neighbours(&self, node: Self::Node) -> impl Iterator<Item = Self::Node> + Clone;
103                #[allow(deprecated)]
104                fn mermaid_string_with_config(&self, config: crate::hugr::views::render::RenderConfig<Self::Node>) -> String;
105                fn dot_string(&self) -> String;
106                fn static_source(&self, node: Self::Node) -> Option<Self::Node>;
107                fn static_targets(&self, node: Self::Node) -> Option<impl Iterator<Item = (Self::Node, crate::IncomingPort)>>;
108                fn value_types(&self, node: Self::Node, dir: crate::Direction) -> impl Iterator<Item = (crate::Port, crate::types::Type)>;
109                fn extensions(&self) -> &crate::extension::ExtensionRegistry;
110                fn validate(&self) -> Result<(), crate::hugr::ValidationError<Self::Node>>;
111                fn extract_hugr(&self, parent: Self::Node) -> (crate::Hugr, impl crate::hugr::views::ExtractionResult<Self::Node> + 'static);
112        }
113    }
114}
115
116impl<H: HugrMutInternals> HugrMutInternals for Rerooted<H> {
117    super::impls::hugr_mut_internal_methods! {this, &mut this.hugr}
118}
119
120impl<H: HugrMut> HugrMut for Rerooted<H> {
121    fn set_entrypoint(&mut self, root: Self::Node) {
122        self.entrypoint = root;
123        self.hugr.set_entrypoint(root);
124    }
125
126    delegate::delegate! {
127        to (&mut self.hugr) {
128                fn get_metadata_mut(&mut self, node: Self::Node, key: impl AsRef<str>) -> &mut crate::hugr::NodeMetadata;
129                fn set_metadata(&mut self, node: Self::Node, key: impl AsRef<str>, metadata: impl Into<crate::hugr::NodeMetadata>);
130                fn remove_metadata(&mut self, node: Self::Node, key: impl AsRef<str>);
131                fn add_node_with_parent(&mut self, parent: Self::Node, op: impl Into<crate::ops::OpType>) -> Self::Node;
132                fn add_node_before(&mut self, sibling: Self::Node, nodetype: impl Into<crate::ops::OpType>) -> Self::Node;
133                fn add_node_after(&mut self, sibling: Self::Node, op: impl Into<crate::ops::OpType>) -> Self::Node;
134                fn remove_node(&mut self, node: Self::Node) -> crate::ops::OpType;
135                fn remove_subtree(&mut self, node: Self::Node);
136                fn copy_descendants(&mut self, root: Self::Node, new_parent: Self::Node, subst: Option<crate::types::Substitution>) -> std::collections::BTreeMap<Self::Node, Self::Node>;
137                fn connect(&mut self, src: Self::Node, src_port: impl Into<crate::OutgoingPort>, dst: Self::Node, dst_port: impl Into<crate::IncomingPort>);
138                fn disconnect(&mut self, node: Self::Node, port: impl Into<crate::Port>);
139                fn add_other_edge(&mut self, src: Self::Node, dst: Self::Node) -> (crate::OutgoingPort, crate::IncomingPort);
140                fn insert_hugr(&mut self, root: Self::Node, other: crate::Hugr) -> crate::hugr::hugrmut::InsertionResult<crate::Node, Self::Node>;
141                fn insert_from_view<Other: crate::hugr::HugrView>(&mut self, root: Self::Node, other: &Other) -> crate::hugr::hugrmut::InsertionResult<Other::Node, Self::Node>;
142                fn insert_subgraph<Other: crate::hugr::HugrView>(&mut self, root: Self::Node, other: &Other, subgraph: &crate::hugr::views::SiblingSubgraph<Other::Node>) -> std::collections::HashMap<Other::Node, Self::Node>;
143                fn use_extension(&mut self, extension: impl Into<std::sync::Arc<crate::extension::Extension>>);
144                fn use_extensions<Reg>(&mut self, registry: impl IntoIterator<Item = Reg>) where crate::extension::ExtensionRegistry: Extend<Reg>;
145        }
146    }
147}
148
149#[cfg(test)]
150mod test {
151    use crate::builder::test::simple_cfg_hugr;
152    use crate::builder::{Dataflow, FunctionBuilder, HugrBuilder, SubContainer};
153    use crate::hugr::HugrMut;
154    use crate::hugr::internal::HugrMutInternals;
155    use crate::hugr::views::ExtractionResult;
156    use crate::ops::handle::NodeHandle;
157    use crate::ops::{DataflowBlock, OpType};
158    use crate::{HugrView, type_row, types::Signature};
159
160    #[test]
161    fn rerooted() {
162        let mut builder = FunctionBuilder::new("main", Signature::new(vec![], vec![])).unwrap();
163        let dfg = builder
164            .dfg_builder_endo([])
165            .unwrap()
166            .finish_sub_container()
167            .unwrap()
168            .node();
169        let mut h = builder.finish_hugr().unwrap();
170        let _func = h.entrypoint();
171
172        // Immutable wrappers
173        let dfg_v = h.with_entrypoint(dfg);
174        assert_eq!(dfg_v.module_root(), h.module_root());
175        assert_eq!(dfg_v.entrypoint(), dfg);
176        assert!(dfg_v.entrypoint_optype().is_dfg());
177        assert!(dfg_v.get_optype(dfg_v.module_root().node()).is_module());
178
179        // Mutable wrappers
180        let mut dfg_v = h.with_entrypoint_mut(dfg);
181        {
182            // That is a HugrMutInternal, so we can try:
183            let root = dfg_v.entrypoint();
184            let bb: OpType = DataflowBlock {
185                inputs: type_row![],
186                other_outputs: type_row![],
187                sum_rows: vec![type_row![]],
188            }
189            .into();
190            dfg_v.replace_op(root, bb.clone());
191
192            assert!(dfg_v.entrypoint_optype().is_dataflow_block());
193            assert!(dfg_v.get_optype(dfg_v.module_root().node()).is_module());
194        }
195        // That modified the original HUGR
196        assert!(h.get_optype(dfg).is_dataflow_block());
197        assert!(h.entrypoint_optype().is_func_defn());
198        assert!(h.get_optype(h.module_root().node()).is_module());
199    }
200
201    #[test]
202    fn extract_rerooted() {
203        let mut hugr = simple_cfg_hugr();
204        let cfg = hugr.entrypoint();
205        let basic_block = hugr.first_child(cfg).unwrap();
206        hugr.set_entrypoint(basic_block);
207        assert!(hugr.get_optype(hugr.entrypoint()).is_dataflow_block());
208
209        let rerooted = hugr.with_entrypoint(cfg);
210        assert!(rerooted.get_optype(rerooted.entrypoint()).is_cfg());
211
212        // Extract the basic block
213        let (extracted_hugr, map) = rerooted.extract_hugr(basic_block);
214        let extracted_cfg = map.extracted_node(cfg);
215        let extracted_bb = map.extracted_node(basic_block);
216        assert_eq!(extracted_hugr.entrypoint(), extracted_bb);
217        assert!(extracted_hugr.get_optype(extracted_cfg).is_cfg());
218        assert_eq!(
219            extracted_hugr.first_child(extracted_cfg),
220            Some(extracted_bb)
221        );
222        assert!(extracted_hugr.get_optype(extracted_bb).is_dataflow_block());
223
224        // Extract the cfg (and current entrypoint)
225        let (extracted_hugr, map) = rerooted.extract_hugr(cfg);
226        let extracted_cfg = map.extracted_node(cfg);
227        let extracted_bb = map.extracted_node(basic_block);
228        assert_eq!(extracted_hugr.entrypoint(), extracted_cfg);
229        assert!(extracted_hugr.get_optype(extracted_cfg).is_cfg());
230        assert_eq!(
231            extracted_hugr.first_child(extracted_cfg),
232            Some(extracted_bb)
233        );
234        assert!(extracted_hugr.get_optype(extracted_bb).is_dataflow_block());
235    }
236
237    #[test]
238    fn mermaid_format() {
239        let h = simple_cfg_hugr();
240        let rerooted = h.with_entrypoint(h.entrypoint());
241        let mermaid_str = rerooted.mermaid_format().finish();
242        assert_eq!(mermaid_str, h.mermaid_format().finish());
243    }
244}
hugr_core/hugr/views/rerooted.rs

hugr_core/hugr/views/
rerooted.rs
