Compute edge orbits while automorphisms are generated, instead of afterward

src/undirected_graph.rs

 use std::collections::{HashMap, HashSet};
+use union_find::{QuickUnionUf, UnionByRank, UnionFind};
 
 #[derive(Debug, Eq, Hash, PartialEq, Clone)]
 pub struct UndirectedGraph {
@@ -6,7 +7,30 @@ pub struct UndirectedGraph {
     edges: Vec<(usize, usize)>,
 }
 
+struct EdgeOrbitData {
+    graph: UndirectedGraph,
+    graph_is_odd: bool,
+    compute_edge_orbits: bool,
+    // Union-Find for edge orbits.
+    edge_orbits_uf: QuickUnionUf<UnionByRank>,
+}
+
+impl EdgeOrbitData {
+    pub fn new() -> Self {
+        EdgeOrbitData {
+            graph: UndirectedGraph::new(0, vec![]),
+            graph_is_odd: false,
+            compute_edge_orbits: false,
+            edge_orbits_uf: QuickUnionUf::<UnionByRank>::new(0),
+        }
+    }
+}
+
 impl UndirectedGraph {
+    thread_local! {
+        static EDGE_ORBIT_DATA: std::cell::RefCell<EdgeOrbitData> = std::cell::RefCell::new(EdgeOrbitData::new());
+    }
+
     pub fn new(num_vertices: usize, edges: Vec<(usize, usize)>) -> Self {
         let mut edges_seen: HashSet<(usize, usize)> = HashSet::new();
         for (a, b) in &edges {
@@ -78,7 +102,11 @@ impl UndirectedGraph {
         fn map_edge(relabel_h_to_g: &[i32], e: &(usize, usize)) -> (usize, usize) {
             let a: usize = relabel_h_to_g[e.0] as usize;
             let b: usize = relabel_h_to_g[e.1] as usize;
-            if a < b { (a, b) } else { (b, a) }
+            if a < b {
+                (a, b)
+            } else {
+                (b, a)
+            }
         }
         let mapped_edges: Vec<(usize, usize)> = h
             .edges
@@ -108,12 +136,8 @@ impl UndirectedGraph {
         sign
     }
 
-    thread_local! {
-        static AUTOMORPHSM_GROUP_GENERATORS: std::cell::RefCell<Vec<Vec<i32>>> = const { std::cell::RefCell::new(vec![]) };
-    }
-
     #[no_mangle]
-    extern "C" fn store_nauty_automorphism(
+    extern "C" fn process_nauty_automorphism(
         _count: i32,
         perm: *mut i32,
         _orbits: *mut i32,
@@ -123,7 +147,37 @@ impl UndirectedGraph {
     ) {
         unsafe {
             let automorphism: Vec<i32> = std::slice::from_raw_parts(perm, n as usize).to_vec();
-            Self::AUTOMORPHSM_GROUP_GENERATORS.with_borrow_mut(|gens| gens.push(automorphism));
+
+            Self::EDGE_ORBIT_DATA.with_borrow_mut(|edge_orbit_data: &mut EdgeOrbitData| {
+                // NOTE: The automorphisms produced by nauty act on the original graph.
+                let sigma: Vec<i32> = Self::induced_edge_permutation(
+                    &edge_orbit_data.graph,
+                    &edge_orbit_data.graph,
+                    &automorphism,
+                );
+
+                if edge_orbit_data.compute_edge_orbits {
+                    // Take powers of sigma and act on edges of graph.
+                    let mut tau: Vec<i32> = sigma.clone();
+                    let num_edges: i32 = edge_orbit_data.graph.edges.len() as i32;
+                    let identity_permutation: Vec<i32> = (0..num_edges).collect();
+                    while tau != identity_permutation {
+                        for (e, &tau_e) in tau.iter().enumerate() {
+                            edge_orbit_data.edge_orbits_uf.union(e, tau_e as usize);
+                        }
+                        for e in tau.iter_mut() {
+                            *e = sigma[*e as usize];
+                        }
+                    }
+                }
+
+                if Self::permutation_sign(sigma) == -1 {
+                    edge_orbit_data.graph_is_odd = true;
+                    // NOTE: Giving up computation of edge orbits; we will not use them anyway.
+                    edge_orbit_data.compute_edge_orbits = false;
+                    // NOTE: There is no way to make nauty cancel the computation here.
+                }
+            });
         }
     }
 
@@ -135,7 +189,7 @@ impl UndirectedGraph {
 
         let mut options: nauty_Traces_sys::optionblk = nauty_Traces_sys::optionblk {
             getcanon: nauty_Traces_sys::TRUE,
-            userautomproc: Some(Self::store_nauty_automorphism),
+            userautomproc: Some(Self::process_nauty_automorphism),
             ..Default::default()
         };
         let mut stats: nauty_Traces_sys::statsblk = nauty_Traces_sys::statsblk::default();
@@ -149,7 +203,12 @@ impl UndirectedGraph {
 
         let mut normal_graph: Vec<u64> = vec![0; n];
 
-        Self::AUTOMORPHSM_GROUP_GENERATORS.set(vec![]);
+        Self::EDGE_ORBIT_DATA.set(EdgeOrbitData {
+            graph: self.clone(),
+            graph_is_odd: false,
+            compute_edge_orbits,
+            edge_orbits_uf: QuickUnionUf::<UnionByRank>::new(self.edges.len()),
+        });
 
         unsafe {
             nauty_Traces_sys::densenauty(
@@ -186,43 +245,17 @@ impl UndirectedGraph {
 
         let mut sign: i32 = 1;
         let mut normal_self_edge_orbits: Vec<usize> = vec![];
-        Self::AUTOMORPHSM_GROUP_GENERATORS.with_borrow(|gens| {
-            // NOTE: Use Union-Find to compute edge orbits.
-            use union_find::{QuickUnionUf, UnionByRank, UnionFind};
-            let mut uf: QuickUnionUf<UnionByRank> =
-                QuickUnionUf::<UnionByRank>::new(self.edges.len());
-
-            for automorphism in gens {
-                // NOTE: The automorphisms produced by nauty act on the original graph.
-                let sigma: Vec<i32> = Self::induced_edge_permutation(self, self, automorphism);
-
-                if compute_edge_orbits {
-                    // Take powers of sigma and act on edges of self.
-                    let mut tau: Vec<i32> = sigma.clone();
-                    let num_edges: i32 = self.edges.len() as i32;
-                    let identity_permutation: Vec<i32> = (0..num_edges).collect();
-                    while tau != identity_permutation {
-                        for (e, &tau_e) in tau.iter().enumerate() {
-                            uf.union(e, tau_e as usize);
-                        }
-                        for e in tau.iter_mut() {
-                            *e = sigma[*e as usize];
-                        }
-                    }
-                }
-
-                if Self::permutation_sign(sigma) == -1 {
-                    sign = 0;
-                    // NOTE: Giving up computation of edge orbits; we will not use them anyway.
-                    break;
-                }
+        Self::EDGE_ORBIT_DATA.with_borrow_mut(|edge_orbit_data: &mut EdgeOrbitData| {
+            if edge_orbit_data.graph_is_odd {
+                sign = 0;
             }
-
             if sign != 0 {
                 let edge_permutation: Vec<i32> =
                     Self::induced_edge_permutation(&normal_self, self, &lab_inv);
-                if compute_edge_orbits {
-                    normal_self_edge_orbits = (0..self.edges.len()).map(|e| uf.find(e)).collect();
+                if edge_orbit_data.compute_edge_orbits {
+                    normal_self_edge_orbits = (0..self.edges.len())
+                        .map(|e| edge_orbit_data.edge_orbits_uf.find(e))
+                        .collect();
                     // So far these are the edge orbits of self. Finally we apply the isomorphism to get the edge orbits of normal_self.
                     for e in &mut normal_self_edge_orbits {
                         *e = edge_permutation[*e] as usize;