daachorse 5.0.0

//! A match-candidate prefilter based on the SOG (Shift-Or with q-Grams) algorithm.
//!
//! The filter scans the haystack with a bit-parallel shift-or automaton over overlapping 2-grams
//! and reports positions where an occurrence of some pattern can start. Sections between
//! candidates are guaranteed to contain no occurrence, so the Aho-Corasick automaton can skip
//! them entirely while it stays in the root state. Candidates may be false positives; they are
//! simply verified by running the Aho-Corasick automaton as usual, so the filter never changes
//! match results.
//!
//! The algorithm is described in the following paper:
//!
//! > Leena Salmela, Jorma Tarhio, and Jari Kytöjoki.
//! > [Multipattern string matching with q-grams](https://doi.org/10.1145/1187436.1187438).
//! > *ACM Journal of Experimental Algorithmics*, 11, 2006.

use alloc::boxed::Box;
use alloc::vec::Vec;

use crate::errors::{DaachorseError, Result};
use crate::serializer::{Serializable, SerializableVec};

/// A SOG (Shift-Or with 2-Grams) prefilter.
#[derive(Clone, Eq, Hash, PartialEq)]
pub struct Prefilter {
    /// Maps a 2-gram to a bit vector whose `i`-th bit is 0 iff the 2-gram occurs at position `i`
    /// in the length-`window_len` prefix of some pattern.
    table: Box<[u8; Self::TABLE_LEN]>,
    /// The window length `m`: the length of the shortest pattern, capped at
    /// [`Prefilter::MAX_WINDOW_LEN`].
    window_len: u8,
    /// The bit at which a candidate is detected: `1 << (window_len - 2)`.
    hit_bit: u8,
}

impl Prefilter {
    /// The maximum window length: an 8-bit state supports up to `8 + 2 - 1` bytes because a window of
    /// `m` bytes yields `m - 1` overlapping 2-grams.
    pub const MAX_WINDOW_LEN: usize = 9;

    /// The minimum window length: a window must contain at least one 2-gram.
    pub const MIN_WINDOW_LEN: usize = 2;

    /// The number of 2-gram entries in the filter table.
    const TABLE_LEN: usize = 65536;

    /// The maximum estimated probability that a uniformly random position becomes a candidate.
    /// A table saturated with the 2-grams of too many patterns reports candidates almost
    /// everywhere and can hardly skip anything, so such a filter is not built. The estimate
    /// assumes random text and is thus optimistic on natural-language text; filters passing
    /// this check but not paying off in practice are disabled by [`PrefilterGate`] at run
    /// time.
    const MAX_EXPECTED_CANDIDATE_RATE: f64 = 1. / 16.;

    /// Returns the smallest position `>= pos` where an occurrence of some pattern can start, or
    /// `haystack.len()` if there is none. The result may be a false positive, but there is never
    /// an occurrence starting in `pos..result`.
    #[inline(always)]
    pub fn next_position(&self, haystack: &[u8], pos: usize) -> usize {
        let Some(&(mut prev)) = haystack.get(pos) else {
            return haystack.len();
        };
        let mut e = u8::MAX;
        for (i, &c) in haystack.iter().enumerate().skip(pos + 1) {
            e = (e << 1) | self.table[usize::from(prev) << 8 | usize::from(c)];
            if e & self.hit_bit == 0 {
                // Starting from an all-one state, the hit bit cannot become zero until
                // `window_len - 1` 2-grams have been consumed, so this subtraction cannot
                // underflow.
                return i + 1 - usize::from(self.window_len);
            }
            prev = c;
        }
        haystack.len()
    }

    /// Same as [`Prefilter::next_position`], but never returns a position in the middle of a
    /// UTF-8 character. Candidates starting in the middle of a character cannot be occurrences
    /// of character-wise patterns, so they are simply skipped.
    #[inline(always)]
    pub fn next_position_at_char_boundary(&self, haystack: &str, mut pos: usize) -> usize {
        loop {
            let candidate_pos = self.next_position(haystack.as_bytes(), pos);
            // is_char_boundary() returns true at the haystack end, so exhausted scans also
            // return here.
            if haystack.is_char_boundary(candidate_pos) {
                return candidate_pos;
            }
            pos = candidate_pos + 1;
        }
    }

    /// Returns the heap size of the filter table in bytes.
    #[allow(clippy::unused_self)]
    pub const fn heap_bytes(&self) -> usize {
        Self::TABLE_LEN
    }

