quick_search/crates/quicksearch-core/src/search/fuzzy.rs

//! Approximate substring matching for the fuzzy cascade stages.
//!
//! Bitap (shift-and with errors, Wu–Manber): finds occurrences of a
//! pattern *within* a haystack with at most `k` Levenshtein edits
//! (insertion / deletion / substitution). Substring semantics are the
//! right fit for search-as-you-type — users type fragments, not whole
//! filenames — and the u64 bit-parallel update costs O(k) word ops per
//! haystack byte with zero allocations.
//!
//! Callers fold both sides to ASCII lowercase first (the pipeline-wide
//! convention). Patterns are limited to 64 bytes by the machine word; the
//! cascade skips fuzzy stages for longer terms.

/// Registers the automaton needs: one per error count `0..=k`.
///
/// `k` is bounded by [`edit_budget`]'s one-edit-per-three-characters ladder
/// against a pattern the machine word caps at 64 bytes, so it never exceeds 21
/// however hostile `fuzzy_max_edits` is (pinned by
/// `edit_budget_stays_within_the_bitap_word_size`). [`Bitap::new`] rejects
/// anything larger, so the array is always big enough.
const MAX_REGISTERS: usize = 22;

pub struct Bitap {
    /// `masks[c]` has bit `i` set iff `pattern[i] == c`.
    masks: [u64; 256],
    /// Pattern length in bytes (1..=64).
    len: usize,
    /// Maximum edit distance.
    k: usize,
}

impl Bitap {
    /// `None` when the pattern is empty, longer than 64 bytes, or the edit
    /// budget does not fit the registers (`reset` shifts by `k`, and the
    /// register array holds [`MAX_REGISTERS`]).
    pub fn new(pattern: &[u8], k: usize) -> Option<Bitap> {
        if pattern.is_empty() || pattern.len() > 64 || k >= MAX_REGISTERS {
            return None;
        }
        let mut masks = [0u64; 256];
        for (i, &b) in pattern.iter().enumerate() {
            masks[b as usize] |= 1u64 << i;
        }
        Some(Bitap {
            masks,
            len: pattern.len(),
            k,
        })
    }

    /// Reset the per-distance state registers. Bit `i` of `r[d]` set means "a
    /// match of pattern[..=i] with ≤ d errors ends at the current text
    /// position". With d errors the first d pattern bytes can be deleted
    /// before any text is read, hence the pre-set low bits.
    ///
    /// Writes into a caller-owned array rather than returning a `Vec`: the
    /// fuzzy passes call this once per scanned row — millions of them on a
    /// whole-table sweep — and a heap allocation per row is pure overhead for
    /// at most [`MAX_REGISTERS`] words.
    fn reset(&self, r: &mut [u64; MAX_REGISTERS]) {
        for (d, reg) in r.iter_mut().enumerate().take(self.k + 1) {
            *reg = if d == 0 { 0 } else { (1u64 << d) - 1 };
        }
    }

    /// Advance all registers by one haystack byte. Returns the smallest
    /// error count d for which the full pattern just matched, if any.
    #[inline]
    fn step(&self, r: &mut [u64], byte: u8) -> Option<usize> {
        let mask = self.masks[byte as usize];
        let done = 1u64 << (self.len - 1);
        let mut hit = None;
        let mut prev_old = r[0]; // R_old[d-1] for the d-th iteration
        // d = 0: exact prefix extension only.
        r[0] = ((r[0] << 1) | 1) & mask;
        if r[0] & done != 0 {
            hit = Some(0);
        }
        for d in 1..=self.k {
            let old = r[d];
            r[d] = (((old << 1) | 1) & mask)      // extend a ≤d-error state
                | prev_old                          // insertion in text
                | (prev_old << 1)                   // substitution
                | ((r[d - 1] << 1) | 1);            // deletion (pattern byte skipped)
            prev_old = old;
            if hit.is_none() && r[d] & done != 0 {
                hit = Some(d);
            }
        }
        hit
    }

    /// Minimum edit distance (≤ k) of any occurrence of the pattern in
    /// `hay`, or `None` if nothing matches within k edits.
    pub fn best_distance(&self, hay: &[u8]) -> Option<usize> {
        self.best_distance_and_first(hay).map(|(d, _)| d)
    }

    /// [`best_distance`](Self::best_distance) plus the first match's
    /// approximate byte range, from one sweep.
    ///
    /// The filename pass needs both — the distance to rank by, the range to
    /// mark in the snippet — and taking them separately meant a second full
    /// scan of the same buffer for every hit. The range carries the same
    /// caveat as [`count_and_first`](Self::count_and_first): it assumes a
    /// pattern-length match, so edits can shift the true start by up to `k`.
    pub fn best_distance_and_first(&self, hay: &[u8]) -> Option<(usize, (usize, usize))> {
        let mut r = [0u64; MAX_REGISTERS];
        self.reset(&mut r);
        let mut best: Option<(usize, (usize, usize))> = None;
        for (i, &b) in hay.iter().enumerate() {
            if let Some(d) = self.step(&mut r, b) {
                let end = i + 1;
                let range = (end.saturating_sub(self.len), end);
                if d == 0 {
                    return Some((0, range));
                }
                if best.is_none_or(|(cur, _)| d < cur) {
                    best = Some((d, range));
                }
            }
        }
        best
    }

