feat: Proptest for Multiportgraph

src/proptest.rs

@@ -2,72 +2,81 @@
 //!
 //! Currently, this module exposes a single function [`gen_portgraph`], which
 //! returns strategies that generate random portgraphs.
-use std::iter::zip;
+use std::collections::{BTreeMap, HashSet};
 
-use crate::{Direction, LinkMut, NodeIndex, PortGraph, PortIndex, PortMut, PortView};
+use crate::{LinkMut, MultiPortGraph, NodeIndex, PortGraph, PortMut, PortOffset, PortView};
 use proptest::prelude::*;
-use rand::seq::SliceRandom;
 
+#[derive(Debug, Clone, Copy)]
+struct Edge {
+    src_v: NodeIndex,
+    target_v: NodeIndex,
+    src_p: PortOffset,
+    target_p: PortOffset,
+}
 prop_compose! {
-    /// A random graph with random number of ports but no edges
-    ///
-    /// The graph has at least one node.
-    ///
-    ///  - `max_n_nodes` is the maximum number of nodes in the generated graph
-    ///  - `max_port` is the maximum number of incoming and outgoing ports at
-    ///    every node
-    fn gen_no_edge_graph(max_n_nodes: usize, max_port: usize)(
-        ports in prop::collection::vec(0..=max_port, 2..=2*max_n_nodes)
-    ) -> PortGraph {
-        let mut g = PortGraph::new();
-        for ind in (0..ports.len() - 1).step_by(2) {
-            g.add_node(ports[ind], ports[ind + 1]);
-        }
-        g
+    /// A random set of edges for a portgraph.
+    fn gen_edges(max_n_nodes: usize, max_n_ports: usize, max_n_edges: usize)(
+        n_edges in 0..=max_n_edges
+    )(
+        out_ports in vec![(1..max_n_nodes, 0..max_n_ports); n_edges],
+        in_ports in vec![(1..max_n_nodes, 0..max_n_ports); n_edges],
+    ) -> Vec<Edge> {
+        out_ports.into_iter().zip(in_ports).map(|((src_v, src_p), (target_v, target_p))| {
+            Edge {
+                src_v: NodeIndex::new(src_v),
+                target_v: NodeIndex::new(target_v),
+                src_p: PortOffset::new_outgoing(src_p),
+                target_p: PortOffset::new_incoming(target_p),
+            }
+        }).collect()
     }
 }
 
-/// A random graph and random edge lists
-///
-/// Returns a tuple of
-///  - `graph`, a graph with ports but without edges
-///  - `in_ports`, a list of incoming ports
-///  - `out_port`, a list of outgoing ports, of the same length as `in_ports`
-fn gen_graph_and_edges(
-    max_n_nodes: usize,
-    max_port: usize,
-    max_n_edges: usize,
-) -> impl Strategy<Value = (Vec<PortIndex>, Vec<PortIndex>, PortGraph)> {
-    let graph = gen_no_edge_graph(max_n_nodes, max_port);
-    (0..=max_n_edges, graph).prop_perturb(|(mut n_edges, graph), mut rng| {
-        let (mut in_ports, mut out_ports): (Vec<_>, Vec<_>) = graph
-            .ports_iter()
-            .partition(|p| graph.port_direction(*p).unwrap() == Direction::Incoming);
-        in_ports.shuffle(&mut rng);
-        out_ports.shuffle(&mut rng);
-
-        n_edges = [n_edges, in_ports.len(), out_ports.len()]
-            .into_iter()
-            .min()
-            .unwrap();
-        in_ports.truncate(n_edges);
-        out_ports.truncate(n_edges);
-        (in_ports, out_ports, graph)
-    })
+prop_compose! {
+    /// A random set of distinct edges for a portgraph.
+    fn gen_unique_edges(max_n_nodes: usize, max_n_ports: usize, max_n_edges: usize)(
+        out_ports in prop::collection::hash_set((1..max_n_nodes, 0..max_n_ports), 0..=max_n_edges),
+        in_ports in prop::collection::hash_set((1..max_n_nodes, 0..max_n_ports), 0..=max_n_edges),
+    ) -> Vec<Edge> {
+        // Unlike the non-unique case, in this case we do not "control" the number of
+        // edges directly. Might make shrinking slightly more difficult.
+        out_ports.into_iter().zip(in_ports).map(|((src_v, src_p), (target_v, target_p))| {
+            Edge {
+                src_v: NodeIndex::new(src_v),
+                target_v: NodeIndex::new(target_v),
+                src_p: PortOffset::new_outgoing(src_p),
+                target_p: PortOffset::new_incoming(target_p),
+            }
+        }).collect()
+    }
 }
 
 prop_compose! {
     /// A random non-empty portgraph
     ///
     /// With at least 1 and at most `max_n_nodes` nodes, with at most `max_port`
     /// incoming and outgoing ports at ever node, and at most `max_n_edges`.
-    pub fn gen_portgraph(max_n_nodes: usize, max_port: usize, max_n_edges: usize)(
-        (in_stubs, out_stubs, mut graph) in gen_graph_and_edges(max_n_nodes, max_port, max_n_edges)
+    pub fn gen_portgraph(max_n_nodes: usize, max_n_ports: usize, max_n_edges: usize)(
+        edges in gen_unique_edges(max_n_nodes, max_n_ports, max_n_edges)
     ) -> PortGraph {
-        for (incoming, outgoing) in zip(in_stubs, out_stubs) {
-            graph.link_ports(outgoing, incoming).unwrap();
-        }
-        graph
+        let mut g: PortGraph = graph_from_edges(&edges);
+        ensure_non_empty(&mut g);
+        g
+    }
+}
+
+prop_compose! {
+    /// A random non-empty multiportgraph
+    ///
+    /// With at least 1 and at most `max_n_nodes` nodes, with at most `max_port`
+    /// incoming and outgoing ports at ever node, and at most `max_n_edges`.
+    pub fn gen_multiportgraph(max_n_nodes: usize, max_n_ports: usize, max_n_edges: usize)(
+        edges in gen_edges(max_n_nodes, max_n_ports, max_n_edges)
+    ) -> MultiPortGraph {
+        let mut g: MultiPortGraph = graph_from_edges(&edges);
+        ensure_non_empty(&mut g);
+        g
     }
 }
 
@@ -97,6 +106,58 @@ where
     })
 }
 