    /// Estimates the probability that a uniformly random position in a random text becomes a
    /// candidate, as the product over all 2-gram positions of the fraction of 2-grams accepted
    /// there.
    #[allow(clippy::as_conversions)]
    fn expected_candidate_rate(&self) -> f64 {
        let mut zeros = [0u32; 8];
        for &bits in self.table.iter() {
            for (i, n) in zeros.iter_mut().enumerate() {
                *n += u32::from((bits >> i) & 1 == 0);
            }
        }
        let mut rate = 1.;
        for &n in &zeros[..usize::from(self.window_len - 1)] {
            rate *= f64::from(n) / Self::TABLE_LEN as f64;
        }
        rate
    }
}

/// An incremental builder of [`Prefilter`]. The automaton builders feed it while the patterns
/// stream through them, so that no copy of the pattern set is kept for the filter.
pub struct PrefilterBuilder {
    /// The table under construction; see [`Prefilter::table`] for the semantics.
    table: Vec<u8>,
    /// The length of the shortest pattern added so far, or `usize::MAX` if none has been added.
    min_len: usize,
}

impl PrefilterBuilder {
    pub fn new() -> Self {
        Self {
            table: vec![u8::MAX; Prefilter::TABLE_LEN],
            min_len: usize::MAX,
        }
    }

    /// Registers the 2-grams of the pattern's window, which is the pattern itself capped at
    /// [`Prefilter::MAX_WINDOW_LEN`] bytes.
    pub fn add(&mut self, pattern: &[u8]) {
        self.min_len = self.min_len.min(pattern.len());
        let window = &pattern[..pattern.len().min(Prefilter::MAX_WINDOW_LEN)];
        for (i, gram) in window.windows(2).enumerate() {
            self.table[usize::from(gram[0]) << 8 | usize::from(gram[1])] &= !(1 << i);
        }
    }

    /// Builds the prefilter, or returns `None` when prefiltering cannot pay off: no pattern was
    /// added, some pattern is shorter than [`Prefilter::MIN_WINDOW_LEN`], or the table is too
    /// dense to filter out enough positions.
    pub fn build(self) -> Option<Prefilter> {
        // usize::MAX means that no pattern has been added.
        if self.min_len < Prefilter::MIN_WINDOW_LEN || self.min_len == usize::MAX {
            return None;
        }
        let window_len = self.min_len.min(Prefilter::MAX_WINDOW_LEN);
        let prefilter = Prefilter {
            table: self.table.into_boxed_slice().try_into().unwrap(),
            window_len: window_len.try_into().unwrap(),
            hit_bit: 1 << (window_len - 2),
        };
        (prefilter.expected_candidate_rate() <= Prefilter::MAX_EXPECTED_CANDIDATE_RATE)
            .then_some(prefilter)
    }
}

impl Serializable for Prefilter {
    fn serialize_to_vec(&self, dst: &mut Vec<u8>) {
        self.window_len.serialize_to_vec(dst);
        self.table.serialize_to_vec(dst);
    }

    fn deserialize_from_slice(src: &[u8]) -> Result<(Self, &[u8])> {
        let (window_len, src) = u8::deserialize_from_slice(src)?;
        if !(Self::MIN_WINDOW_LEN..=Self::MAX_WINDOW_LEN).contains(&usize::from(window_len)) {
            return Err(DaachorseError::invalid_automaton());
        }
        let (table, rest) = Box::<[u8; Self::TABLE_LEN]>::deserialize_from_slice(src)?;
        Ok((
            Self {
                table,
                window_len,
                hit_bit: 1 << (window_len - 2),
            },
            rest,
        ))
    }

    fn serialized_bytes() -> usize {
        u8::serialized_bytes() + Self::TABLE_LEN
    }
}

/// A runtime gate that disables prefiltering when it does not pay off on the current haystack.
///
/// Each iterator owns a gate. Every filter run records how many bytes it allowed the automaton to
/// skip, and once the average gain within a measurement window falls below a threshold, the gate
/// closes and the iterator falls back to plain automaton scanning.
#[derive(Clone)]
pub struct PrefilterGate {
    calls_in_window: u32,
    gain_in_window: usize,
    enabled: bool,
}

impl PrefilterGate {
    /// The number of filter runs in one measurement window of [`PrefilterGate`].
    const GATE_WINDOW_CALLS: u32 = 64;

    /// The minimum number of bytes that filter runs in one measurement window must have skipped in
    /// total. Below this, filtering is likely slower than plain automaton scanning.
    const GATE_MIN_WINDOW_GAIN: usize = 512;

    pub(crate) const fn new() -> Self {
        Self {
            calls_in_window: 0,
            gain_in_window: 0,
            enabled: true,
        }
    }

