From 061b1e6d7cb35b46a563b0f9b8505beed063b7f2 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Lo=C3=AFc=20Lecrenier?= <loic.lecrenier@me.com>
Date: Thu, 30 Mar 2023 14:49:25 +0200
Subject: [PATCH] Tiny refactor of query graph remove_nodes method

---
 milli/src/search/new/mod.rs         |  2 +-
 milli/src/search/new/query_graph.rs | 65 ++++++++++++++++++++---------
 milli/src/search/new/words.rs       |  2 +-
 3 files changed, 48 insertions(+), 21 deletions(-)

diff --git a/milli/src/search/new/mod.rs b/milli/src/search/new/mod.rs
index 5628ee1a9..92e8882be 100644
--- a/milli/src/search/new/mod.rs
+++ b/milli/src/search/new/mod.rs
@@ -84,7 +84,7 @@ fn resolve_maximally_reduced_query_graph(
             .collect(),
         TermsMatchingStrategy::All => vec![],
     };
-    graph.remove_nodes(&nodes_to_remove);
+    graph.remove_nodes_keep_edges(&nodes_to_remove);
 
     logger.query_for_universe(&graph);
     let docids = compute_query_graph_docids(ctx, &graph, universe)?;
diff --git a/milli/src/search/new/query_graph.rs b/milli/src/search/new/query_graph.rs
index 368e22fda..faa487299 100644
--- a/milli/src/search/new/query_graph.rs
+++ b/milli/src/search/new/query_graph.rs
@@ -177,11 +177,34 @@ impl QueryGraph {
             node.data = node_data;
         }
         let mut graph = QueryGraph { root_node, end_node, nodes };
-        graph.rebuild_edges();
+        graph.build_initial_edges();
 
         Ok(graph)
     }
 
+    /// Remove the given nodes, connecting all their predecessors to all their successors.
+    pub fn remove_nodes_keep_edges(&mut self, nodes: &[Interned<QueryNode>]) {
+        for &node_id in nodes {
+            let node = self.nodes.get(node_id);
+            let old_node_pred = node.predecessors.clone();
+            let old_node_succ = node.successors.clone();
+            for pred in old_node_pred.iter() {
+                let pred_successors = &mut self.nodes.get_mut(pred).successors;
+                pred_successors.remove(node_id);
+                pred_successors.union(&old_node_succ);
+            }
+            for succ in old_node_succ.iter() {
+                let succ_predecessors = &mut self.nodes.get_mut(succ).predecessors;
+                succ_predecessors.remove(node_id);
+                succ_predecessors.union(&old_node_pred);
+            }
+            let node = self.nodes.get_mut(node_id);
+            node.data = QueryNodeData::Deleted;
+            node.predecessors.clear();
+            node.successors.clear();
+        }
+    }
+
     /// Remove the given nodes and all their edges from the query graph.
     pub fn remove_nodes(&mut self, nodes: &[Interned<QueryNode>]) {
         for &node_id in nodes {
@@ -201,10 +224,30 @@ impl QueryGraph {
             node.predecessors.clear();
             node.successors.clear();
         }
-        self.rebuild_edges();
+    }
+    /// Simplify the query graph by removing all nodes that are disconnected from
+    /// the start or end nodes.
+    pub fn simplify(&mut self) {
+        loop {
+            let mut nodes_to_remove = vec![];
+            for (node_idx, node) in self.nodes.iter() {
+                if (!matches!(node.data, QueryNodeData::End | QueryNodeData::Deleted)
+                    && node.successors.is_empty())
+                    || (!matches!(node.data, QueryNodeData::Start | QueryNodeData::Deleted)
+                        && node.predecessors.is_empty())
+                {
+                    nodes_to_remove.push(node_idx);
+                }
+            }
+            if nodes_to_remove.is_empty() {
+                break;
+            } else {
+                self.remove_nodes(&nodes_to_remove);
+            }
+        }
     }
 
-    fn rebuild_edges(&mut self) {
+    fn build_initial_edges(&mut self) {
         for (_, node) in self.nodes.iter_mut() {
             node.successors.clear();
             node.predecessors.clear();
@@ -253,22 +296,6 @@ impl QueryGraph {
         }
     }
 
-    /// Remove all the nodes that correspond to a word starting at the given position and rebuild
-    /// the edges of the graph appropriately.
-    pub fn remove_words_starting_at_position(&mut self, position: i8) -> bool {
-        let mut nodes_to_remove = vec![];
-        for (node_idx, node) in self.nodes.iter() {
-            let QueryNodeData::Term(LocatedQueryTermSubset { term_subset: _, positions, term_ids: _  }) = &node.data else { continue };
-            if positions.start() == &position {
-                nodes_to_remove.push(node_idx);
-            }
-        }
-
-        self.remove_nodes(&nodes_to_remove);
-
-        !nodes_to_remove.is_empty()
-    }
-
     pub fn removal_order_for_terms_matching_strategy_last(&self) -> Vec<SmallBitmap<QueryNode>> {
         let (first_term_idx, last_term_idx) = {
             let mut first_term_idx = u8::MAX;
diff --git a/milli/src/search/new/words.rs b/milli/src/search/new/words.rs
index b5ee5e085..710af2243 100644
--- a/milli/src/search/new/words.rs
+++ b/milli/src/search/new/words.rs
@@ -76,7 +76,7 @@ impl<'ctx> RankingRule<'ctx, QueryGraph> for Words {
             self.exhausted = true;
         } else {
             let nodes_to_remove = self.nodes_to_remove.pop().unwrap();
-            query_graph.remove_nodes(&nodes_to_remove.iter().collect::<Vec<_>>());
+            query_graph.remove_nodes_keep_edges(&nodes_to_remove.iter().collect::<Vec<_>>());
         }
 
         Ok(Some(RankingRuleOutput { query: child_query_graph, candidates: this_bucket }))