Add settings to force milli to exhaustively compute the total number of hits

milli/src/search/criteria/initial.rs

@@ -1,17 +1,22 @@
 use roaring::RoaringBitmap;
 
 use super::{Criterion, CriterionParameters, CriterionResult};
+use crate::search::criteria::{resolve_query_tree, Context};
 use crate::search::query_tree::Operation;
 use crate::Result;
 
-pub struct Initial {
+pub struct Initial<'t> {
+    ctx: &'t dyn Context<'t>,
     answer: Option<CriterionResult>,
+    exhaustive_number_hits: bool,
 }
 
-impl Initial {
+impl<'t> Initial<'t> {
     pub fn new(
+        ctx: &'t dyn Context<'t>,
         query_tree: Option<Operation>,
         filtered_candidates: Option<RoaringBitmap>,
+        exhaustive_number_hits: bool,
     ) -> Initial {
         let answer = CriterionResult {
             query_tree,
@@ -19,13 +24,28 @@ impl Initial {
             filtered_candidates,
             bucket_candidates: None,
         };
-        Initial { answer: Some(answer) }
+        Initial { ctx, answer: Some(answer), exhaustive_number_hits }
     }
 }
 
-impl Criterion for Initial {
+impl Criterion for Initial<'_> {
     #[logging_timer::time("Initial::{}")]
-    fn next(&mut self, _: &mut CriterionParameters) -> Result<Option<CriterionResult>> {
-        Ok(self.answer.take())
+    fn next(&mut self, params: &mut CriterionParameters) -> Result<Option<CriterionResult>> {
+        self.answer
+            .take()
+            .map(|mut answer| {
+                if self.exhaustive_number_hits && answer.query_tree.is_some() {
+                    let candidates = resolve_query_tree(
+                        self.ctx,
+                        answer.query_tree.as_ref().unwrap(),
+                        &mut params.wdcache,
+                    )?;
+
+                    answer.candidates = Some(candidates.clone());
+                    answer.bucket_candidates = Some(candidates);
+                }
+                Ok(answer)
+            })
+            .transpose()
     }
 }