    #[inline(always)]
    pub(crate) const fn is_enabled(&self) -> bool {
        self.enabled
    }

    /// Records that a filter run let the automaton skip `gain` bytes.
    #[inline(always)]
    pub(crate) fn record(&mut self, gain: usize) {
        self.calls_in_window += 1;
        self.gain_in_window += gain;
        if self.calls_in_window == Self::GATE_WINDOW_CALLS {
            if self.gain_in_window < Self::GATE_MIN_WINDOW_GAIN {
                self.enabled = false;
            }
            self.calls_in_window = 0;
            self.gain_in_window = 0;
        }
    }
}

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

    fn build(patterns: &[&[u8]]) -> Option<Prefilter> {
        let mut builder = PrefilterBuilder::new();
        patterns.iter().for_each(|pattern| builder.add(pattern));
        builder.build()
    }

    #[test]
    fn test_build_none_for_short_min_pattern() {
        assert!(build(&[b"aqua", b"a"]).is_none());
        assert!(build(&[b"aqua", b""]).is_none());
    }

    #[test]
    fn test_build_none_for_empty_pattern_set() {
        assert!(build(&[]).is_none());
    }

    #[test]
    fn test_build_none_when_table_saturated() {
        // All 2-byte patterns clear every table entry at position 0, making the expected
        // candidate rate 1.
        let mut builder = PrefilterBuilder::new();
        (0..=u16::MAX).for_each(|gram| builder.add(&gram.to_be_bytes()));
        assert!(builder.build().is_none());
    }

    #[test]
    fn test_long_pattern_windows_are_capped() {
        // Patterns longer than MAX_WINDOW_LEN must not overflow the 8-bit table entries.
        let long = b"undine".repeat(20);
        let pf = build(&[&long, b"neovenezia"]).unwrap();
        let haystack = [&[b'x'; 50][..], &long].concat();
        assert_eq!(pf.next_position(&haystack, 0), 50);
    }

    #[test]
    fn test_next_position_reports_occurrence() {
        let pf = build(&[b"aria", b"iris"]).unwrap();
        // The candidate is exactly the occurrence of "aria".
        assert_eq!(pf.next_position(b"xxxxariaxx", 0), 4);
        // Resuming after the candidate reaches the end without further candidates.
        assert_eq!(pf.next_position(b"xxxxariaxx", 5), 10);
        // Starting at the haystack end is allowed.
        assert_eq!(pf.next_position(b"xxxxariaxx", 10), 10);
    }

    #[test]
    fn test_next_position_without_occurrence() {
        let pf = build(&[b"aria", b"iris"]).unwrap();
        assert_eq!(pf.next_position(b"xxxxxxxxxx", 0), 10);
        // 2-grams of the patterns appearing at wrong window positions do not form a candidate.
        assert_eq!(pf.next_position(b"xxarxrixia", 0), 10);
    }

    #[test]
    fn test_next_position_may_report_false_positive() {
        let pf = build(&[b"aria", b"iris"]).unwrap();
        // "aris" chains 2-grams of both patterns ("ar" and "ri" from "aria", "is" from "iris")
        // at consistent window positions, so it is reported as a candidate although no pattern
        // occurs there. The caller verifies candidates on the automaton, so false positives are
        // harmless.
        assert_eq!(pf.next_position(b"xxarisxxxx", 0), 2);
    }

    #[test]
    fn test_min_window_pattern_builds() {
        // A 2-byte shortest pattern is exactly the minimum window and must build a filter that
        // reports its occurrences.
        let pf = build(&[b"ai", b"aika"]).unwrap();
        assert_eq!(pf.next_position(b"xxxxaixxxx", 0), 4);
        assert_eq!(pf.next_position(b"xxxxxxxxxx", 0), 10);
    }

    #[test]
    fn test_window_is_capped_at_exactly_nine_bytes() {
        // The 10-byte pattern is filtered through its 9-byte window "neovenezi".
        let pf = build(&[b"neovenezia"]).unwrap();
        // The full 9-byte window is a candidate even where the pattern's tenth byte differs.
        assert_eq!(pf.next_position(b"xxneovenezix", 0), 2);
        // A text sharing only the first 8 window bytes is not a candidate: the window must be
        // 9 bytes long, not shorter.
        assert_eq!(pf.next_position(b"xxneovenezxx", 0), 12);
    }

    #[test]
    fn test_char_boundary_reports_occurrence() {
        let pf = build(&["火星猫".as_bytes(), b"undine"]).unwrap();
        let haystack = "アリア社長は火星猫です";
        // "火星猫" occurs at byte position 18, which is a character boundary.
        assert_eq!(pf.next_position_at_char_boundary(haystack, 0), 18);
        assert_eq!(
            pf.next_position_at_char_boundary(haystack, 19),
            haystack.len()
        );
    }

    #[test]
    fn test_char_boundary_skips_mid_char_candidates() {
        // This pattern occurs in "星猫" only at byte position 1, in the middle of a character.
        let pf = build(&[b"\x98\x9f\xe7\x8c\xab"]).unwrap();
        let haystack = "星猫";
        assert_eq!(pf.next_position(haystack.as_bytes(), 0), 1);
        // A mid-character candidate cannot be an occurrence of a character-wise pattern and is
        // skipped.
        assert_eq!(
            pf.next_position_at_char_boundary(haystack, 0),
            haystack.len()
        );
    }

    #[test]
    fn test_char_boundary_candidate_right_after_mid_char_candidate() {
        // Byte-wise patterns can put a candidate in the middle of a character immediately
        // before a boundary candidate; skipping the former must not lose the latter.
        let pf = build(&[b"\x9f\xe7\x8c", "猫".as_bytes()]).unwrap();
        let haystack = "星猫"; // bytes: e6 98 9f | e7 8c ab
        assert_eq!(pf.next_position(haystack.as_bytes(), 0), 2);
        assert_eq!(pf.next_position_at_char_boundary(haystack, 0), 3);
    }

    #[test]
    fn test_next_position_never_goes_backward() {
        // The scan loops rely on next_position never returning a position before `pos`;
        // otherwise next_position_at_char_boundary could loop forever. The bit-parallel state
        // guarantees this: a hit needs `window_len - 1` shifts after the reset at `pos`.
        let pf = build(&[b"\x9f\xe7\x8c", "猫".as_bytes()]).unwrap();
        let haystack = "星猫".as_bytes();
        for pos in 0..=haystack.len() {
            assert!(pf.next_position(haystack, pos) >= pos, "pos {pos}");
        }
    }

    #[test]
    fn test_gate_closes_after_low_gain_window() {
        let mut gate = PrefilterGate::new();
        // 64 runs skipping 7 bytes each stay below the 512-byte window threshold, so the gate
        // closes exactly at the window end.
        for _ in 0..63 {
            gate.record(7);
            assert!(gate.is_enabled());
        }
        gate.record(7);
        assert!(!gate.is_enabled());
    }

    #[test]
    fn test_gate_stays_open_at_exact_threshold() {
        let mut gate = PrefilterGate::new();
        // 64 runs skipping 8 bytes each reach exactly the 512-byte threshold: not below it, so
        // the gate stays open.
        for _ in 0..64 {
            gate.record(8);
        }
        assert!(gate.is_enabled());
    }

    #[test]
    fn test_gate_evaluates_each_window_independently() {
        let mut gate = PrefilterGate::new();
        // A single huge skip keeps the first window far above the threshold.
        gate.record(100_000);
        for _ in 0..63 {
            gate.record(0);
        }
        assert!(gate.is_enabled());
        // The gain of the first window must not leak into the second: 64 gainless runs must
        // close the gate.
        for _ in 0..64 {
            gate.record(0);
        }
        assert!(!gate.is_enabled());
    }

    #[test]
    fn test_serialize_roundtrip() {
        let pf = build(&[b"aqua", b"aria"]).unwrap();
        let mut data = vec![];
        pf.serialize_to_vec(&mut data);
        assert_eq!(data.len(), Prefilter::serialized_bytes());
        data.push(42);
        let (other, rest) = Prefilter::deserialize_from_slice(&data).unwrap();
        assert_eq!(&[42], rest);
        assert!(pf == other);
    }

    #[test]
    fn test_deserialize_rejects_invalid_data() {
        let pf = build(&[b"gondola"]).unwrap();
        let mut data = vec![];
        pf.serialize_to_vec(&mut data);
        // The window length must stay within its valid range.
        let with_window_len = |window_len: u8| [&[window_len], &data[1..]].concat();
        assert!(Prefilter::deserialize_from_slice(&with_window_len(0)).is_err());
        assert!(Prefilter::deserialize_from_slice(&with_window_len(1)).is_err());
        assert!(Prefilter::deserialize_from_slice(&with_window_len(10)).is_err());
        assert!(Prefilter::deserialize_from_slice(&with_window_len(u8::MAX)).is_err());
        // A truncated table must also be rejected.
        assert!(Prefilter::deserialize_from_slice(&data[..data.len() - 1]).is_err());
    }
}