    /// Count non-overlapping occurrences (at ≤ k edits) and report the
    /// first match's approximate byte range in `hay`. After each hit the
    /// automaton resets, so an exact match followed by trailing bytes
    /// counts once, and overlapping suffix matches don't inflate counts.
    /// The reported range assumes pattern-length matches — edits can shift
    /// the true start by up to k bytes, which is fine for snippet windows.
    pub fn count_and_first(&self, hay: &[u8]) -> (usize, Option<(usize, usize)>) {
        let mut r = [0u64; MAX_REGISTERS];
        self.reset(&mut r);
        let mut count = 0usize;
        let mut first: Option<(usize, usize)> = None;
        for (i, &b) in hay.iter().enumerate() {
            if self.step(&mut r, b).is_some() {
                count += 1;
                if first.is_none() {
                    let end = i + 1;
                    first = Some((end.saturating_sub(self.len), end));
                }
                self.reset(&mut r);
            }
        }
        (count, first)
    }
}

/// The cascade's edit-distance budget for a folded term: one edit per
/// three characters, capped by `[search].fuzzy_max_edits`. Terms outside
/// 3..=64 bytes skip the fuzzy stages entirely (< 3 is noise, > 64 exceeds
/// the word size), and a cap of 0 disables them everywhere.
///
/// At the default cap of 2 this is the historic ladder: 3–5 bytes get one
/// edit, 6–64 get two.
pub fn edit_budget(term_len: usize, max_edits: usize) -> Option<usize> {
    if !(3..=64).contains(&term_len) || max_edits == 0 {
        return None;
    }
    Some((term_len / 3).min(max_edits))
}

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

    fn best(pattern: &str, hay: &str, k: usize) -> Option<usize> {
        Bitap::new(pattern.as_bytes(), k)
            .unwrap()
            .best_distance(hay.as_bytes())
    }

    #[test]
    fn exact_substring_is_distance_zero() {
        assert_eq!(best("hello", "say hello world", 2), Some(0));
        assert_eq!(best("hello", "hello", 0), Some(0));
    }

    #[test]
    fn single_edits_are_distance_one() {
        assert_eq!(best("hello", "xx hxllo xx", 2), Some(1), "substitution");
        assert_eq!(best("hello", "xx helo xx", 2), Some(1), "deletion");
        assert_eq!(best("hello", "xx heXllo xx", 2), Some(1), "insertion");
    }

    #[test]
    fn two_edits() {
        assert_eq!(best("hello", "xx hxlo xx", 2), Some(2));
        assert_eq!(best("hello", "xx ho xx", 2), None, "3 edits > k");
    }

    #[test]
    fn no_match_within_budget() {
        assert_eq!(best("hello", "completely different", 1), None);
        assert_eq!(best("abc", "", 1), None);
    }

    #[test]
    fn k_zero_is_exact_search() {
        assert_eq!(best("abc", "xxabcxx", 0), Some(0));
        assert_eq!(best("abc", "xxabxcx", 0), None);
    }

    #[test]
    fn pattern_length_limits() {
        assert!(Bitap::new(b"", 1).is_none());
        assert!(Bitap::new(&[b'a'; 65], 1).is_none());
        assert!(Bitap::new(&[b'a'; 64], 1).is_some());
    }

    #[test]
    fn oversized_k_is_rejected_not_shifted() {
        // `reset` shifts by k and writes k+1 registers, so both the u64 shift
        // and the fixed array have to be respected.
        assert!(Bitap::new(b"abc", MAX_REGISTERS).is_none());
        assert!(Bitap::new(b"abc", 64).is_none());
        assert!(Bitap::new(b"abc", usize::MAX).is_none());
        assert!(Bitap::new(b"abc", MAX_REGISTERS - 1).is_some());
    }

    #[test]
    fn count_non_overlapping() {
        let b = Bitap::new(b"ab", 0).unwrap();
        let (count, first) = b.count_and_first(b"ab ab ab");
        assert_eq!(count, 3);
        assert_eq!(first, Some((0, 2)));

        // "aaaa" contains "aaa" once non-overlapping.
        let b = Bitap::new(b"aaa", 0).unwrap();
        let (count, _) = b.count_and_first(b"aaaa");
        assert_eq!(count, 1);
    }

    #[test]
    fn count_fuzzy_and_range_sane() {
        let b = Bitap::new(b"hello", 1).unwrap();
        let hay = b"say helo and hxllo again";
        let (count, first) = b.count_and_first(hay);
        assert_eq!(count, 2);
        let (s, e) = first.unwrap();
        assert!(s < e && e <= hay.len());
        let window = &hay[s..e];
        assert!(
            std::str::from_utf8(window).unwrap().contains("hel"),
            "first range should cover the first hit, got {:?}",
            std::str::from_utf8(window)
        );
    }

    #[test]
    fn edit_budget_default_cap_is_the_historic_ladder() {
        assert_eq!(edit_budget(0, 2), None);
        assert_eq!(edit_budget(2, 2), None);
        assert_eq!(edit_budget(3, 2), Some(1));
        assert_eq!(edit_budget(5, 2), Some(1));
        assert_eq!(edit_budget(6, 2), Some(2));
        assert_eq!(edit_budget(64, 2), Some(2));
        assert_eq!(edit_budget(65, 2), None);
        assert_eq!(edit_budget(usize::MAX, 2), None);
    }

    #[test]
    fn edit_budget_scales_with_length_up_to_the_cap() {
        assert_eq!(edit_budget(3, 4), Some(1));
        assert_eq!(edit_budget(6, 4), Some(2));
        assert_eq!(edit_budget(9, 4), Some(3));
        assert_eq!(edit_budget(12, 4), Some(4));
        assert_eq!(edit_budget(64, 4), Some(4), "cap wins over length");
    }

    #[test]
    fn edit_budget_cap_of_one_stays_strict() {
        for len in 3..=64 {
            assert_eq!(edit_budget(len, 1), Some(1));
        }
    }

    #[test]
    fn edit_budget_zero_disables_fuzzy() {
        for len in 0..=70 {
            assert_eq!(edit_budget(len, 0), None);
        }
    }

    /// Even a hostile config value can't produce a k the bitap rejects:
    /// the length ladder caps it at 21 for the longest legal term.
    #[test]
    fn edit_budget_stays_within_the_bitap_word_size() {
        for len in 3..=64 {
            let k = edit_budget(len, usize::MAX).unwrap();
            assert!(k <= 21, "len={} gave k={}", len, k);
            assert!(Bitap::new(&vec![b'a'; len], k).is_some());
        }
    }

    /// Brute-force oracle: minimum Levenshtein distance between `pattern`
    /// and any substring of `hay`, capped at k.
    fn oracle(pattern: &[u8], hay: &[u8], k: usize) -> Option<usize> {
        // An occurrence must end at some text position; empty text has
        // none. Without this, k >= pattern-length "matches" empty text by
        // deleting every pattern byte — a degenerate non-occurrence the
        // automaton rightly never reports. (Production keeps k < len via
        // the len/3 budget, so only the oracle ever saw this edge.)
        if hay.is_empty() {
            return None;
        }
        // Standard DP where row 0 is all zeros (match can start anywhere).
        let m = pattern.len();
        let mut prev: Vec<usize> = vec![0; hay.len() + 1];
        let mut cur = vec![0; hay.len() + 1];
        let mut best = usize::MAX;
        // dp[i][j] = min edits to match pattern[..i] ending at hay[..j]
        for i in 1..=m {
            cur[0] = i;
            for j in 1..=hay.len() {
                let cost = if pattern[i - 1] == hay[j - 1] { 0 } else { 1 };
                cur[j] = (prev[j - 1] + cost).min(prev[j] + 1).min(cur[j - 1] + 1);
            }
            std::mem::swap(&mut prev, &mut cur);
        }
        for j in 0..=hay.len() {
            best = best.min(prev[j]);
        }
        if best <= k {
            Some(best)
        } else {
            None
        }
    }

    #[test]
    fn matches_brute_force_oracle() {
        // Deterministic LCG so the test is reproducible.
        let mut seed: u64 = 0x2545F4914F6CDD1D;
        let mut rng = move || {
            seed = seed.wrapping_mul(6364136223846793005).wrapping_add(1442695040888963407);
            (seed >> 33) as usize
        };
        let alphabet = b"abcx";
        for _ in 0..500 {
            let plen = 3 + rng() % 6;
            let hlen = rng() % 20;
            let pattern: Vec<u8> = (0..plen).map(|_| alphabet[rng() % 4]).collect();
            let hay: Vec<u8> = (0..hlen).map(|_| alphabet[rng() % 4]).collect();
            for k in 0..=4 {
                let got = Bitap::new(&pattern, k).unwrap().best_distance(&hay);
                let want = oracle(&pattern, &hay, k);
                assert_eq!(
                    got,
                    want,
                    "pattern={:?} hay={:?} k={}",
                    std::str::from_utf8(&pattern),
                    std::str::from_utf8(&hay),
                    k
                );
            }
        }
    }
}