Simplify some search function by reducing the number of parameters

src/search.rs

@@ -80,8 +80,7 @@ impl<'a> Search<'a> {
     /// Fetch the words from the given FST related to the given DFAs along with
     /// the associated documents ids.
     fn fetch_words_docids(
-        rtxn: &heed::RoTxn,
+        &self,
-        index: &Index,
         fst: &fst::Set<&[u8]>,
         dfas: Vec<(String, bool, DFA)>,
     ) -> anyhow::Result<Vec<(HashMap<String, (u8, RoaringBitmap)>, RoaringBitmap)>>
@@ -98,7 +97,7 @@ impl<'a> Search<'a> {
             while let Some((word, state)) = stream.next() {
 
                 let word = std::str::from_utf8(word)?;
-                let docids = index.word_docids.get(rtxn, word)?.unwrap();
+                let docids = self.index.word_docids.get(self.rtxn, word)?.unwrap();
                 let distance = dfa.distance(state);
                 unions_docids.union_with(&docids);
                 acc_derived_words.insert(word.to_string(), (distance.to_u8(), docids));
@@ -134,8 +133,7 @@ impl<'a> Search<'a> {
     }
 
     fn fecth_keywords(
-        rtxn: &heed::RoTxn,
+        &self,
-        index: &Index,
         derived_words: &[(HashMap<String, (u8, RoaringBitmap)>, RoaringBitmap)],
         candidate: DocumentId,
     ) -> anyhow::Result<Vec<IntoIter>>
@@ -148,7 +146,7 @@ impl<'a> Search<'a> {
             for (word, (_distance, docids)) in words {
 
                 if !docids.contains(candidate) { continue; }
-                if let Some(positions) = index.docid_word_positions.get(rtxn, &(candidate, word))? {
+                if let Some(positions) = self.index.docid_word_positions.get(self.rtxn, &(candidate, word))? {
                     union_positions.union_with(&positions);
                 }
             }
@@ -158,43 +156,10 @@ impl<'a> Search<'a> {
         Ok(keywords)
     }
 
-    pub fn execute(&self) -> anyhow::Result<SearchResult> {
+    fn words_pair_combinations<'h>(
-        let rtxn = self.rtxn;
+        w1: &'h HashMap<String, (u8, RoaringBitmap)>,
-        let index = self.index;
+        w2: &'h HashMap<String, (u8, RoaringBitmap)>,
-        let limit = self.limit;
+    ) -> Vec<(&'h str, &'h str)>
-
-        let fst = match index.fst(rtxn)? {
-            Some(fst) => fst,
-            None => return Ok(Default::default()),
-        };
-
-        // Construct the DFAs related to the query words.
-        // TODO do a placeholder search when query string isn't present.
-        let dfas = match &self.query {
-            Some(q) => Self::generate_query_dfas(q),
-            None => return Ok(Default::default()),
-        };
-
-        if dfas.is_empty() {
-            return Ok(Default::default());
-        }
-
-        let derived_words = Self::fetch_words_docids(rtxn, index, &fst, dfas)?;
-        let candidates = Self::compute_candidates(&derived_words);
-
-        debug!("candidates: {:?}", candidates);
-
-        // If there is only one query word, no need to compute the best proximities.
-        if derived_words.len() == 1 || candidates.is_empty() {
-            let found_words = derived_words.into_iter().flat_map(|(w, _)| w).map(|(w, _)| w).collect();
-            let documents_ids = candidates.iter().take(limit).collect();
-            return Ok(SearchResult { found_words, documents_ids });
-        }
-
-        fn words_pair_combinations<'a>(
-            w1: &'a HashMap<String, (u8, RoaringBitmap)>,
-            w2: &'a HashMap<String, (u8, RoaringBitmap)>,
-        ) -> Vec<(&'a str, &'a str)>
     {
         let mut pairs = Vec::new();
         for (w1, (_typos, docids1)) in w1 {
@@ -208,8 +173,7 @@ impl<'a> Search<'a> {
     }
 
     fn depth_first_search(
-            index: &Index,
+        &self,
-            rtxn: &heed::RoTxn,
         words: &[(HashMap<String, (u8, RoaringBitmap)>, RoaringBitmap)],
         candidates: &RoaringBitmap,
         parent_docids: Option<&RoaringBitmap>,
@@ -217,7 +181,7 @@ impl<'a> Search<'a> {
     ) -> anyhow::Result<Option<RoaringBitmap>>
     {
         let (words1, words2) = (&words[0].0, &words[1].0);
-            let pairs = words_pair_combinations(words1, words2);
+        let pairs = Self::words_pair_combinations(words1, words2);
 
         for proximity in 1..=8 {
             let mut docids = match union_cache.entry((words.len(), proximity)) {
@@ -229,7 +193,7 @@ impl<'a> Search<'a> {
                     } else {
                         for (w1, w2) in pairs.iter().cloned() {
                             let key = (w1, w2, proximity);
-                                if let Some(di) = index.word_pair_proximity_docids.get(rtxn, &key)? {
+                            if let Some(di) = self.index.word_pair_proximity_docids.get(self.rtxn, &key)? {
                                 docids.union_with(&di);
                             }
                         }
@@ -246,7 +210,7 @@ impl<'a> Search<'a> {
                 let words = &words[1..];
                 // We are the last word.
                 if words.len() < 2 { return Ok(Some(docids)) }
-                    if let Some(di) = depth_first_search(index, rtxn, words, candidates, Some(&docids), union_cache)? {
+                if let Some(di) = self.depth_first_search(words, candidates, Some(&docids), union_cache)? {
                     return Ok(Some(di))
                 }
             }
@@ -255,9 +219,38 @@ impl<'a> Search<'a> {
         Ok(None)
     }
 
-        let mut union_cache = HashMap::new();
+    pub fn execute(&self) -> anyhow::Result<SearchResult> {
-        let answer = depth_first_search(index, rtxn, &derived_words, &candidates, None, &mut union_cache)?;
+        let limit = self.limit;
 
+        let fst = match self.index.fst(self.rtxn)? {
+            Some(fst) => fst,
+            None => return Ok(Default::default()),
+        };
+
+        // Construct the DFAs related to the query words.
+        // TODO do a placeholder search when query string isn't present.
+        let dfas = match &self.query {
+            Some(q) => Self::generate_query_dfas(q),
+            None => return Ok(Default::default()),
+        };
+
+        if dfas.is_empty() {
+            return Ok(Default::default());
+        }
+
+        let derived_words = self.fetch_words_docids(&fst, dfas)?;
+        let mut candidates = Self::compute_candidates(&derived_words);
+
+        debug!("candidates: {:?}", candidates);
+
+        // If there is only one query word, no need to compute the best proximities.
+        if derived_words.len() == 1 || candidates.is_empty() {
+            let found_words = derived_words.into_iter().flat_map(|(w, _)| w).map(|(w, _)| w).collect();
+            let documents_ids = candidates.iter().take(limit).collect();
+            return Ok(SearchResult { found_words, documents_ids });
+        }
+
+        let mut union_cache = HashMap::new();
         let mut documents = Vec::new();
         if let Some(answer) = answer {
             documents.push(answer);