Skip to main content

hermes_core/query/
boolean.rs

1//! Boolean query with MUST, SHOULD, and MUST_NOT clauses
2
3use std::sync::Arc;
4
5use crate::segment::SegmentReader;
6use crate::structures::TERMINATED;
7use crate::{DocId, Score};
8
9use super::planner::{
10    build_combined_bitset, build_sparse_bmp_results, build_sparse_bmp_results_filtered,
11    build_sparse_maxscore_executor, chain_predicates, combine_sparse_results, compute_idf,
12    extract_all_sparse_infos, finish_text_maxscore, prepare_per_field_grouping,
13    prepare_text_maxscore,
14};
15use super::{CountFuture, EmptyScorer, GlobalStats, Query, Scorer, ScorerFuture};
16
17/// Boolean query with MUST, SHOULD, and MUST_NOT clauses
18///
19/// When all clauses are SHOULD term queries on the same field, automatically
20/// uses MaxScore optimization for efficient top-k retrieval.
21#[derive(Default, Clone)]
22pub struct BooleanQuery {
23    pub must: Vec<Arc<dyn Query>>,
24    pub should: Vec<Arc<dyn Query>>,
25    pub must_not: Vec<Arc<dyn Query>>,
26    /// Optional global statistics for cross-segment IDF
27    global_stats: Option<Arc<GlobalStats>>,
28}
29
30impl std::fmt::Debug for BooleanQuery {
31    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
32        f.debug_struct("BooleanQuery")
33            .field("must_count", &self.must.len())
34            .field("should_count", &self.should.len())
35            .field("must_not_count", &self.must_not.len())
36            .field("has_global_stats", &self.global_stats.is_some())
37            .finish()
38    }
39}
40
41impl std::fmt::Display for BooleanQuery {
42    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
43        write!(f, "Boolean(")?;
44        let mut first = true;
45        for q in &self.must {
46            if !first {
47                write!(f, " ")?;
48            }
49            write!(f, "+{}", q)?;
50            first = false;
51        }
52        for q in &self.should {
53            if !first {
54                write!(f, " ")?;
55            }
56            write!(f, "{}", q)?;
57            first = false;
58        }
59        for q in &self.must_not {
60            if !first {
61                write!(f, " ")?;
62            }
63            write!(f, "-{}", q)?;
64            first = false;
65        }
66        write!(f, ")")
67    }
68}
69
70impl BooleanQuery {
71    pub fn new() -> Self {
72        Self::default()
73    }
74
75    pub fn must(mut self, query: impl Query + 'static) -> Self {
76        self.must.push(Arc::new(query));
77        self
78    }
79
80    pub fn should(mut self, query: impl Query + 'static) -> Self {
81        self.should.push(Arc::new(query));
82        self
83    }
84
85    pub fn must_not(mut self, query: impl Query + 'static) -> Self {
86        self.must_not.push(Arc::new(query));
87        self
88    }
89
90    /// Set global statistics for cross-segment IDF
91    pub fn with_global_stats(mut self, stats: Arc<GlobalStats>) -> Self {
92        self.global_stats = Some(stats);
93        self
94    }
95}
96
97/// Build a SHOULD-only scorer from a vec of optimized scorers.
98fn build_should_scorer<'a>(scorers: Vec<Box<dyn Scorer + 'a>>) -> Box<dyn Scorer + 'a> {
99    if scorers.is_empty() {
100        return Box::new(EmptyScorer);
101    }
102    if scorers.len() == 1 {
103        return scorers.into_iter().next().unwrap();
104    }
105    let mut scorer = BooleanScorer {
106        must: vec![],
107        should: scorers,
108        must_not: vec![],
109        current_doc: 0,
110    };
111    scorer.current_doc = scorer.find_next_match();
112    Box::new(scorer)
113}
114
115// ── Planner macro ────────────────────────────────────────────────────────
116//
117// Unified planner for both async and sync paths.  Parameterised on:
118//   $scorer_fn      – scorer_with_options | scorer_sync_with_options
119//   $get_postings_fn – get_postings | get_postings_sync
120//   $execute_fn     – execute | execute_sync
121//   $($aw)*         – .await  (present for async, absent for sync)
122//
123// Decision order:
124//   1. Single-clause unwrap
125//   2. Pure OR → text MaxScore | sparse MaxScore | per-field MaxScore
126//   3. Filter push-down → predicate-aware sparse MaxScore | PredicatedScorer
127//   4. Standard BooleanScorer fallback
128macro_rules! boolean_plan {
129    ($must:expr, $should:expr, $must_not:expr, $global_stats:expr,
130     $reader:expr, $limit:expr, $scorer_options:expr,
131     $scorer_fn:ident, $get_postings_fn:ident, $execute_fn:ident
132     $(, $aw:tt)*) => {{
133        let must: &[Arc<dyn Query>] = &$must;
134        let should_all: &[Arc<dyn Query>] = &$should;
135        let must_not: &[Arc<dyn Query>] = &$must_not;
136        let global_stats: Option<&Arc<GlobalStats>> = $global_stats;
137        let reader: &SegmentReader = $reader;
138        let limit: usize = $limit;
139        let scorer_options: super::ScorerOptions = $scorer_options;
140
141        // Cap SHOULD clauses to MAX_QUERY_TERMS, but only count queries that need
142        // posting-list cursors. Fast-field predicates (O(1) per doc) are exempt.
143        let should_capped: Vec<Arc<dyn Query>>;
144        let should: &[Arc<dyn Query>] = if should_all.len() > super::MAX_QUERY_TERMS {
145            let is_predicate: Vec<bool> = should_all
146                .iter()
147                .map(|q| q.is_filter() || q.as_doc_predicate(reader).is_some())
148                .collect();
149            let cursor_count = is_predicate.iter().filter(|&&p| !p).count();
150
151            if cursor_count > super::MAX_QUERY_TERMS {
152                let mut kept = Vec::with_capacity(should_all.len());
153                let mut cursor_kept = 0usize;
154                for (q, &is_pred) in should_all.iter().zip(is_predicate.iter()) {
155                    if is_pred {
156                        kept.push(q.clone());
157                    } else if cursor_kept < super::MAX_QUERY_TERMS {
158                        kept.push(q.clone());
159                        cursor_kept += 1;
160                    }
161                }
162                log::debug!(
163                    "BooleanQuery: capping cursor SHOULD from {} to {} ({} fast-field predicates exempt)",
164                    cursor_count,
165                    super::MAX_QUERY_TERMS,
166                    kept.len() - cursor_kept,
167                );
168                should_capped = kept;
169                &should_capped
170            } else {
171                log::debug!(
172                    "BooleanQuery: {} SHOULD clauses OK ({} need cursors, {} fast-field predicates)",
173                    should_all.len(),
174                    cursor_count,
175                    should_all.len() - cursor_count,
176                );
177                should_all
178            }
179        } else {
180            should_all
181        };
182
183        // ── 1. Single-clause optimisation ────────────────────────────────
184        if must_not.is_empty() {
185            if must.len() == 1 && should.is_empty() {
186                return must[0].$scorer_fn(reader, limit, scorer_options) $(.  $aw)* ;
187            }
188            if should.len() == 1 && must.is_empty() {
189                return should[0].$scorer_fn(reader, limit, scorer_options) $(. $aw)* ;
190            }
191        }
192
193        // ── 2. Pure OR → MaxScore optimisations ──────────────────────────
194        if must.is_empty() && must_not.is_empty() && should.len() >= 2 {
195            // 2a. Text MaxScore (single-field, all term queries)
196            if !scorer_options.collect_positions
197                && let Some((mut infos, text_field, avg_field_len, num_docs)) =
198                prepare_text_maxscore(should, reader, global_stats)
199            {
200                let mut posting_lists = Vec::with_capacity(infos.len());
201                for info in infos.drain(..) {
202                    if let Some(pl) = reader.$get_postings_fn(info.field, &info.term)
203                        $(. $aw)* ?
204                    {
205                        let idf = compute_idf(&pl, info.field, &info.term, num_docs, global_stats);
206                        posting_lists.push((pl, idf));
207                    }
208                }
209                // Seed from the cross-segment floor: this path scores final
210                // per-doc BM25 into a top-`limit` heap, so a floor carried from
211                // an already-searched segment prunes exactly (see
212                // SharedThreshold). The per-field path below stays at 0.0 —
213                // its per-field partial scores are not the final doc score.
214                let shared_threshold = std::cell::Cell::new(scorer_options.initial_threshold);
215                return finish_text_maxscore(
216                    posting_lists,
217                    avg_field_len,
218                    limit,
219                    &shared_threshold,
220                    reader.schema().index_label(),
221                    reader.schema().get_field_name(text_field).unwrap_or("?"),
222                );
223            }
224
225            // 2b. Sparse (single-field, all sparse term queries)
226            // Auto-detect: BMP executor if field has BMP index, else MaxScore
227            if let Some(infos) = extract_all_sparse_infos(should) {
228                if let Some((raw, info)) =
229                    build_sparse_bmp_results(&infos, reader, limit, &scorer_options)
230                {
231                    return Ok(combine_sparse_results(raw, info.combiner, info.field, limit));
232                }
233                if let Some((executor, info)) =
234                    build_sparse_maxscore_executor(&infos, reader, limit, None)
235                {
236                    let raw = executor.$execute_fn() $(. $aw)* ?;
237                    return Ok(combine_sparse_results(raw, info.combiner, info.field, limit));
238                }
239            }
240
241            // 2c. Per-field text MaxScore (multi-field term grouping)
242            if !scorer_options.collect_positions
243                && let Some(grouping) =
244                    prepare_per_field_grouping(should, reader, limit, global_stats)
245            {
246                let mut scorers: Vec<Box<dyn Scorer + '_>> = Vec::new();
247                // Query-local cross-group threshold seeding (see finish_text_maxscore)
248                let shared_threshold = std::cell::Cell::new(0.0f32);
249                for (field, avg_field_len, infos) in &grouping.multi_term_groups {
250                    let mut posting_lists = Vec::with_capacity(infos.len());
251                    for info in infos {
252                        if let Some(pl) = reader.$get_postings_fn(info.field, &info.term)
253                            $(. $aw)* ?
254                        {
255                            let idf = compute_idf(
256                                &pl, *field, &info.term, grouping.num_docs, global_stats,
257                            );
258                            posting_lists.push((pl, idf));
259                        }
260                    }
261                    if !posting_lists.is_empty() {
262                        scorers.push(finish_text_maxscore(
263                            posting_lists,
264                            *avg_field_len,
265                            grouping.per_field_limit,
266                            &shared_threshold,
267                            reader.schema().index_label(),
268                            reader.schema().get_field_name(*field).unwrap_or("?"),
269                        )?);
270                    }
271                }
272                for &idx in &grouping.fallback_indices {
273                    scorers.push(should[idx].$scorer_fn(
274                        reader,
275                        limit,
276                        scorer_options.without_threshold(),
277                    ) $(. $aw)* ?);
278                }
279                return Ok(build_should_scorer(scorers));
280            }
281        }
282
283        // ── 3. Filter push-down (MUST + SHOULD) ─────────────────────────
284        if !scorer_options.collect_positions
285            && !should.is_empty()
286            && !must.is_empty()
287        {
288            // Pre-check: is SHOULD all-sparse? This determines whether we can
289            // use bitset fallback for MUST clauses that lack fast-field predicates.
290            // For sparse SHOULD, the predicate is pushed into BMP/MaxScore traversal
291            // so all qualifying docs are found. For text SHOULD, we must NOT convert
292            // MUST to a predicate (PredicatedScorer would drop MUST-only docs that
293            // don't match SHOULD), so those go to verifier → BooleanScorer.
294            let should_is_sparse = extract_all_sparse_infos(should).is_some();
295
296            // 3a. Compile MUST → predicates (O(1)) vs verifier scorers (seek)
297            //
298            // Priority: as_doc_predicate (fast-field O(1)) > as_doc_bitset
299            // (posting-list materialization, O(1) lookup, sparse-SHOULD only)
300            // > verifier scorer (seek).
301            let mut predicates: Vec<super::DocPredicate<'_>> = Vec::new();
302            let mut must_verifiers: Vec<Box<dyn super::Scorer + '_>> = Vec::new();
303            for q in must {
304                if let Some(pred) = q.as_doc_predicate(reader) {
305                    log::debug!("BooleanQuery planner 3a: MUST clause → predicate ({})", q);
306                    predicates.push(pred);
307                } else if should_is_sparse {
308                    if let Some(bitset) = q.as_doc_bitset(reader) {
309                        log::debug!("BooleanQuery planner 3a: MUST clause → bitset predicate ({})", q);
310                        predicates.push(Box::new(move |doc_id| bitset.contains(doc_id)));
311                    } else {
312                        log::debug!("BooleanQuery planner 3a: MUST clause → verifier scorer ({})", q);
313                        must_verifiers.push(q.$scorer_fn(
314                            reader, limit, scorer_options.without_threshold()
315                        ) $(. $aw)* ?);
316                    }
317                } else {
318                    log::debug!("BooleanQuery planner 3a: MUST clause → verifier scorer ({})", q);
319                    must_verifiers.push(q.$scorer_fn(
320                        reader, limit, scorer_options.without_threshold()
321                    ) $(. $aw)* ?);
322                }
323            }
324            // Compile MUST_NOT → negated predicates vs verifier scorers
325            let mut must_not_verifiers: Vec<Box<dyn super::Scorer + '_>> = Vec::new();
326            for q in must_not {
327                if let Some(pred) = q.as_doc_predicate(reader) {
328                    let negated: super::DocPredicate<'_> =
329                        Box::new(move |doc_id| !pred(doc_id));
330                    predicates.push(negated);
331                } else if should_is_sparse {
332                    if let Some(bitset) = q.as_doc_bitset(reader) {
333                        log::debug!("BooleanQuery planner 3a: MUST_NOT clause → bitset predicate ({})", q);
334                        predicates.push(Box::new(move |doc_id| !bitset.contains(doc_id)));
335                    } else {
336                        must_not_verifiers.push(q.$scorer_fn(
337                            reader, limit, scorer_options.without_threshold()
338                        ) $(. $aw)* ?);
339                    }
340                } else {
341                    must_not_verifiers.push(q.$scorer_fn(
342                        reader, limit, scorer_options.without_threshold()
343                    ) $(. $aw)* ?);
344                }
345            }
346
347            // 3b. Fast path: pure predicates + sparse SHOULD → BMP or MaxScore w/ predicate
348            if must_verifiers.is_empty()
349                && must_not_verifiers.is_empty()
350                && !predicates.is_empty()
351            {
352                if let Some(infos) = extract_all_sparse_infos(should) {
353                    // Try BMP with bitset first: build compact bitset from MUST/MUST_NOT
354                    // posting lists (O(M) for term queries) for fast per-slot lookup.
355                    let bitset_result = build_combined_bitset(must, must_not, reader);
356                    if let Some(ref bitset) = bitset_result {
357                        let bitset_pred = |doc_id: crate::DocId| bitset.contains(doc_id);
358                        if let Some((raw, info)) =
359                            build_sparse_bmp_results_filtered(
360                                &infos, reader, limit, &bitset_pred, &scorer_options
361                            )
362                        {
363                            log::debug!(
364                                "BooleanQuery planner: bitset-aware sparse BMP, {} dims, {} matching docs",
365                                infos.len(),
366                                bitset.count()
367                            );
368                            return Ok(combine_sparse_results(raw, info.combiner, info.field, limit));
369                        }
370                    }
371
372                    // Fallback: closure predicate (for queries that don't support bitsets)
373                    let combined = chain_predicates(predicates);
374                    if let Some((raw, info)) =
375                        build_sparse_bmp_results_filtered(
376                            &infos, reader, limit, &*combined, &scorer_options
377                        )
378                    {
379                        log::debug!(
380                            "BooleanQuery planner: predicate-aware sparse BMP, {} dims",
381                            infos.len()
382                        );
383                        return Ok(combine_sparse_results(raw, info.combiner, info.field, limit));
384                    }
385                    // Try MaxScore with predicate
386                    if let Some((executor, info)) =
387                        build_sparse_maxscore_executor(&infos, reader, limit, Some(combined))
388                    {
389                        log::debug!(
390                            "BooleanQuery planner: predicate-aware sparse MaxScore, {} dims",
391                            infos.len()
392                        );
393                        let raw = executor.$execute_fn() $(. $aw)* ?;
394                        return Ok(combine_sparse_results(raw, info.combiner, info.field, limit));
395                    }
396                    // predicates consumed — cannot fall through; rebuild them
397                    // (this path only triggers if neither sparse index exists)
398                    // should_is_sparse is true here (we're inside extract_all_sparse_infos)
399                    predicates = Vec::new();
400                    for q in must {
401                        if let Some(pred) = q.as_doc_predicate(reader) {
402                            predicates.push(pred);
403                        } else if let Some(bitset) = q.as_doc_bitset(reader) {
404                            predicates.push(Box::new(move |doc_id| bitset.contains(doc_id)));
405                        }
406                    }
407                    for q in must_not {
408                        if let Some(pred) = q.as_doc_predicate(reader) {
409                            let negated: super::DocPredicate<'_> =
410                                Box::new(move |doc_id| !pred(doc_id));
411                            predicates.push(negated);
412                        } else if let Some(bitset) = q.as_doc_bitset(reader) {
413                            predicates.push(Box::new(move |doc_id| !bitset.contains(doc_id)));
414                        }
415                    }
416                }
417            }
418
419            // 3c. PredicatedScorer fallback. Filters can discard candidates,
420            // so use the same bounded candidate budget as other query paths.
421            let has_filters = !predicates.is_empty()
422                || !must_verifiers.is_empty()
423                || !must_not_verifiers.is_empty();
424            let should_limit = if has_filters {
425                super::max_candidate_limit(limit)
426            } else {
427                limit
428            };
429            let should_scorer = if should.len() == 1 {
430                should[0].$scorer_fn(
431                    reader, should_limit, scorer_options.without_threshold()
432                ) $(. $aw)* ?
433            } else {
434                let sub = BooleanQuery {
435                    must: Vec::new(),
436                    should: should.to_vec(),
437                    must_not: Vec::new(),
438                    global_stats: global_stats.cloned(),
439                };
440                sub.$scorer_fn(
441                    reader, should_limit, scorer_options.without_threshold()
442                ) $(. $aw)* ?
443            };
444
445            let use_predicated =
446                must_verifiers.is_empty() || should_scorer.size_hint() >= limit as u32;
447
448            if use_predicated {
449                log::debug!(
450                    "BooleanQuery planner: PredicatedScorer {} preds + {} must_v + {} must_not_v, \
451                     SHOULD size_hint={}, over_fetch={}",
452                    predicates.len(), must_verifiers.len(), must_not_verifiers.len(),
453                    should_scorer.size_hint(), should_limit
454                );
455                return Ok(Box::new(super::PredicatedScorer::new(
456                    should_scorer, predicates, must_verifiers, must_not_verifiers,
457                )));
458            }
459
460            // size_hint < limit with verifiers → BooleanScorer
461            log::debug!(
462                "BooleanQuery planner: BooleanScorer fallback, size_hint={} < limit={}, \
463                 {} must_v + {} must_not_v",
464                should_scorer.size_hint(), limit,
465                must_verifiers.len(), must_not_verifiers.len()
466            );
467            let mut scorer = BooleanScorer {
468                must: must_verifiers,
469                should: vec![should_scorer],
470                must_not: must_not_verifiers,
471                current_doc: 0,
472            };
473            scorer.current_doc = scorer.find_next_match();
474            return Ok(Box::new(scorer));
475        }
476
477        // ── 4. Standard BooleanScorer fallback ───────────────────────────
478        let mut must_scorers = Vec::with_capacity(must.len());
479        for q in must {
480            must_scorers.push(q.$scorer_fn(
481                reader, limit, scorer_options.without_threshold()
482            ) $(. $aw)* ?);
483        }
484        let mut should_scorers = Vec::with_capacity(should.len());
485        for q in should {
486            should_scorers.push(q.$scorer_fn(
487                reader, limit, scorer_options.without_threshold()
488            ) $(. $aw)* ?);
489        }
490        let mut must_not_scorers = Vec::with_capacity(must_not.len());
491        for q in must_not {
492            must_not_scorers.push(q.$scorer_fn(
493                reader, limit, scorer_options.without_threshold()
494            ) $(. $aw)* ?);
495        }
496        let mut scorer = BooleanScorer {
497            must: must_scorers,
498            should: should_scorers,
499            must_not: must_not_scorers,
500            current_doc: 0,
501        };
502        scorer.current_doc = scorer.find_next_match();
503        Ok(Box::new(scorer) as Box<dyn Scorer + '_>)
504    }};
505}
506
507impl Query for BooleanQuery {
508    fn scorer<'a>(&self, reader: &'a SegmentReader, limit: usize) -> ScorerFuture<'a> {
509        self.scorer_with_options(reader, limit, super::ScorerOptions::with_positions())
510    }
511
512    fn scorer_with_options<'a>(
513        &self,
514        reader: &'a SegmentReader,
515        limit: usize,
516        options: super::ScorerOptions,
517    ) -> ScorerFuture<'a> {
518        let must = self.must.clone();
519        let should = self.should.clone();
520        let must_not = self.must_not.clone();
521        let global_stats = self.global_stats.clone();
522        Box::pin(async move {
523            boolean_plan!(
524                must,
525                should,
526                must_not,
527                global_stats.as_ref(),
528                reader,
529                limit,
530                options,
531                scorer_with_options,
532                get_postings,
533                execute,
534                await
535            )
536        })
537    }
538
539    #[cfg(feature = "sync")]
540    fn scorer_sync<'a>(
541        &self,
542        reader: &'a SegmentReader,
543        limit: usize,
544    ) -> crate::Result<Box<dyn Scorer + 'a>> {
545        self.scorer_sync_with_options(reader, limit, super::ScorerOptions::with_positions())
546    }
547
548    #[cfg(feature = "sync")]
549    fn scorer_sync_with_options<'a>(
550        &self,
551        reader: &'a SegmentReader,
552        limit: usize,
553        options: super::ScorerOptions,
554    ) -> crate::Result<Box<dyn Scorer + 'a>> {
555        boolean_plan!(
556            self.must,
557            self.should,
558            self.must_not,
559            self.global_stats.as_ref(),
560            reader,
561            limit,
562            options,
563            scorer_sync_with_options,
564            get_postings_sync,
565            execute_sync
566        )
567    }
568
569    fn as_doc_bitset(&self, reader: &SegmentReader) -> Option<super::DocBitset> {
570        if self.must.is_empty() && self.should.is_empty() {
571            return None;
572        }
573
574        let num_docs = reader.num_docs();
575
576        // MUST clauses: intersect bitsets (AND)
577        let mut result: Option<super::DocBitset> = None;
578        for q in &self.must {
579            let bs = q.as_doc_bitset(reader)?;
580            match result {
581                None => result = Some(bs),
582                Some(ref mut acc) => acc.intersect_with(&bs),
583            }
584        }
585
586        // SHOULD clauses: union bitsets (OR), then intersect with MUST result
587        if !self.should.is_empty() {
588            let mut should_union = super::DocBitset::new(num_docs);
589            for q in &self.should {
590                let bs = q.as_doc_bitset(reader)?;
591                should_union.union_with(&bs);
592            }
593            match result {
594                None => result = Some(should_union),
595                Some(ref mut acc) => {
596                    // When MUST clauses exist, SHOULD is optional (doesn't filter).
597                    // When no MUST clauses, at least one SHOULD must match.
598                    if self.must.is_empty() {
599                        *acc = should_union;
600                    }
601                }
602            }
603        }
604
605        // MUST_NOT clauses: subtract bitsets (ANDNOT)
606        if let Some(ref mut acc) = result {
607            for q in &self.must_not {
608                {
609                    let bs = q.as_doc_bitset(reader)?;
610                    acc.subtract(&bs);
611                }
612            }
613        }
614
615        result
616    }
617
618    fn as_doc_predicate<'a>(&self, reader: &'a SegmentReader) -> Option<super::DocPredicate<'a>> {
619        // Need at least some clauses
620        if self.must.is_empty() && self.should.is_empty() {
621            return None;
622        }
623
624        // Try converting all clauses to predicates; bail if any child can't
625        let must_preds: Vec<_> = self
626            .must
627            .iter()
628            .map(|q| q.as_doc_predicate(reader))
629            .collect::<Option<Vec<_>>>()?;
630        let should_preds: Vec<_> = self
631            .should
632            .iter()
633            .map(|q| q.as_doc_predicate(reader))
634            .collect::<Option<Vec<_>>>()?;
635        let must_not_preds: Vec<_> = self
636            .must_not
637            .iter()
638            .map(|q| q.as_doc_predicate(reader))
639            .collect::<Option<Vec<_>>>()?;
640
641        let has_must = !must_preds.is_empty();
642
643        Some(Box::new(move |doc_id| {
644            // All MUST predicates must pass
645            if !must_preds.iter().all(|p| p(doc_id)) {
646                return false;
647            }
648            // When there are no MUST clauses, at least one SHOULD must pass
649            if !has_must && !should_preds.is_empty() && !should_preds.iter().any(|p| p(doc_id)) {
650                return false;
651            }
652            // No MUST_NOT predicate should pass
653            must_not_preds.iter().all(|p| !p(doc_id))
654        }))
655    }
656
657    fn count_estimate<'a>(&self, reader: &'a SegmentReader) -> CountFuture<'a> {
658        let must = self.must.clone();
659        let should = self.should.clone();
660
661        Box::pin(async move {
662            if !must.is_empty() {
663                let mut estimates = Vec::with_capacity(must.len());
664                for q in &must {
665                    estimates.push(q.count_estimate(reader).await?);
666                }
667                estimates
668                    .into_iter()
669                    .min()
670                    .ok_or_else(|| crate::Error::Corruption("Empty must clause".to_string()))
671            } else if !should.is_empty() {
672                let mut sum = 0u32;
673                for q in &should {
674                    sum = sum.saturating_add(q.count_estimate(reader).await?);
675                }
676                Ok(sum)
677            } else {
678                Ok(0)
679            }
680        })
681    }
682}
683
684struct BooleanScorer<'a> {
685    must: Vec<Box<dyn Scorer + 'a>>,
686    should: Vec<Box<dyn Scorer + 'a>>,
687    must_not: Vec<Box<dyn Scorer + 'a>>,
688    current_doc: DocId,
689}
690
691impl BooleanScorer<'_> {
692    fn find_next_match(&mut self) -> DocId {
693        if self.must.is_empty() && self.should.is_empty() {
694            return TERMINATED;
695        }
696
697        loop {
698            let candidate = if !self.must.is_empty() {
699                let mut max_doc = self
700                    .must
701                    .iter()
702                    .map(|s| s.doc())
703                    .max()
704                    .unwrap_or(TERMINATED);
705
706                if max_doc == TERMINATED {
707                    return TERMINATED;
708                }
709
710                loop {
711                    let mut all_match = true;
712                    for scorer in &mut self.must {
713                        let doc = scorer.seek(max_doc);
714                        if doc == TERMINATED {
715                            return TERMINATED;
716                        }
717                        if doc > max_doc {
718                            max_doc = doc;
719                            all_match = false;
720                            break;
721                        }
722                    }
723                    if all_match {
724                        break;
725                    }
726                }
727                max_doc
728            } else {
729                self.should
730                    .iter()
731                    .map(|s| s.doc())
732                    .filter(|&d| d != TERMINATED)
733                    .min()
734                    .unwrap_or(TERMINATED)
735            };
736
737            if candidate == TERMINATED {
738                return TERMINATED;
739            }
740
741            let excluded = self.must_not.iter_mut().any(|scorer| {
742                let doc = scorer.seek(candidate);
743                doc == candidate
744            });
745
746            if !excluded {
747                // Seek SHOULD scorers to candidate so score() can see their contributions
748                for scorer in &mut self.should {
749                    scorer.seek(candidate);
750                }
751                self.current_doc = candidate;
752                return candidate;
753            }
754
755            // Advance past excluded candidate
756            if !self.must.is_empty() {
757                for scorer in &mut self.must {
758                    scorer.advance();
759                }
760            } else {
761                // For SHOULD-only: seek all scorers past the excluded candidate
762                for scorer in &mut self.should {
763                    if scorer.doc() <= candidate && scorer.doc() != TERMINATED {
764                        scorer.seek(candidate + 1);
765                    }
766                }
767            }
768        }
769    }
770}
771
772impl super::docset::DocSet for BooleanScorer<'_> {
773    fn doc(&self) -> DocId {
774        self.current_doc
775    }
776
777    fn advance(&mut self) -> DocId {
778        if !self.must.is_empty() {
779            for scorer in &mut self.must {
780                scorer.advance();
781            }
782        } else {
783            for scorer in &mut self.should {
784                if scorer.doc() == self.current_doc {
785                    scorer.advance();
786                }
787            }
788        }
789
790        self.current_doc = self.find_next_match();
791        self.current_doc
792    }
793
794    fn seek(&mut self, target: DocId) -> DocId {
795        for scorer in &mut self.must {
796            scorer.seek(target);
797        }
798
799        for scorer in &mut self.should {
800            scorer.seek(target);
801        }
802
803        self.current_doc = self.find_next_match();
804        self.current_doc
805    }
806
807    fn size_hint(&self) -> u32 {
808        if !self.must.is_empty() {
809            self.must.iter().map(|s| s.size_hint()).min().unwrap_or(0)
810        } else {
811            self.should.iter().map(|s| s.size_hint()).sum()
812        }
813    }
814}
815
816impl Scorer for BooleanScorer<'_> {
817    fn score(&self) -> Score {
818        let mut total = 0.0;
819
820        for scorer in &self.must {
821            if scorer.doc() == self.current_doc {
822                total += scorer.score();
823            }
824        }
825
826        for scorer in &self.should {
827            if scorer.doc() == self.current_doc {
828                total += scorer.score();
829            }
830        }
831
832        total
833    }
834
835    fn matched_positions(&self) -> Option<super::MatchedPositions> {
836        let mut all_positions: super::MatchedPositions = Vec::new();
837
838        for scorer in &self.must {
839            if scorer.doc() == self.current_doc
840                && let Some(positions) = scorer.matched_positions()
841            {
842                all_positions.extend(positions);
843            }
844        }
845
846        for scorer in &self.should {
847            if scorer.doc() == self.current_doc
848                && let Some(positions) = scorer.matched_positions()
849            {
850                all_positions.extend(positions);
851            }
852        }
853
854        if all_positions.is_empty() {
855            None
856        } else {
857            Some(all_positions)
858        }
859    }
860}
861
862#[cfg(test)]
863mod tests {
864    use super::*;
865    use crate::dsl::Field;
866    use crate::query::{QueryDecomposition, TermQuery};
867
868    #[test]
869    fn test_maxscore_eligible_pure_or_same_field() {
870        // Pure OR query with multiple terms in same field should be MaxScore-eligible
871        let query = BooleanQuery::new()
872            .should(TermQuery::text(Field(0), "hello"))
873            .should(TermQuery::text(Field(0), "world"))
874            .should(TermQuery::text(Field(0), "foo"));
875
876        // All clauses should return term info
877        assert!(
878            query
879                .should
880                .iter()
881                .all(|q| matches!(q.decompose(), QueryDecomposition::TextTerm(_)))
882        );
883
884        // All should be same field
885        let infos: Vec<_> = query
886            .should
887            .iter()
888            .filter_map(|q| match q.decompose() {
889                QueryDecomposition::TextTerm(info) => Some(info),
890                _ => None,
891            })
892            .collect();
893        assert_eq!(infos.len(), 3);
894        assert!(infos.iter().all(|i| i.field == Field(0)));
895    }
896
897    #[test]
898    fn test_maxscore_not_eligible_different_fields() {
899        // OR query with terms in different fields should NOT use MaxScore
900        let query = BooleanQuery::new()
901            .should(TermQuery::text(Field(0), "hello"))
902            .should(TermQuery::text(Field(1), "world")); // Different field!
903
904        let infos: Vec<_> = query
905            .should
906            .iter()
907            .filter_map(|q| match q.decompose() {
908                QueryDecomposition::TextTerm(info) => Some(info),
909                _ => None,
910            })
911            .collect();
912        assert_eq!(infos.len(), 2);
913        // Fields are different, MaxScore should not be used
914        assert!(infos[0].field != infos[1].field);
915    }
916
917    #[test]
918    fn test_maxscore_not_eligible_with_must() {
919        // Query with MUST clause should NOT use MaxScore optimization
920        let query = BooleanQuery::new()
921            .must(TermQuery::text(Field(0), "required"))
922            .should(TermQuery::text(Field(0), "hello"))
923            .should(TermQuery::text(Field(0), "world"));
924
925        // Has MUST clause, so MaxScore optimization should not kick in
926        assert!(!query.must.is_empty());
927    }
928
929    #[test]
930    fn test_maxscore_not_eligible_with_must_not() {
931        // Query with MUST_NOT clause should NOT use MaxScore optimization
932        let query = BooleanQuery::new()
933            .should(TermQuery::text(Field(0), "hello"))
934            .should(TermQuery::text(Field(0), "world"))
935            .must_not(TermQuery::text(Field(0), "excluded"));
936
937        // Has MUST_NOT clause, so MaxScore optimization should not kick in
938        assert!(!query.must_not.is_empty());
939    }
940
941    #[test]
942    fn test_maxscore_not_eligible_single_term() {
943        // Single SHOULD clause should NOT use MaxScore (no benefit)
944        let query = BooleanQuery::new().should(TermQuery::text(Field(0), "hello"));
945
946        // Only one term, MaxScore not beneficial
947        assert_eq!(query.should.len(), 1);
948    }
949
950    #[test]
951    fn test_term_query_info_extraction() {
952        let term_query = TermQuery::text(Field(42), "test");
953        match term_query.decompose() {
954            QueryDecomposition::TextTerm(info) => {
955                assert_eq!(info.field, Field(42));
956                assert_eq!(info.term, b"test");
957            }
958            _ => panic!("Expected TextTerm decomposition"),
959        }
960    }
961
962    #[test]
963    fn test_boolean_query_no_term_info() {
964        // BooleanQuery itself should not return term info
965        let query = BooleanQuery::new().should(TermQuery::text(Field(0), "hello"));
966
967        assert!(matches!(query.decompose(), QueryDecomposition::Opaque));
968    }
969}