Optimize find_shortest_chains

CHANGELOG.rst

@@ -7,6 +7,7 @@ Unreleased
 
 * Add closed layers to layer contract.
 * Rename default repository branch to 'main'.
+* Optimise `find_shortest_chains` query.
 
 3.9 (2025-05-05)
 ----------------

rust/src/graph/higher_order_queries.rs

@@ -2,9 +2,8 @@ use crate::errors::GrimpResult;
 use crate::graph::{ExtendWithDescendants, Graph, ModuleToken};
 use derive_new::new;
 use getset::Getters;
-use itertools::Itertools;
 use rayon::prelude::*;
-use rustc_hash::{FxHashMap, FxHashSet};
+use rustc_hash::FxHashSet;
 
 use tap::prelude::*;
 
@@ -146,35 +145,18 @@ impl Graph {
         let excluded_modules =
             all_layers_modules - &(&from_layer_with_descendants | &to_layer_with_descendants);
 
-        // Disallow chains via these imports.
-        // We'll add chains to this set as we discover them.
-        let mut excluded_imports = FxHashMap::default();
+        let chains = self._find_shortest_chains(
+            &from_layer_with_descendants,
+            &to_layer_with_descendants,
+            &excluded_modules,
+        )?;
 
         // Collect direct imports...
         let mut direct_imports = vec![];
         // ...and the middles of any indirect imports.
         let mut middles = vec![];
-        loop {
-            let chain = self.find_shortest_chain_with_excluded_modules_and_imports(
-                &from_layer_with_descendants,
-                &to_layer_with_descendants,
-                &excluded_modules,
-                &excluded_imports,
-            )?;
-
-            if chain.is_none() {
-                break;
-            }
-            let chain = chain.unwrap();
-
-            // Exclude this chain from further searching.
-            for (importer, imported) in chain.iter().tuple_windows() {
-                excluded_imports
-                    .entry(*importer)
-                    .or_default()
-                    .insert(*imported);
-            }
 
+        for chain in chains {
             let (head, middle, tail) = self.split_chain(&chain);
             match middle {
                 Some(middle) => middles.push(middle),

rust/src/graph/import_chain_queries.rs

@@ -100,24 +100,55 @@ impl Graph {
             downstream_modules.extend_with_descendants(self);
             upstream_modules.extend_with_descendants(self);
         }
-        let all_modules = &downstream_modules | &upstream_modules;
-
-        let chains = downstream_modules
-            .iter()
-            .cartesian_product(upstream_modules.iter())
-            .filter_map(|(downstream_module, upstream_module)| {
-                let excluded_modules =
-                    &all_modules - &FxHashSet::from_iter([*downstream_module, *upstream_module]);
-                self.find_shortest_chain_with_excluded_modules_and_imports(
-                    &(*downstream_module).into(),
-                    &(*upstream_module).into(),
-                    &excluded_modules,
-                    &FxHashMap::default(),
-                )
-                .unwrap()
-            })
+
+        let chains = self
+            ._find_shortest_chains(
+                &downstream_modules,
+                &upstream_modules,
+                &FxHashSet::from_iter([]),
+            )?
+            .into_iter()
             .collect();
 
         Ok(chains)
     }
+
+    pub(crate) fn _find_shortest_chains(
+        &self,
+        from_modules: &FxHashSet<ModuleToken>,
+        to_modules: &FxHashSet<ModuleToken>,
+        excluded_modules: &FxHashSet<ModuleToken>,
+    ) -> GrimpResult<Vec<Vec<ModuleToken>>> {
+        let mut chains = vec![];
+
+        // Disallow chains via these imports.
+        // We'll add chains to this set as we discover them.
+        let mut excluded_imports = FxHashMap::default();
+
+        loop {
+            let chain = self.find_shortest_chain_with_excluded_modules_and_imports(
+                from_modules,
+                to_modules,
+                excluded_modules,
+                &excluded_imports,
+            )?;
+
+            if chain.is_none() {
+                break;
+            }
+            let chain = chain.unwrap();
+
+            // Exclude this chain from further searching.
+            for (importer, imported) in chain.iter().tuple_windows() {
+                excluded_imports
+                    .entry(*importer)
+                    .or_default()
+                    .insert(*imported);
+            }
+
+            chains.push(chain);
+        }
+
+        Ok(chains)
+    }
 }