lucisearch 0.8.0

//! Filter-mode disjunction over many sub-scorers using a windowed bitset.
//!
//! Matches Lucene's `BooleanScorer` and Tantivy's `BufferedUnionScorer`
//! pattern: refills a fixed-size bitset window from each sub-scorer, then
//! iterates set bits. Designed for filter-mode (constant score) use cases
//! like prefix/wildcard/regexp rewrites and `bool/filter` clauses.
//!
//! Sub-scorer scores are NOT accumulated — `score()` returns the
//! configured constant boost (default 1.0). This matches Lucene's
//! `MultiTermQueryConstantScoreBlendedWrapper` semantics where the
//! disjunction is wrapped in a `ConstantScoreScorer`.
//!
//! Window size matches `MaxScoreBulkScorer` at 2048 docs (32 × u64) for
//! consistency with Luci's existing windowed-bitset infrastructure.
//!
//! See [[fix-disjunction-heap-inefficiency]].

use crate::core::{DocId, NO_MORE_DOCS, Scorer, TwoPhaseIterator};

const WINDOW_SIZE: usize = 2048;
const WINDOW_WORDS: usize = WINDOW_SIZE / 64;
const WINDOW_SIZE_U32: u32 = WINDOW_SIZE as u32;

/// Filter-mode disjunction over many sub-scorers using a windowed bitset.
pub struct BufferedUnionScorer {
    scorers: Vec<Box<dyn Scorer>>,
    /// Bitset for the current window (2048 bits).
    bitset: [u64; WINDOW_WORDS],
    /// Doc ID corresponding to bit 0 of the current window.
    window_base: u32,
    /// Currently emitted doc. NO_MORE_DOCS if exhausted.
    current: DocId,
    /// Position in the current window: bits at `cursor..` are not yet emitted.
    cursor: usize,
    /// Constant score returned by `score()`.
    boost: f32,
}

impl BufferedUnionScorer {
    /// Build a filter-mode union over the given scorers, returning a
    /// constant score of 1.0 per matching doc.
    pub fn new(scorers: Vec<Box<dyn Scorer>>) -> Self {
        Self::with_boost(scorers, 1.0)
    }

    /// Build with a custom constant boost.
    pub fn with_boost(mut scorers: Vec<Box<dyn Scorer>>, boost: f32) -> Self {
        // Drop already-exhausted scorers in place — no extra allocation.
        scorers.retain(|s| s.doc_id() != NO_MORE_DOCS);

        let mut union = Self {
            scorers,
            bitset: [0u64; WINDOW_WORDS],
            window_base: 0,
            current: NO_MORE_DOCS,
            cursor: 0,
            boost,
        };

        if !union.scorers.is_empty() {
            union.advance_to_next_set_bit();
        }

        union
    }

    /// Find the minimum doc_id across all sub-scorers.
    fn min_scorer_doc(&self) -> DocId {
        let mut min = NO_MORE_DOCS;
        for s in &self.scorers {
            let d = s.doc_id();
            if d != NO_MORE_DOCS && d < min {
                min = d;
            }
        }
        min
    }

    /// Refill the bitset window starting at `base`. Each sub-scorer is
    /// advanced to the window range and its matching docs are added to
    /// the bitset. After this call, every sub-scorer is at a doc
    /// `>= base + WINDOW_SIZE` (or exhausted).
    fn refill_from(&mut self, base: u32) {
        self.window_base = base;
        self.bitset = [0u64; WINDOW_WORDS];
        self.cursor = 0;

        let window_end = base.saturating_add(WINDOW_SIZE_U32);

        for scorer in &mut self.scorers {
            // Advance to base if behind
            let mut doc = scorer.doc_id();
            if doc != NO_MORE_DOCS && doc.as_u32() < base {
                doc = scorer.advance(DocId::new(base));
            }
            // Walk through docs in [base, window_end)
            while doc != NO_MORE_DOCS && doc.as_u32() < window_end {
                let idx = (doc.as_u32() - base) as usize;
                self.bitset[idx / 64] |= 1u64 << (idx % 64);
                doc = scorer.next();
            }
        }
    }

    /// Find the next set bit at or after `self.cursor` in the current
    /// window. Returns the bit position if found.
    fn next_set_bit_in_window(&self) -> Option<usize> {
        let mut word_idx = self.cursor / 64;
        if word_idx >= WINDOW_WORDS {
            return None;
        }

        // Mask out bits before cursor in the current word
        let bit_offset = self.cursor % 64;
        let mask = u64::MAX << bit_offset;
        let mut word = self.bitset[word_idx] & mask;

        loop {
            if word != 0 {
                let bit = word.trailing_zeros() as usize;
                return Some(word_idx * 64 + bit);
            }
            word_idx += 1;
            if word_idx >= WINDOW_WORDS {
                return None;
            }
            word = self.bitset[word_idx];
        }
    }