+fn graph_from_edges<G: PortMut + LinkMut + Default>(edges: &[Edge]) -> G {
+    let mut map_vertices = BTreeMap::new();
+    let mut in_port_counts = BTreeMap::new();
+    let mut out_port_counts = BTreeMap::new();
+    let mut add_port = |v, p| {
+        let (max, port) = match p {
+            PortOffset::Incoming(p) => (in_port_counts.entry(v).or_insert(0), p),
+            PortOffset::Outgoing(p) => (out_port_counts.entry(v).or_insert(0), p),
+        };
+        *max = (*max).max(port as usize + 1);
+    };
+    for &Edge {
+        src_v,
+        target_v,
+        src_p,
+        target_p,
+    } in edges
+    {
+        add_port(src_v, src_p);
+        add_port(target_v, target_p);
+    }
+    let mut g = G::default();
+    let vertices: HashSet<_> = edges
+        .iter()
+        .flat_map(|e| vec![e.src_v, e.target_v])
+        .collect();
+    for v in vertices {
+        let in_ports = in_port_counts.get(&v).copied().unwrap_or(0);
+        let out_ports = out_port_counts.get(&v).copied().unwrap_or(0);
+        let new_v = g.add_node(in_ports, out_ports);
+        map_vertices.insert(v, new_v);
+    }
+    for &Edge {
+        src_v,
+        target_v,
+        src_p,
+        target_p,
+    } in edges
+    {
+        let src_v = map_vertices[&src_v];
+        let target_v = map_vertices[&target_v];
+        g.link_offsets(src_v, src_p, target_v, target_p).unwrap();
+    }
+    g
+}
+
+fn ensure_non_empty<G: PortView + PortMut>(g: &mut G) {
+    if g.node_count() == 0 {
+        g.add_node(0, 0);
+    }
+}
+
 #[cfg(test)]
 mod tests {
     use super::*;
@@ -111,4 +172,14 @@ mod tests {
             prop_assert!(graph.port_count() <= 4 * 2 * graph.node_count());
         }
     }
+
+    proptest! {
+        #[test]
+        fn gen_basic_multigraphs(graph in gen_multiportgraph(10, 4, 20)) {
+            prop_assert!(graph.node_count() <= 10);
+            prop_assert!(graph.node_count() >= 1);
+            prop_assert!(graph.link_count() <= 20);
+            prop_assert!(graph.port_count() <= 4 * 2 * graph.node_count());
+        }
+    }
 }