Edge representation in model: make ports refer to the id of the edge.

hugr-core/src/export.rs

@@ -8,10 +8,10 @@ use crate::{
         CustomType, FuncTypeBase, MaybeRV, PolyFuncTypeBase, RowVariable, SumType, TypeArg,
         TypeBase, TypeEnum,
     },
-    Direction, Hugr, HugrView, Node, Port,
+    Hugr, HugrView, Node, Port,
 };
 use hugr_model::v0 as model;
-use indexmap::IndexMap;
+use indexmap::IndexSet;
 use smol_str::ToSmolStr;
 use tinyvec::TinyVec;
 
@@ -30,15 +30,15 @@ struct Context<'a> {
     module: model::Module,
     /// Mapping from ports to edge indices.
     /// This only includes the minimum port among groups of linked ports.
-    edges: IndexMap<(Node, Port), usize>,
+    edges: IndexSet<(Node, Port)>,
 }
 
 impl<'a> Context<'a> {
     pub fn new(hugr: &'a Hugr) -> Self {
         Self {
             hugr,
             module: model::Module::default(),
-            edges: IndexMap::new(),
+            edges: IndexSet::new(),
         }
     }
 
@@ -54,34 +54,32 @@ impl<'a> Context<'a> {
         model::NodeId(index as _)
     }
 
+    /// Returns the edge id for a given port, creating a new edge if necessary.
+    ///
+    /// Any two ports that are linked will be represented by the same edge.
+    fn get_edge_id(&mut self, node: Node, port: Port) -> model::EdgeId {
+        // To ensure that linked ports are represented by the same edge, we take the minimum port
+        // among all the linked ports, including the one we started with.
+        let linked_ports = self.hugr.linked_ports(node, port);
+        let all_ports = std::iter::once((node, port)).chain(linked_ports);
+        let repr = all_ports.min().unwrap();
+        let edge = self.edges.insert_full(repr).1 as _;
+        model::EdgeId(edge)
+    }
+
     pub fn make_port(&mut self, node: Node, port: impl Into<Port>) -> model::PortId {
         let port = port.into();
         let index = self.module.ports.len();
         let port_id = model::PortId(index as _);
         let r#type = self.make_term(model::Term::Wildcard); // TODO
+        let edge = self.get_edge_id(node, port);
+
         self.module.ports.push(model::Port {
             r#type,
             meta: Vec::new(),
+            edge,
         });
 
-        // To ensure that linked ports are represented by the same edge, we take the minimum port
-        // among all the linked ports, including the one we started with.
-        let linked_ports = self.hugr.linked_ports(node, port);
-        let all_ports = std::iter::once((node, port)).chain(linked_ports);
-        let repr = all_ports.min().unwrap();
-
-        let edge_id = *self.edges.entry(repr).or_insert_with(|| {
-            let edge_id = self.module.edges.len();
-            let edge = model::Edge::default();
-            self.module.edges.push(edge);
-            edge_id
-        });
-
-        match port.direction() {
-            Direction::Incoming => self.module.edges[edge_id].inputs.push(port_id),
-            Direction::Outgoing => self.module.edges[edge_id].outputs.push(port_id),
-        }
-
         port_id
     }
 

hugr-model/src/v0/mod.rs

@@ -34,6 +34,10 @@ pub struct NodeId(pub u32);
 #[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash, Default)]
 pub struct PortId(pub u32);
 
+/// Index of an edge in a hugr graph.
+#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash, Default)]
+pub struct EdgeId(pub u32);
+
 /// Index of a term in a hugr graph.
 #[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash, Default)]
 pub struct TermId(pub u32);
@@ -61,8 +65,6 @@ pub struct Module {
     pub nodes: Vec<Node>,
     /// Table of [`Port`]s.
     pub ports: Vec<Port>,
-    /// Table of [`Edge`]s.
-    pub edges: Vec<Edge>,
     /// Table of [`Term`]s.
     pub terms: Vec<Term>,
 }
@@ -195,6 +197,10 @@ pub struct MetaItem {
 /// A port in the graph.
 #[derive(Debug, Clone, PartialEq, Eq, Hash)]
 pub struct Port {
+    /// The id of the edge that the port is connected to.
+    /// All ports referencing the same edge are connected.
+    pub edge: EdgeId,
+
     /// The type of the port.
     ///
     /// This must be a term of type `Type`.
@@ -205,18 +211,6 @@ pub struct Port {
     pub meta: Vec<MetaItem>,
 }
 
-/// An edge in the graph.
-///
-/// An edge connects input ports to output ports.
-/// A port may only be part of at most one edge.
-#[derive(Debug, Clone, Default, PartialEq, Eq, Hash)]
-pub struct Edge {
-    /// The input ports of the edge.
-    pub inputs: TinyVec<[PortId; 3]>,
-    /// The output ports of the edge.
-    pub outputs: TinyVec<[PortId; 3]>,
-}
-
 /// Schemes are parameterized terms.
 #[derive(Debug, Clone, PartialEq, Eq, Hash)]
 pub struct Scheme {