    /// Advance to the next set bit, refilling the window if needed.
    fn advance_to_next_set_bit(&mut self) {
        // Try the current window first
        if let Some(idx) = self.next_set_bit_in_window() {
            self.cursor = idx + 1;
            self.current = DocId::new(self.window_base + idx as u32);
            return;
        }

        // Window exhausted — refill from the next min doc.
        loop {
            let next_min = self.min_scorer_doc();
            if next_min == NO_MORE_DOCS {
                self.current = NO_MORE_DOCS;
                return;
            }
            self.refill_from(next_min.as_u32());
            // After refilling from min_scorer_doc, bit 0 of the new window
            // MUST be set (we picked the smallest current doc as base).
            // The loop is defensive against edge cases.
            if let Some(idx) = self.next_set_bit_in_window() {
                self.cursor = idx + 1;
                self.current = DocId::new(self.window_base + idx as u32);
                return;
            }
        }
    }
}

impl Scorer for BufferedUnionScorer {
    fn doc_id(&self) -> DocId {
        self.current
    }

    fn next(&mut self) -> DocId {
        if self.current == NO_MORE_DOCS {
            return NO_MORE_DOCS;
        }
        self.advance_to_next_set_bit();
        self.current
    }

    fn advance(&mut self, target: DocId) -> DocId {
        if self.current == NO_MORE_DOCS {
            return NO_MORE_DOCS;
        }
        if self.current >= target {
            return self.current;
        }
        let target_u32 = target.as_u32();
        let window_end = self.window_base.saturating_add(WINDOW_SIZE_U32);

        if target_u32 >= self.window_base && target_u32 < window_end {
            // In current window — scan forward from target
            self.cursor = (target_u32 - self.window_base) as usize;
            self.advance_to_next_set_bit();
        } else {
            // Outside current window — refill from target
            self.refill_from(target_u32);
            if let Some(idx) = self.next_set_bit_in_window() {
                self.cursor = idx + 1;
                self.current = DocId::new(self.window_base + idx as u32);
            } else {
                // Refilling from target may have skipped past — try min
                self.advance_to_next_set_bit();
            }
        }
        self.current
    }

    fn score(&mut self) -> f32 {
        self.boost
    }

