feat: Implement petgraph traits on Hugr

src/hugr/views.rs

@@ -1,6 +1,7 @@
 //! Read-only access into HUGR graphs and subgraphs.
 
 pub mod hierarchy;
+pub mod petgraph;
 pub mod sibling;
 
 #[cfg(test)]

src/hugr/views/hierarchy.rs

@@ -14,8 +14,6 @@
 //! interchangeably. They implement [`HugrView`] as well as petgraph's _visit_
 //! traits.
 
-pub mod petgraph;
-
 use std::iter;
 
 use ::petgraph::visit as pv;
@@ -27,6 +25,8 @@ use crate::ops::handle::NodeHandle;
 use crate::ops::OpTrait;
 use crate::{hugr::NodeType, hugr::OpType, Direction, Hugr, Node, Port};
 
+use super::petgraph::PetgraphHugr;
+
 use super::{sealed::HugrInternals, HugrView, NodeMetadata};
 
 type FlatRegionGraph<'g> = portgraph::view::FlatRegion<'g, &'g MultiPortGraph>;
@@ -406,16 +406,7 @@ where
 }
 
 /// A common trait for views of a HUGR hierarchical subgraph.
-pub trait HierarchyView<'a>:
-    HugrView
-    + pv::GraphBase<NodeId = Node>
-    + pv::GraphProp
-    + pv::NodeCount
-    + pv::NodeIndexable
-    + pv::EdgeCount
-    + pv::Visitable
-    + pv::GetAdjacencyMatrix
-    + pv::Visitable
+pub trait HierarchyView<'a>: HugrView + PetgraphHugr
 where
     for<'g> &'g Self: pv::IntoNeighborsDirected + pv::IntoNodeIdentifiers,
 {

src/hugr/views/hierarchy/petgraph.rs → src/hugr/views/petgraph.rs

@@ -6,12 +6,141 @@ use crate::hugr::HugrView;
 use crate::ops::handle::NodeHandle;
 use crate::ops::OpType;
 use crate::types::EdgeKind;
-use crate::{Node, Port};
+use crate::{Hugr, Node, Port};
 
 use context_iterators::{ContextIterator, IntoContextIterator, MapWithCtx};
 use petgraph::visit as pv;
 use portgraph::NodeIndex;
 
+/// A trait grouping the multiple petgraph traits implemented on Hugrs
+//
+// TODO: We could just require these in `HugrView`.
+pub trait PetgraphHugr:
+    HugrView
+    + pv::GraphBase<NodeId = Node>
+    + pv::GraphProp
+    + pv::NodeCount
+    + pv::NodeIndexable
+    + pv::EdgeCount
+    + pv::Visitable
+    + pv::GetAdjacencyMatrix
+    + pv::Visitable
+where
+    for<'g> &'g Self: pv::IntoNeighborsDirected + pv::IntoNodeIdentifiers,
+{
+}
+
+impl pv::GraphBase for Hugr {
+    type NodeId = Node;
+    type EdgeId = ((Node, Port), (Node, Port));
+}
+
+impl pv::GraphProp for Hugr {
+    type EdgeType = petgraph::Directed;
+}
+
+impl pv::NodeCount for Hugr {
+    fn node_count(&self) -> usize {
+        HugrView::node_count(self)
+    }
+}
+
+impl pv::NodeIndexable for Hugr {
+    fn node_bound(&self) -> usize {
+        HugrView::node_count(self)
+    }
+
+    fn to_index(&self, ix: Self::NodeId) -> usize {
+        ix.index.into()
+    }
+
+    fn from_index(&self, ix: usize) -> Self::NodeId {
+        NodeIndex::new(ix).into()
+    }
+}
+
+impl pv::EdgeCount for Hugr {
+    fn edge_count(&self) -> usize {
+        HugrView::edge_count(self)
+    }
+}
+
+impl pv::Data for Hugr {
+    type NodeWeight = OpType;
+    type EdgeWeight = EdgeKind;
+}
+
+impl<'g> pv::IntoNodeIdentifiers for &'g Hugr {
+    type NodeIdentifiers = <Hugr as HugrView>::Nodes<'g>;
+
+    fn node_identifiers(self) -> Self::NodeIdentifiers {
+        self.nodes()
+    }
+}
+
+impl<'g> pv::IntoNodeReferences for &'g Hugr {
+    type NodeRef = HugrNodeRef<'g>;
+    type NodeReferences = MapWithCtx<<Hugr as HugrView>::Nodes<'g>, Self, HugrNodeRef<'g>>;
+
+    fn node_references(self) -> Self::NodeReferences {
+        self.nodes()
+            .with_context(self)
+            .map_with_context(|n, &hugr| HugrNodeRef::from_node(n, hugr))
+    }
+}
+
+impl<'g> pv::IntoNeighbors for &'g Hugr {
+    type Neighbors = <Hugr as HugrView>::Neighbours<'g>;
+
+    fn neighbors(self, n: Self::NodeId) -> Self::Neighbors {
+        self.output_neighbours(n)
+    }
+}
+
+impl<'g> pv::IntoNeighborsDirected for &'g Hugr {
+    type NeighborsDirected = <Hugr as HugrView>::Neighbours<'g>;
+
+    fn neighbors_directed(
+        self,
+        n: Self::NodeId,
+        d: petgraph::Direction,
+    ) -> Self::NeighborsDirected {
+        self.neighbours(n, d.into())
+    }
+}
+
+impl pv::Visitable for Hugr {
+    type Map = std::collections::HashSet<Self::NodeId>;
+
+    fn visit_map(&self) -> Self::Map {
+        std::collections::HashSet::new()
+    }
+
+    fn reset_map(&self, map: &mut Self::Map) {
+        map.clear();
+    }
+}
+
+impl pv::GetAdjacencyMatrix for Hugr {
+    type AdjMatrix = std::collections::HashSet<(Self::NodeId, Self::NodeId)>;
+
+    fn adjacency_matrix(&self) -> Self::AdjMatrix {
+        let mut matrix = std::collections::HashSet::new();
+        for node in self.nodes() {
+            for neighbour in self.output_neighbours(node) {
+                matrix.insert((node, neighbour));
+            }
+        }
+        matrix
+    }
+
+    fn is_adjacent(&self, matrix: &Self::AdjMatrix, a: Self::NodeId, b: Self::NodeId) -> bool {
+        matrix.contains(&(a, b))
+    }
+}
+
+impl PetgraphHugr for Hugr {}
+
 macro_rules! impl_region_petgraph_traits {
     ($hugr:ident) => {
         impl<'a, Root, Base> pv::GraphBase for $hugr<'a, Root, Base>
@@ -177,6 +306,13 @@ macro_rules! impl_region_petgraph_traits {
                 matrix.contains(&(a, b))
             }
         }
+
+        impl<'a, Root, Base> PetgraphHugr for $hugr<'a, Root, Base>
+        where
+            Root: NodeHandle,
+            Base: HugrInternals + HugrView,
+        {
+        }
     };
 }