rustworkx_core/
token_swapper.rs

1// Licensed under the Apache License, Version 2.0 (the "License"); you may
2// not use this file except in compliance with the License. You may obtain
3// a copy of the License at
4//
5//     http://www.apache.org/licenses/LICENSE-2.0
6//
7// Unless required by applicable law or agreed to in writing, software
8// distributed under the License is distributed on an "AS IS" BASIS, WITHOUT
9// WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the
10// License for the specific language governing permissions and limitations
11// under the License.
12
13use rand::distributions::{Standard, Uniform};
14use rand::prelude::*;
15use rand_pcg::Pcg64;
16use std::error::Error;
17use std::fmt;
18use std::hash::Hash;
19
20use hashbrown::HashMap;
21use petgraph::stable_graph::{NodeIndex, StableGraph};
22use petgraph::visit::{
23    EdgeCount, GraphBase, IntoEdges, IntoNeighborsDirected, IntoNodeIdentifiers, NodeCount,
24    NodeIndexable, Visitable,
25};
26use petgraph::Directed;
27use petgraph::Direction::{Incoming, Outgoing};
28use rayon::prelude::*;
29use rayon_cond::CondIterator;
30
31use crate::connectivity::find_cycle;
32use crate::dictmap::*;
33use crate::shortest_path::dijkstra;
34use crate::traversal::dfs_edges;
35
36type Swap = (NodeIndex, NodeIndex);
37type Edge = (NodeIndex, NodeIndex);
38
39/// Error returned by token swapper if the request mapping
40/// is impossible
41#[derive(Debug, PartialEq, Eq, Ord, PartialOrd, Copy, Clone)]
42pub struct MapNotPossible;
43
44impl Error for MapNotPossible {}
45impl fmt::Display for MapNotPossible {
46    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
47        write!(f, "No mapping possible.")
48    }
49}
50
51struct TokenSwapper<G: GraphBase>
52where
53    G::NodeId: Eq + Hash,
54{
55    // The input graph
56    graph: G,
57    // The user-supplied mapping to use for swapping tokens
58    mapping: HashMap<G::NodeId, G::NodeId>,
59    // Number of trials
60    trials: usize,
61    // Seed for random selection of a node for a trial
62    seed: Option<u64>,
63    // Threshold for how many nodes will trigger parallel iterator
64    parallel_threshold: usize,
65    // Map of NodeId to NodeIndex
66    node_map: HashMap<G::NodeId, NodeIndex>,
67    // Map of NodeIndex to NodeId
68    rev_node_map: HashMap<NodeIndex, G::NodeId>,
69}
70
71impl<G> TokenSwapper<G>
72where
73    G: NodeCount
74        + EdgeCount
75        + IntoEdges
76        + Visitable
77        + NodeIndexable
78        + IntoNeighborsDirected
79        + IntoNodeIdentifiers
80        + Send
81        + Sync,
82    G::NodeId: Hash + Eq + Send + Sync,
83{
84    fn new(
85        graph: G,
86        mapping: HashMap<G::NodeId, G::NodeId>,
87        trials: Option<usize>,
88        seed: Option<u64>,
89        parallel_threshold: Option<usize>,
90    ) -> Self {
91        TokenSwapper {
92            graph,
93            mapping,
94            trials: trials.unwrap_or(4),
95            seed,
96            parallel_threshold: parallel_threshold.unwrap_or(50),
97            node_map: HashMap::with_capacity(graph.node_count()),
98            rev_node_map: HashMap::with_capacity(graph.node_count()),
99        }
100    }
101
102    fn map(&mut self) -> Result<Vec<Swap>, MapNotPossible> {
103        let num_nodes = self.graph.node_bound();
104        let num_edges = self.graph.edge_count();
105
106        // Directed graph with nodes matching ``graph`` and
107        // edges for neighbors closer than nodes
108        let mut digraph = StableGraph::with_capacity(num_nodes, num_edges);
109
110        // First fill the digraph with nodes. Then since it's a stable graph,
111        // must go through and remove nodes that were removed in original graph
112        for _ in 0..self.graph.node_bound() {
113            digraph.add_node(());
114        }
115        let mut count: usize = 0;
116        for gnode in self.graph.node_identifiers() {
117            let gidx = self.graph.to_index(gnode);
118            if gidx != count {
119                for idx in count..gidx {
120                    digraph.remove_node(NodeIndex::new(idx));
121                }
122                count = gidx;
123            }
124            count += 1;
125        }
126
127        // Create maps between NodeId and NodeIndex
128        for node in self.graph.node_identifiers() {
129            self.node_map
130                .insert(node, NodeIndex::new(self.graph.to_index(node)));
131            self.rev_node_map
132                .insert(NodeIndex::new(self.graph.to_index(node)), node);
133        }
134        // sub will become same as digraph but with no self edges in add_token_edges
135        let mut sub_digraph = digraph.clone();
136
137        // The mapping in HashMap form using NodeIndex
138        let mut tokens: HashMap<NodeIndex, NodeIndex> = self
139            .mapping
140            .iter()
141            .map(|(k, v)| (self.node_map[k], self.node_map[v]))
142            .collect();
143
144        // todo_nodes are all the mapping entries where left != right
145        let mut todo_nodes: Vec<NodeIndex> = tokens
146            .iter()
147            .filter_map(|(node, dest)| if node != dest { Some(*node) } else { None })
148            .collect();
149        todo_nodes.par_sort();
150
151        // Add initial edges to the digraph/sub_digraph
152        for node in self.graph.node_identifiers() {
153            self.add_token_edges(
154                self.node_map[&node],
155                &mut digraph,
156                &mut sub_digraph,
157                &mut tokens,
158            )?;
159        }
160        // First collect the self.trial number of random numbers
161        // into a Vec based on the given seed
162        let outer_rng: Pcg64 = match self.seed {
163            Some(rng_seed) => Pcg64::seed_from_u64(rng_seed),
164            None => Pcg64::from_entropy(),
165        };
166        let trial_seeds_vec: Vec<u64> =
167            outer_rng.sample_iter(&Standard).take(self.trials).collect();
168
169        CondIterator::new(
170            trial_seeds_vec,
171            self.graph.node_count() >= self.parallel_threshold,
172        )
173        .map(|trial_seed| {
174            self.trial_map(
175                digraph.clone(),
176                sub_digraph.clone(),
177                tokens.clone(),
178                todo_nodes.clone(),
179                trial_seed,
180            )
181        })
182        .min_by_key(|result| match result {
183            Ok(res) => Ok(res.len()),
184            Err(e) => Err(*e),
185        })
186        .unwrap()
187    }
188
189    fn add_token_edges(
190        &self,
191        node: NodeIndex,
192        digraph: &mut StableGraph<(), (), Directed>,
193        sub_digraph: &mut StableGraph<(), (), Directed>,
194        tokens: &mut HashMap<NodeIndex, NodeIndex>,
195    ) -> Result<(), MapNotPossible> {
196        // Adds an edge to digraph if distance from the token to a neighbor is
197        // less than distance from token to node. sub_digraph is same except
198        // for self-edges.
199        if !tokens.contains_key(&node) {
200            return Ok(());
201        }
202        if tokens[&node] == node {
203            digraph.update_edge(node, node, ());
204            return Ok(());
205        }
206        let id_node = self.rev_node_map[&node];
207        let id_token = self.rev_node_map[&tokens[&node]];
208
209        if self.graph.neighbors(id_node).next().is_none() {
210            return Err(MapNotPossible {});
211        }
212
213        for id_neighbor in self.graph.neighbors(id_node) {
214            let neighbor = self.node_map[&id_neighbor];
215            let dist_neighbor: DictMap<G::NodeId, usize> = dijkstra(
216                &self.graph,
217                id_neighbor,
218                Some(id_token),
219                |_| Ok::<usize, &str>(1),
220                None,
221            )
222            .unwrap();
223
224            let dist_node: DictMap<G::NodeId, usize> = dijkstra(
225                &self.graph,
226                id_node,
227                Some(id_token),
228                |_| Ok::<usize, &str>(1),
229                None,
230            )
231            .unwrap();
232            let neigh_dist = dist_neighbor.get(&id_token);
233            let node_dist = dist_node.get(&id_token);
234            if neigh_dist.is_none() {
235                return Err(MapNotPossible {});
236            }
237            if node_dist.is_none() {
238                return Err(MapNotPossible {});
239            }
240
241            if neigh_dist < node_dist {
242                digraph.update_edge(node, neighbor, ());
243                sub_digraph.update_edge(node, neighbor, ());
244            }
245        }
246        Ok(())
247    }
248
249    fn trial_map(
250        &self,
251        mut digraph: StableGraph<(), (), Directed>,
252        mut sub_digraph: StableGraph<(), (), Directed>,
253        mut tokens: HashMap<NodeIndex, NodeIndex>,
254        mut todo_nodes: Vec<NodeIndex>,
255        trial_seed: u64,
256    ) -> Result<Vec<Swap>, MapNotPossible> {
257        // Create a random trial list of swaps to move tokens to optimal positions
258        let mut steps = 0;
259        let mut swap_edges: Vec<Swap> = vec![];
260        let mut rng_seed: Pcg64 = Pcg64::seed_from_u64(trial_seed);
261        while !todo_nodes.is_empty() && steps <= 4 * digraph.node_count().pow(2) {
262            // Choose a random todo_node
263            let between = Uniform::new(0, todo_nodes.len());
264            let random: usize = between.sample(&mut rng_seed);
265            let todo_node = todo_nodes[random];
266
267            // If there's a cycle in sub_digraph, add it to swap_edges and do swap
268            let cycle = find_cycle(&sub_digraph, Some(todo_node));
269            if !cycle.is_empty() {
270                for edge in cycle[1..].iter().rev() {
271                    swap_edges.push(*edge);
272                    self.swap(
273                        edge.0,
274                        edge.1,
275                        &mut digraph,
276                        &mut sub_digraph,
277                        &mut tokens,
278                        &mut todo_nodes,
279                    )?;
280                }
281                steps += cycle.len() - 1;
282            // If there's no cycle, see if there's an edge target that matches a token key.
283            // If so, add to swap_edges and do swap
284            } else {
285                let mut found = false;
286                let sub2 = &sub_digraph.clone();
287                for edge in dfs_edges(sub2, Some(todo_node)) {
288                    let new_edge = (NodeIndex::new(edge.0), NodeIndex::new(edge.1));
289                    if !tokens.contains_key(&new_edge.1) {
290                        swap_edges.push(new_edge);
291                        self.swap(
292                            new_edge.0,
293                            new_edge.1,
294                            &mut digraph,
295                            &mut sub_digraph,
296                            &mut tokens,
297                            &mut todo_nodes,
298                        )?;
299                        steps += 1;
300                        found = true;
301                        break;
302                    }
303                }
304                // If none found, look for cycle in digraph which will result in
305                // an unhappy node. Look for a predecessor and add node and pred
306                // to swap_edges and do swap
307                if !found {
308                    let cycle: Vec<Edge> = find_cycle(&digraph, Some(todo_node));
309                    let unhappy_node = cycle[0].0;
310                    let mut found = false;
311                    let di2 = &mut digraph.clone();
312                    for predecessor in di2.neighbors_directed(unhappy_node, Incoming) {
313                        if predecessor != unhappy_node {
314                            swap_edges.push((unhappy_node, predecessor));
315                            self.swap(
316                                unhappy_node,
317                                predecessor,
318                                &mut digraph,
319                                &mut sub_digraph,
320                                &mut tokens,
321                                &mut todo_nodes,
322                            )?;
323                            steps += 1;
324                            found = true;
325                            break;
326                        }
327                    }
328                    assert!(
329                        found,
330                        "The token swap process has ended unexpectedly, this points to a bug in rustworkx, please open an issue."
331                    );
332                }
333            }
334        }
335        assert!(
336            todo_nodes.is_empty(),
337            "The output final swap map is incomplete, this points to a bug in rustworkx, please open an issue."
338        );
339        Ok(swap_edges)
340    }
341
342    fn swap(
343        &self,
344        node1: NodeIndex,
345        node2: NodeIndex,
346        digraph: &mut StableGraph<(), (), Directed>,
347        sub_digraph: &mut StableGraph<(), (), Directed>,
348        tokens: &mut HashMap<NodeIndex, NodeIndex>,
349        todo_nodes: &mut Vec<NodeIndex>,
350    ) -> Result<(), MapNotPossible> {
351        // Get token values for the 2 nodes and remove them
352        let token1 = tokens.remove(&node1);
353        let token2 = tokens.remove(&node2);
354
355        // Swap the token edge values
356        if let Some(t2) = token2 {
357            tokens.insert(node1, t2);
358        }
359        if let Some(t1) = token1 {
360            tokens.insert(node2, t1);
361        }
362        // For each node, remove the (node, successor) from digraph and
363        // sub_digraph. Then add new token edges back in.
364        for node in [node1, node2] {
365            let edge_nodes: Vec<(NodeIndex, NodeIndex)> = digraph
366                .neighbors_directed(node, Outgoing)
367                .map(|successor| (node, successor))
368                .collect();
369            for (edge_node1, edge_node2) in edge_nodes {
370                let edge = digraph.find_edge(edge_node1, edge_node2).unwrap();
371                digraph.remove_edge(edge);
372            }
373            let edge_nodes: Vec<(NodeIndex, NodeIndex)> = sub_digraph
374                .neighbors_directed(node, Outgoing)
375                .map(|successor| (node, successor))
376                .collect();
377            for (edge_node1, edge_node2) in edge_nodes {
378                let edge = sub_digraph.find_edge(edge_node1, edge_node2).unwrap();
379                sub_digraph.remove_edge(edge);
380            }
381            self.add_token_edges(node, digraph, sub_digraph, tokens)?;
382
383            // If a node is a token key and not equal to the value, add it to todo_nodes
384            if tokens.contains_key(&node) && tokens[&node] != node {
385                if !todo_nodes.contains(&node) {
386                    todo_nodes.push(node);
387                }
388            // Otherwise if node is in todo_nodes, remove it
389            } else if todo_nodes.contains(&node) {
390                todo_nodes.swap_remove(todo_nodes.iter().position(|x| *x == node).unwrap());
391            }
392        }
393        Ok(())
394    }
395}
396
397/// Module to perform an approximately optimal Token Swapping algorithm. Supports partial
398/// mappings (i.e. not-permutations) for graphs with missing tokens.
399///
400/// Based on the paper: Approximation and Hardness for Token Swapping by Miltzow et al. (2016)
401/// ArXiV: <https://arxiv.org/abs/1602.05150>
402///
403/// Arguments:
404///
405/// * `graph` - The graph on which to perform the token swapping.
406/// * `mapping` - A partial mapping to be implemented in swaps.
407/// * `trials` - Optional number of trials. If None, defaults to 4.
408/// * `seed` - Optional integer seed. If None, the internal rng will be initialized from system entropy.
409/// * `parallel_threshold` - Optional integer for the number of nodes in the graph that will
410///     trigger the use of parallel threads. If the number of nodes in the graph is less than this value
411///     it will run in a single thread. The default value is 50.
412///
413/// It returns a list of tuples representing the swaps to perform. The result will be an
414/// `Err(MapNotPossible)` if the `token_swapper()` function can't find a mapping.
415///
416/// This function is multithreaded and will launch a thread pool with threads equal to
417/// the number of CPUs by default. You can tune the number of threads with
418/// the ``RAYON_NUM_THREADS`` environment variable. For example, setting ``RAYON_NUM_THREADS=4``
419/// would limit the thread pool to 4 threads.
420///
421/// # Example
422/// ```rust
423/// use hashbrown::HashMap;
424/// use rustworkx_core::petgraph;
425/// use rustworkx_core::token_swapper::token_swapper;
426/// use rustworkx_core::petgraph::graph::NodeIndex;
427///
428///  let g = petgraph::graph::UnGraph::<(), ()>::from_edges(&[(0, 1), (1, 2), (2, 3)]);
429///  let mapping = HashMap::from([
430///   (NodeIndex::new(0), NodeIndex::new(0)),
431///   (NodeIndex::new(1), NodeIndex::new(3)),
432///   (NodeIndex::new(3), NodeIndex::new(1)),
433///   (NodeIndex::new(2), NodeIndex::new(2)),
434///  ]);
435///  // Do the token swap
436///  let output = token_swapper(&g, mapping, Some(4), Some(4), Some(50)).expect("Swap mapping failed.");
437///  assert_eq!(3, output.len());
438///
439/// ```
440pub fn token_swapper<G>(
441    graph: G,
442    mapping: HashMap<G::NodeId, G::NodeId>,
443    trials: Option<usize>,
444    seed: Option<u64>,
445    parallel_threshold: Option<usize>,
446) -> Result<Vec<Swap>, MapNotPossible>
447where
448    G: NodeCount
449        + EdgeCount
450        + IntoEdges
451        + Visitable
452        + NodeIndexable
453        + IntoNeighborsDirected
454        + IntoNodeIdentifiers
455        + Send
456        + Sync,
457    G::NodeId: Hash + Eq + Send + Sync,
458{
459    let mut swapper = TokenSwapper::new(graph, mapping, trials, seed, parallel_threshold);
460    swapper.map()
461}
462
463#[cfg(test)]
464mod test_token_swapper {
465
466    use crate::petgraph;
467    use crate::token_swapper::token_swapper;
468    use hashbrown::HashMap;
469    use petgraph::graph::NodeIndex;
470
471    fn do_swap(mapping: &mut HashMap<NodeIndex, NodeIndex>, swaps: &Vec<(NodeIndex, NodeIndex)>) {
472        // Apply the swaps to the mapping to get final result
473        for (swap1, swap2) in swaps {
474            //Need to create temp nodes in case of partial mapping
475            let mut temp_node1: Option<NodeIndex> = None;
476            let mut temp_node2: Option<NodeIndex> = None;
477            if mapping.contains_key(swap1) {
478                temp_node1 = Some(mapping[swap1]);
479                mapping.remove(swap1);
480            }
481            if mapping.contains_key(swap2) {
482                temp_node2 = Some(mapping[swap2]);
483                mapping.remove(swap2);
484            }
485            if let Some(t1) = temp_node1 {
486                mapping.insert(*swap2, t1);
487            }
488            if let Some(t2) = temp_node2 {
489                mapping.insert(*swap1, t2);
490            }
491        }
492    }
493
494    #[test]
495    fn test_simple_swap() {
496        // Simple arbitrary swap
497        let g = petgraph::graph::UnGraph::<(), ()>::from_edges([(0, 1), (1, 2), (2, 3)]);
498        let mapping = HashMap::from([
499            (NodeIndex::new(0), NodeIndex::new(0)),
500            (NodeIndex::new(1), NodeIndex::new(3)),
501            (NodeIndex::new(3), NodeIndex::new(1)),
502            (NodeIndex::new(2), NodeIndex::new(2)),
503        ]);
504        let swaps = token_swapper(&g, mapping, Some(4), Some(4), Some(50));
505        assert_eq!(3, swaps.expect("swap mapping errored").len());
506    }
507
508    #[test]
509    fn test_small_swap() {
510        // Reverse all small swap
511        let g = petgraph::graph::UnGraph::<(), ()>::from_edges([
512            (0, 1),
513            (1, 2),
514            (2, 3),
515            (3, 4),
516            (4, 5),
517            (5, 6),
518            (6, 7),
519        ]);
520        let mut mapping = HashMap::with_capacity(8);
521        for i in 0..8 {
522            mapping.insert(NodeIndex::new(i), NodeIndex::new(7 - i));
523        }
524        // Do the token swap
525        let mut new_map = mapping.clone();
526        let swaps =
527            token_swapper(&g, mapping, Some(4), Some(4), Some(50)).expect("swap mapping errored");
528        do_swap(&mut new_map, &swaps);
529        let mut expected = HashMap::with_capacity(8);
530        for i in 0..8 {
531            expected.insert(NodeIndex::new(i), NodeIndex::new(i));
532        }
533        assert_eq!(expected, new_map);
534    }
535
536    #[test]
537    fn test_happy_swap_chain() {
538        // Reverse all happy swap chain > 2
539        let g = petgraph::graph::UnGraph::<(), ()>::from_edges([
540            (0, 1),
541            (0, 2),
542            (0, 3),
543            (0, 4),
544            (1, 2),
545            (1, 3),
546            (1, 4),
547            (2, 3),
548            (2, 4),
549            (3, 4),
550            (3, 6),
551        ]);
552        let mapping = HashMap::from([
553            (NodeIndex::new(0), NodeIndex::new(4)),
554            (NodeIndex::new(1), NodeIndex::new(0)),
555            (NodeIndex::new(2), NodeIndex::new(3)),
556            (NodeIndex::new(3), NodeIndex::new(6)),
557            (NodeIndex::new(4), NodeIndex::new(2)),
558            (NodeIndex::new(6), NodeIndex::new(1)),
559        ]);
560        // Do the token swap
561        let mut new_map = mapping.clone();
562        let swaps =
563            token_swapper(&g, mapping, Some(4), Some(4), Some(50)).expect("swap mapping errored");
564        do_swap(&mut new_map, &swaps);
565        let mut expected = HashMap::with_capacity(6);
566        for i in (0..5).chain(6..7) {
567            expected.insert(NodeIndex::new(i), NodeIndex::new(i));
568        }
569        assert_eq!(expected, new_map);
570    }
571
572    #[test]
573    fn test_partial_simple() {
574        // Simple partial swap
575        let g = petgraph::graph::UnGraph::<(), ()>::from_edges([(0, 1), (1, 2), (2, 3)]);
576        let mapping = HashMap::from([(NodeIndex::new(0), NodeIndex::new(3))]);
577        let mut new_map = mapping.clone();
578        let swaps =
579            token_swapper(&g, mapping, Some(4), Some(4), Some(1)).expect("swap mapping errored");
580        do_swap(&mut new_map, &swaps);
581        let mut expected = HashMap::with_capacity(4);
582        expected.insert(NodeIndex::new(3), NodeIndex::new(3));
583        assert_eq!(expected, new_map);
584    }
585
586    #[test]
587    fn test_partial_simple_remove_node() {
588        // Simple partial swap
589        let mut g =
590            petgraph::graph::UnGraph::<(), ()>::from_edges([(0, 1), (1, 2), (2, 3), (3, 4)]);
591        let mapping = HashMap::from([(NodeIndex::new(0), NodeIndex::new(3))]);
592        g.remove_node(NodeIndex::new(2));
593        g.add_edge(NodeIndex::new(1), NodeIndex::new(3), ());
594        let mut new_map = mapping.clone();
595        let swaps =
596            token_swapper(&g, mapping, Some(4), Some(4), Some(1)).expect("swap mapping errored");
597        do_swap(&mut new_map, &swaps);
598        let mut expected = HashMap::with_capacity(4);
599        expected.insert(NodeIndex::new(3), NodeIndex::new(3));
600        assert_eq!(expected, new_map);
601    }
602
603    #[test]
604    fn test_partial_small() {
605        // Partial inverting on small path graph
606        let g = petgraph::graph::UnGraph::<(), ()>::from_edges([(0, 1), (1, 2), (2, 3)]);
607        let mapping = HashMap::from([
608            (NodeIndex::new(0), NodeIndex::new(3)),
609            (NodeIndex::new(1), NodeIndex::new(2)),
610        ]);
611        let mut new_map = mapping.clone();
612        let swaps =
613            token_swapper(&g, mapping, Some(4), Some(4), Some(50)).expect("swap mapping errored");
614        do_swap(&mut new_map, &swaps);
615        let expected = HashMap::from([
616            (NodeIndex::new(2), NodeIndex::new(2)),
617            (NodeIndex::new(3), NodeIndex::new(3)),
618        ]);
619        assert_eq!(5, swaps.len());
620        assert_eq!(expected, new_map);
621    }
622
623    #[test]
624    fn test_large_partial_random() {
625        // Test a random (partial) mapping on a large randomly generated graph
626        use crate::generators::gnm_random_graph;
627        use rand::prelude::*;
628        use rand_pcg::Pcg64;
629        use std::iter::zip;
630
631        let mut rng: Pcg64 = Pcg64::seed_from_u64(4);
632
633        // Note that graph may have "gaps" in the node counts, i.e. the numbering is noncontiguous.
634        let size = 100;
635        let mut g: petgraph::stable_graph::StableGraph<(), ()> =
636            gnm_random_graph(size, size.pow(2) / 10, Some(4), || (), || ()).unwrap();
637
638        // Remove self-loops
639        let nodes: Vec<_> = g.node_indices().collect();
640        for node in nodes {
641            if let Some(edge) = g.find_edge(node, node) {
642                g.remove_edge(edge);
643            }
644        }
645        // Make sure the graph is connected by adding C_n
646        for i in 0..(g.node_count() - 1) {
647            g.add_edge(NodeIndex::new(i), NodeIndex::new(i + 1), ());
648        }
649
650        // Get node indices and randomly shuffle
651        let mut mapped_nodes: Vec<usize> = g.node_indices().map(|node| node.index()).collect();
652        let nodes = mapped_nodes.clone();
653        mapped_nodes.shuffle(&mut rng);
654
655        // Zip nodes and shuffled nodes and remove every other one
656        let mut mapping: Vec<(usize, usize)> = zip(nodes, mapped_nodes).collect();
657        mapping.retain(|(a, _)| a % 2 == 0);
658
659        // Convert mapping to HashMap of NodeIndex's
660        let mapping: HashMap<NodeIndex, NodeIndex> = mapping
661            .into_iter()
662            .map(|(a, b)| (NodeIndex::new(a), NodeIndex::new(b)))
663            .collect();
664        let mut new_map = mapping.clone();
665        let expected: HashMap<NodeIndex, NodeIndex> =
666            mapping.values().map(|val| (*val, *val)).collect();
667
668        let swaps =
669            token_swapper(&g, mapping, Some(4), Some(4), Some(50)).expect("swap mapping errored");
670        do_swap(&mut new_map, &swaps);
671        assert_eq!(expected, new_map)
672    }
673
674    #[test]
675    fn test_disjoint_graph_works() {
676        let g = petgraph::graph::UnGraph::<(), ()>::from_edges([(0, 1), (2, 3)]);
677        let mapping = HashMap::from([
678            (NodeIndex::new(1), NodeIndex::new(0)),
679            (NodeIndex::new(0), NodeIndex::new(1)),
680            (NodeIndex::new(2), NodeIndex::new(3)),
681            (NodeIndex::new(3), NodeIndex::new(2)),
682        ]);
683        let mut new_map = mapping.clone();
684        let swaps =
685            token_swapper(&g, mapping, Some(10), Some(4), Some(50)).expect("swap mapping errored");
686        do_swap(&mut new_map, &swaps);
687        let expected = HashMap::from([
688            (NodeIndex::new(2), NodeIndex::new(2)),
689            (NodeIndex::new(3), NodeIndex::new(3)),
690            (NodeIndex::new(1), NodeIndex::new(1)),
691            (NodeIndex::new(0), NodeIndex::new(0)),
692        ]);
693        assert_eq!(2, swaps.len());
694        assert_eq!(expected, new_map);
695    }
696
697    #[test]
698    fn test_disjoint_graph_fails() {
699        let g = petgraph::graph::UnGraph::<(), ()>::from_edges([(0, 1), (2, 3)]);
700        let mapping = HashMap::from([
701            (NodeIndex::new(2), NodeIndex::new(0)),
702            (NodeIndex::new(1), NodeIndex::new(1)),
703            (NodeIndex::new(0), NodeIndex::new(2)),
704            (NodeIndex::new(3), NodeIndex::new(3)),
705        ]);
706        assert!(
707            token_swapper(&g, mapping, Some(10), Some(4), Some(50)).is_err(),
708            "This should error"
709        );
710    }
711
712    #[test]
713    fn test_edgeless_graph_fails() {
714        let mut g = petgraph::graph::UnGraph::<(), ()>::new_undirected();
715        let a = g.add_node(());
716        let b = g.add_node(());
717        let c = g.add_node(());
718        let d = g.add_node(());
719        g.add_edge(c, d, ());
720        let mapping = HashMap::from([(a, b), (b, a)]);
721        assert!(
722            token_swapper(&g, mapping, Some(10), Some(4), Some(50)).is_err(),
723            "This should error"
724        );
725    }
726}
rustworkx_core/token_swapper.rs

rustworkx_core/
token_swapper.rs