    fn two_phase(&mut self) -> Option<&mut dyn TwoPhaseIterator> {
        None
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    /// Simple Vec-backed scorer for tests.
    struct VecScorer {
        docs: Vec<DocId>,
        pos: usize,
    }

    impl VecScorer {
        fn new(docs: Vec<u32>) -> Box<dyn Scorer> {
            Box::new(Self {
                docs: docs.into_iter().map(DocId::new).collect(),
                pos: 0,
            })
        }
    }

    impl Scorer for VecScorer {
        fn doc_id(&self) -> DocId {
            if self.pos < self.docs.len() {
                self.docs[self.pos]
            } else {
                NO_MORE_DOCS
            }
        }
        fn next(&mut self) -> DocId {
            if self.pos < self.docs.len() {
                self.pos += 1;
            }
            self.doc_id()
        }
        fn advance(&mut self, target: DocId) -> DocId {
            while self.pos < self.docs.len() && self.docs[self.pos] < target {
                self.pos += 1;
            }
            self.doc_id()
        }
        fn score(&mut self) -> f32 {
            1.0
        }
        fn two_phase(&mut self) -> Option<&mut dyn TwoPhaseIterator> {
            None
        }
    }

    fn collect(scorer: &mut dyn Scorer) -> Vec<u32> {
        let mut out = Vec::new();
        while scorer.doc_id() != NO_MORE_DOCS {
            out.push(scorer.doc_id().as_u32());
            scorer.next();
        }
        out
    }

    #[test]
    fn buffered_union_two_scorers() {
        let s1 = VecScorer::new(vec![0, 2, 4]);
        let s2 = VecScorer::new(vec![1, 2, 3]);
        let mut union = BufferedUnionScorer::new(vec![s1, s2]);
        assert_eq!(collect(&mut union), vec![0, 1, 2, 3, 4]);
    }

    #[test]
    fn buffered_union_three_scorers_overlap() {
        let s1 = VecScorer::new(vec![5]);
        let s2 = VecScorer::new(vec![5]);
        let s3 = VecScorer::new(vec![5]);
        let mut union = BufferedUnionScorer::new(vec![s1, s2, s3]);
        assert_eq!(collect(&mut union), vec![5]);
    }

    #[test]
    fn buffered_union_single_scorer() {
        let s = VecScorer::new(vec![0, 1, 2]);
        let mut union = BufferedUnionScorer::new(vec![s]);
        assert_eq!(collect(&mut union), vec![0, 1, 2]);
    }

    #[test]
    fn buffered_union_no_overlap() {
        let s1 = VecScorer::new(vec![0, 2, 4, 6]);
        let s2 = VecScorer::new(vec![1, 3, 5, 7]);
        let mut union = BufferedUnionScorer::new(vec![s1, s2]);
        assert_eq!(collect(&mut union), vec![0, 1, 2, 3, 4, 5, 6, 7]);
    }

    #[test]
    fn buffered_union_empty_scorer_filtered() {
        let s1 = VecScorer::new(vec![5]);
        let s2 = VecScorer::new(vec![]); // empty (already at NO_MORE_DOCS)
        let mut union = BufferedUnionScorer::new(vec![s1, s2]);
        assert_eq!(collect(&mut union), vec![5]);
    }

    #[test]
    fn buffered_union_advance_within_window() {
        let s1 = VecScorer::new(vec![0, 5, 10, 15]);
        let s2 = VecScorer::new(vec![1, 6, 11, 16]);
        let mut union = BufferedUnionScorer::new(vec![s1, s2]);
        // Initially at doc 0
        assert_eq!(union.doc_id(), DocId::new(0));
        // Advance to 8 — should land on 10
        assert_eq!(union.advance(DocId::new(8)), DocId::new(10));
        // Continue iteration
        assert_eq!(union.next(), DocId::new(11));
        assert_eq!(union.next(), DocId::new(15));
        assert_eq!(union.next(), DocId::new(16));
        assert_eq!(union.next(), NO_MORE_DOCS);
    }

    #[test]
    fn buffered_union_advance_past_window() {
        // Sparse docs with gaps larger than window size
        let s1 = VecScorer::new(vec![0, 5000, 10000]);
        let s2 = VecScorer::new(vec![100, 5100, 10100]);
        let mut union = BufferedUnionScorer::new(vec![s1, s2]);
        // Should iterate all 6 docs
        assert_eq!(union.doc_id(), DocId::new(0));
        assert_eq!(union.next(), DocId::new(100));
        assert_eq!(union.next(), DocId::new(5000));
        assert_eq!(union.next(), DocId::new(5100));
        assert_eq!(union.next(), DocId::new(10000));
        assert_eq!(union.next(), DocId::new(10100));
        assert_eq!(union.next(), NO_MORE_DOCS);
    }

    #[test]
    fn buffered_union_window_jump() {
        // Two scorers with docs in distant windows
        let s1 = VecScorer::new(vec![5, 100000]);
        let s2 = VecScorer::new(vec![10, 100010]);
        let mut union = BufferedUnionScorer::new(vec![s1, s2]);
        assert_eq!(collect(&mut union), vec![5, 10, 100000, 100010]);
    }

    #[test]
    fn buffered_union_advance_to_existing_doc() {
        let s1 = VecScorer::new(vec![0, 5, 10]);
        let s2 = VecScorer::new(vec![1, 6, 11]);
        let mut union = BufferedUnionScorer::new(vec![s1, s2]);
        // Advance to 5 — should land on 5
        assert_eq!(union.advance(DocId::new(5)), DocId::new(5));
        assert_eq!(union.next(), DocId::new(6));
    }

    #[test]
    fn buffered_union_advance_past_end() {
        let s1 = VecScorer::new(vec![0, 5]);
        let mut union = BufferedUnionScorer::new(vec![s1]);
        assert_eq!(union.advance(DocId::new(100)), NO_MORE_DOCS);
    }

    #[test]
    fn buffered_union_score_is_constant() {
        let s1 = VecScorer::new(vec![0, 1, 2]);
        let mut union = BufferedUnionScorer::new(vec![s1]);
        assert_eq!(union.score(), 1.0);
        union.next();
        assert_eq!(union.score(), 1.0);
        union.next();
        assert_eq!(union.score(), 1.0);
    }

    #[test]
    fn buffered_union_custom_boost() {
        let s1 = VecScorer::new(vec![0]);
        let mut union = BufferedUnionScorer::with_boost(vec![s1], 2.5);
        assert_eq!(union.score(), 2.5);
    }

    #[test]
    fn buffered_union_dense_window() {
        // 100 sequential docs all in one window
        let docs: Vec<u32> = (0..100).collect();
        let s = VecScorer::new(docs.clone());
        let mut union = BufferedUnionScorer::new(vec![s]);
        assert_eq!(collect(&mut union), docs);
    }

    #[test]
    fn buffered_union_window_boundary() {
        // Doc at exact window boundary
        let s1 = VecScorer::new(vec![2047, 2048, 2049]);
        let mut union = BufferedUnionScorer::new(vec![s1]);
        assert_eq!(collect(&mut union), vec![2047, 2048, 2049]);
    }

    #[test]
    fn buffered_union_many_scorers() {
        // 10 scorers, sparse, mostly distinct docs
        let scorers: Vec<Box<dyn Scorer>> = (0..10)
            .map(|i| VecScorer::new(vec![i * 100, i * 100 + 50, i * 100 + 99]))
            .collect();
        let mut union = BufferedUnionScorer::new(scorers);
        let docs = collect(&mut union);
        // Each scorer contributes 3 distinct docs → 30 total, all distinct
        assert_eq!(docs.len(), 30);
        // Should be sorted
        for w in docs.windows(2) {
            assert!(w[0] < w[1], "not sorted: {:?}", w);
        }
    }
}