From 50a7393c554333de51494b59834845b0fb51dc3d Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Cl=C3=A9ment=20Renault?= <clement@meilisearch.com>
Date: Fri, 21 Jun 2024 11:43:04 +0200
Subject: [PATCH] Modify the compute_query_term_subset_docids function to
 accept the universe

---
 milli/src/search/new/query_graph.rs                   |  2 +-
 .../search/new/ranking_rule_graph/exactness/mod.rs    |  2 +-
 .../ranking_rule_graph/proximity/compute_docids.rs    |  4 +---
 milli/src/search/new/ranking_rule_graph/typo/mod.rs   |  3 +--
 milli/src/search/new/ranking_rule_graph/words/mod.rs  |  3 +--
 milli/src/search/new/resolve_query_graph.rs           | 11 ++++++-----
 6 files changed, 11 insertions(+), 14 deletions(-)

diff --git a/milli/src/search/new/query_graph.rs b/milli/src/search/new/query_graph.rs
index 40e324a7e..9ab5d9dad 100644
--- a/milli/src/search/new/query_graph.rs
+++ b/milli/src/search/new/query_graph.rs
@@ -302,7 +302,7 @@ impl QueryGraph {
             for (_, node) in self.nodes.iter() {
                 match &node.data {
                     QueryNodeData::Term(t) => {
-                        let docids = compute_query_term_subset_docids(ctx, &t.term_subset)?;
+                        let docids = compute_query_term_subset_docids(ctx, None, &t.term_subset)?;
                         for id in t.term_ids.clone() {
                             term_docids
                                 .entry(id)
diff --git a/milli/src/search/new/ranking_rule_graph/exactness/mod.rs b/milli/src/search/new/ranking_rule_graph/exactness/mod.rs
index d3fa040ee..274688877 100644
--- a/milli/src/search/new/ranking_rule_graph/exactness/mod.rs
+++ b/milli/src/search/new/ranking_rule_graph/exactness/mod.rs
@@ -59,7 +59,7 @@ impl RankingRuleGraphTrait for ExactnessGraph {
             }
             ExactnessCondition::Any(dest_node) => {
                 let docids =
-                    universe & compute_query_term_subset_docids(ctx, &dest_node.term_subset)?;
+                    compute_query_term_subset_docids(ctx, Some(universe), &dest_node.term_subset)?;
                 (docids, dest_node.clone())
             }
         };
diff --git a/milli/src/search/new/ranking_rule_graph/proximity/compute_docids.rs b/milli/src/search/new/ranking_rule_graph/proximity/compute_docids.rs
index b348c682d..e73e5e6aa 100644
--- a/milli/src/search/new/ranking_rule_graph/proximity/compute_docids.rs
+++ b/milli/src/search/new/ranking_rule_graph/proximity/compute_docids.rs
@@ -22,10 +22,8 @@ pub fn compute_docids(
             (left_term, right_term, *cost)
         }
         ProximityCondition::Term { term } => {
-            let mut docids = compute_query_term_subset_docids(ctx, &term.term_subset)?;
-            docids &= universe;
             return Ok(ComputedCondition {
-                docids,
+                docids: compute_query_term_subset_docids(ctx, Some(universe), &term.term_subset)?,
                 universe_len: universe.len(),
                 start_term_subset: None,
                 end_term_subset: term.clone(),
diff --git a/milli/src/search/new/ranking_rule_graph/typo/mod.rs b/milli/src/search/new/ranking_rule_graph/typo/mod.rs
index c1256f6a1..b19163ba2 100644
--- a/milli/src/search/new/ranking_rule_graph/typo/mod.rs
+++ b/milli/src/search/new/ranking_rule_graph/typo/mod.rs
@@ -27,8 +27,7 @@ impl RankingRuleGraphTrait for TypoGraph {
     ) -> Result<ComputedCondition> {
         let TypoCondition { term, .. } = condition;
         // maybe compute_query_term_subset_docids should accept a universe as argument
-        let mut docids = compute_query_term_subset_docids(ctx, &term.term_subset)?;
-        docids &= universe;
+        let docids = compute_query_term_subset_docids(ctx, Some(universe), &term.term_subset)?;
 
         Ok(ComputedCondition {
             docids,
diff --git a/milli/src/search/new/ranking_rule_graph/words/mod.rs b/milli/src/search/new/ranking_rule_graph/words/mod.rs
index bb6baf669..19b54e76f 100644
--- a/milli/src/search/new/ranking_rule_graph/words/mod.rs
+++ b/milli/src/search/new/ranking_rule_graph/words/mod.rs
@@ -26,8 +26,7 @@ impl RankingRuleGraphTrait for WordsGraph {
     ) -> Result<ComputedCondition> {
         let WordsCondition { term, .. } = condition;
         // maybe compute_query_term_subset_docids should accept a universe as argument
-        let mut docids = compute_query_term_subset_docids(ctx, &term.term_subset)?;
-        docids &= universe;
+        let docids = compute_query_term_subset_docids(ctx, Some(universe), &term.term_subset)?;
 
         Ok(ComputedCondition {
             docids,
diff --git a/milli/src/search/new/resolve_query_graph.rs b/milli/src/search/new/resolve_query_graph.rs
index 8d89e8275..3f24319f1 100644
--- a/milli/src/search/new/resolve_query_graph.rs
+++ b/milli/src/search/new/resolve_query_graph.rs
@@ -31,6 +31,7 @@ impl<'ctx> SearchContext<'ctx> {
 }
 pub fn compute_query_term_subset_docids(
     ctx: &mut SearchContext<'_>,
+    universe: Option<&RoaringBitmap>,
     term: &QueryTermSubset,
 ) -> Result<RoaringBitmap> {
     let mut docids = RoaringBitmap::new();
@@ -49,7 +50,10 @@ pub fn compute_query_term_subset_docids(
         }
     }
 
-    Ok(docids)
+    match universe {
+        Some(universe) => Ok(docids & universe),
+        None => Ok(docids),
+    }
 }
 
 pub fn compute_query_term_subset_docids_within_field_id(
@@ -147,10 +151,7 @@ pub fn compute_query_graph_docids(
                 term_subset,
                 positions: _,
                 term_ids: _,
-            }) => {
-                let node_docids = compute_query_term_subset_docids(ctx, term_subset)?;
-                predecessors_docids & node_docids
-            }
+            }) => compute_query_term_subset_docids(ctx, Some(&predecessors_docids), term_subset)?,
             QueryNodeData::Deleted => {
                 panic!()
             }