zoe 0.0.31

use crate::{
    DEFAULT_SIMD_LANES,
    private::Sealed,
    search::{
        RangeSearch, find_k_repeating_mapped_inner, inexact::fuzzy_substring_match_simd, k_repeating::find_k_repeating,
        position_by_byte, position_by_byte_inner, position_by_byte_mapped_inner,
    },
};
use std::{ops::Range, simd::prelude::*};

// Short-circuiting note:
// 0 < needle ≤ haystack IMPLIES 0 < haystack

/// Trait for searching byte substrings.
///
/// [`Nucleotides`](crate::data::types::nucleotides::Nucleotides),
/// [`AminoAcids`](crate::data::types::amino_acids::AminoAcids), and
/// [`QualityScores`](crate::data::types::phred::QualityScores) are printable
/// subsets of ASCII underneath the hood and therefore be efficiently searched
/// in this manner.
///
/// This trait is also compatible with [`RangeSearch`]. See [Restricting the
/// search range](crate::search#restricting-the-search-range) for more details.
pub trait ByteSubstring: Sealed {
    /// Returns `true` is the substring is found and `false` if not. Searches in
    /// the forward direction.
    #[must_use]
    fn contains_substring(&self, needle: impl AsRef<[u8]>) -> bool;

    /// Returns the substring's index range if found or [`None`] if not. Searches in
    /// the forward direction.
    #[must_use]
    fn find_substring(&self, needle: impl AsRef<[u8]>) -> Option<Range<usize>>;

    /// Finds the `needle` in the haystack using inexact matching up to the
    /// `DIFFERENCES_ALLOWED`.
    ///
    /// See [`fuzzy_substring_match_simd`] for more details.
    #[must_use]
    fn find_fuzzy_substring<const DIFFERENCES_ALLOWED: usize>(&self, needle: impl AsRef<[u8]>) -> Option<Range<usize>>;

    /// Finds the first occurrence of the `byte` in the haystack.
    #[must_use]
    fn find_byte(&self, byte: u8) -> Option<usize>;

    /// Find contiguous, repeating byte characters in a byte slice. See [`find_k_repeating`]
    /// for more details.
    #[must_use]
    fn find_repeating_byte(&self, needle: u8, size: usize) -> Option<Range<usize>>;

    /// Checks if the slice begins with some minimal number of consecutive byte
    /// repetitions and returns the largest contiguous range if so.
    #[must_use]
    fn find_repeating_at_start(&self, needle: u8, min_reps: usize) -> Option<Range<usize>>;

    /// Checks if the slice ends with some minimal number of consecutive byte
    /// repetitions and returns the largest contiguous range if so.
    #[must_use]
    fn find_repeating_at_end(&self, needle: u8, min_reps: usize) -> Option<Range<usize>>;
}

impl<T: AsRef<[u8]> + ?Sized + Sealed> ByteSubstring for T {
    #[inline]
    fn contains_substring(&self, needle: impl AsRef<[u8]>) -> bool {
        let haystack = self.as_ref();
        let needle = needle.as_ref();

        substring_match_simd::<{ DEFAULT_SIMD_LANES }>(haystack, needle).is_some()
    }

    #[inline]
    fn find_substring(&self, needle: impl AsRef<[u8]>) -> Option<Range<usize>> {
        let haystack = self.as_ref();
        let needle = needle.as_ref();
        substring_match_simd::<{ DEFAULT_SIMD_LANES }>(haystack, needle).map(|s| s..s + needle.len())
    }

    #[inline]
    fn find_fuzzy_substring<const DIFFERENCES_ALLOWED: usize>(&self, needle: impl AsRef<[u8]>) -> Option<Range<usize>> {
        let haystack = self.as_ref();
        let needle = needle.as_ref();

        fuzzy_substring_match_simd::<{ DEFAULT_SIMD_LANES }, DIFFERENCES_ALLOWED>(haystack, needle)
            .map(|s| s..s + needle.len())
    }

    #[inline]
    fn find_byte(&self, byte: u8) -> Option<usize> {
        let haystack = self.as_ref();

        position_by_byte::<{ DEFAULT_SIMD_LANES }>(haystack, byte)
    }

    #[inline]
    fn find_repeating_byte(&self, needle: u8, size: usize) -> Option<Range<usize>> {
        let haystack = self.as_ref();

        find_k_repeating::<{ DEFAULT_SIMD_LANES }>(haystack, needle, size).map(|s| s..s + size)
    }

    #[inline]
    fn find_repeating_at_start(&self, needle: u8, min_reps: usize) -> Option<Range<usize>> {
        let (head, tail) = self.as_ref().split_at_checked(min_reps)?;
        if head.iter().all(|b| *b == needle) {
            let offset = tail.iter().take_while(|b| **b == needle).count();
            Some(0..min_reps + offset)
        } else {
            None
        }
    }

    #[inline]
    fn find_repeating_at_end(&self, needle: u8, min_reps: usize) -> Option<Range<usize>> {
        let bytes = self.as_ref();
        if min_reps > bytes.len() {
            return None;
        }
        let (head, tail) = bytes.split_at(bytes.len() - min_reps);
        if tail.iter().all(|b| *b == needle) {
            let offset = head.iter().rev().take_while(|b| **b == needle).count();
            Some(head.len() - offset..bytes.len())
        } else {
            None
        }
    }
}

impl<Q: ?Sized> ByteSubstring for RangeSearch<'_, Q> {
    #[inline]
    fn contains_substring(&self, needle: impl AsRef<[u8]>) -> bool {
        self.slice.contains_substring(needle)
    }

    #[inline]
    fn find_substring(&self, needle: impl AsRef<[u8]>) -> Option<Range<usize>> {
        self.slice.find_substring(needle).map(|r| self.adjust_to_context(&r))
    }

    #[inline]
    fn find_fuzzy_substring<const DIFFERENCES_ALLOWED: usize>(&self, needle: impl AsRef<[u8]>) -> Option<Range<usize>> {
        self.slice
            .find_fuzzy_substring::<DIFFERENCES_ALLOWED>(needle)
            .map(|r| self.adjust_to_context(&r))
    }

    #[inline]
    fn find_byte(&self, byte: u8) -> Option<usize> {
        self.slice.find_byte(byte).map(|i| self.adjust_to_context(&i))
    }

    #[inline]
    fn find_repeating_byte(&self, needle: u8, size: usize) -> Option<Range<usize>> {
        self.slice
            .find_repeating_byte(needle, size)
            .map(|r| self.adjust_to_context(&r))
    }

    #[inline]
    fn find_repeating_at_start(&self, needle: u8, min_reps: usize) -> Option<Range<usize>> {
        self.slice
            .find_repeating_at_start(needle, min_reps)
            .map(|r| self.adjust_to_context(&r))
    }

    #[inline]
    fn find_repeating_at_end(&self, needle: u8, min_reps: usize) -> Option<Range<usize>> {
        self.slice
            .find_repeating_at_end(needle, min_reps)
            .map(|r| self.adjust_to_context(&r))
    }
}

/// Similar to [`ByteSubstring`] but requires the input byte string to be
/// mutable.
pub trait ByteSubstringMut: Sealed {
    /// If the substring is found, the first instance is removed and index range
    /// of the removed substring is returned, otherwise [`None`]. Searches in
    /// the forward direction.
    fn remove_first_substring(&mut self, needle: impl AsRef<[u8]>) -> Option<Range<usize>>;

    /// Replace a single byte of the stored sequence. See the [`Recode`] trait
    /// for a more wholistic approach.
    ///
    /// [`Recode`]: crate::data::Recode
    fn replace_all_bytes(&mut self, needle: u8, replacement: u8);
}

impl<T: AsMut<Vec<u8>> + AsRef<[u8]> + Sealed> ByteSubstringMut for T {
    #[inline]
    fn remove_first_substring(&mut self, needle: impl AsRef<[u8]>) -> Option<Range<usize>> {
        if let Some(r) = self.find_substring(needle) {
            self.as_mut().drain(r.clone());
            Some(r)
        } else {
            None
        }
    }

    #[inline]
    fn replace_all_bytes(&mut self, needle: u8, replacement: u8) {
        crate::search::replace_all_bytes(self.as_mut(), needle, replacement);
    }
}

/// Finds the `needle` byte substring in the `haystack`, returning the starting
/// index or [`None`] otherwise.
///
/// ## Limitations
///
/// This is a naïve exact match implementation and should only be used for small
/// byte strings.
#[inline]
#[must_use]
pub fn substring_match(haystack: &[u8], needle: &[u8]) -> Option<usize> {
    if needle.len() > haystack.len() || needle.is_empty() {
        return None;
    }

    for (i, w) in haystack.windows(needle.len()).enumerate() {
        if needle == w {
            return Some(i);
        }
    }

    None
}

/// Returns the starting index of the matched substring or [`None`] otherwise.
///
/// ## Parameters
///
/// `N`: The number of SIMD lanes to use for the search. This must be greater
/// than 2.
///
/// ## Limitations
///
/// The current version is not optimized for larger needles (> 30 bp).
///
/// ## Citation
///
/// 1. Muła, Wojciech (2018). "SIMD-friendly algorithms for substring
///    searching." Available at:
///    <http://0x80.pl/notesen/2016-11-28-simd-strfind.html#algorithm-1-generic-simd>.
///    Accessed September 3, 2024.
#[inline]
#[must_use]
#[cfg_attr(feature = "multiversion", multiversion::multiversion(targets = "simd"))]
pub fn substring_match_simd<const N: usize>(haystack: &[u8], needle: &[u8]) -> Option<usize> {
    if needle.len() > haystack.len() || needle.is_empty() {
        return None;
    }

    let first = needle[0];
    if needle.len() == 1 {
        return position_by_byte_inner::<N>(haystack, first);
    }

    let Some((n2_offset, last)) = needle.iter().copied().enumerate().rev().find(|(_, b)| *b != first) else {
        return find_k_repeating::<N>(haystack, first, needle.len());
    };

    let n1 = Simd::from_array([first; N]);
    let n2 = Simd::from_array([last; N]);

    // In order to verify the needle, we need to subtract it off. However, the
    // last character in the vector counts.
    let chunks1 = haystack[..=(haystack.len() - needle.len())]
        .as_chunks::<N>()
        .0
        .iter()
        .copied()
        .map(Simd::from_array);
    let chunks2 = haystack[n2_offset..].as_chunks::<N>().0.iter().copied().map(Simd::from_array);
    let z = std::iter::zip(chunks1, chunks2);

    let mut i = 0;
    for (c1, c2) in z {
        let f1 = n1.simd_eq(c1);
        let f2 = n2.simd_eq(c2);

        let mut m = (f1 & f2).to_bitmask();

        while m > 0 {
            let bit_position = m.trailing_zeros() as usize;
            let candidate_index = i + bit_position;

            if &haystack[candidate_index..candidate_index + needle.len()] == needle {
                return Some(candidate_index);
            }
            m &= m - 1;
        }
        i += N;
    }

    if N <= 8 {
        substring_match(&haystack[i..], needle).map(|j| i + j)
    } else {
        substring_match_simd::<8>(&haystack[i..], needle).map(|j| i + j)
    }
}

/// Similar to [`substring_match_simd`] but takes a transformation of the
/// haystack before matching.
///
/// A transformation mapping for both bytes and SIMD vectors is required in
/// order to verify potential matches.
///
/// ## Parameters
///
/// - `N`: The number of SIMD lanes to use for the search. This must be greater
///   than 2.
/// - `S`: The type of the `simd_transform` closure (this can often be inferred
///   by the compiler).
/// - `B`: The type of the `byte_transform` closure (this can often be inferred
///   by the compiler).
///
/// ## Limitations
///
/// The current version is not optimized for larger needles (> 30 bp).
#[inline]
#[must_use]
#[cfg_attr(feature = "multiversion", multiversion::multiversion(targets = "simd"))]
pub fn find_mapped_match_simd<const N: usize, S, B>(
    haystack: &[u8], needle: &[u8], simd_transform: S, byte_transform: B,
) -> Option<usize>
where
    S: Fn(Simd<u8, N>) -> Simd<u8, N>,
    B: Fn(u8) -> u8, {
    if needle.len() > haystack.len() || needle.is_empty() {
        return None;
    }

    let first = needle[0];
    if needle.len() == 1 {
        return position_by_byte_mapped_inner(haystack, first, simd_transform, byte_transform);
    }

    let Some((n2_offset, last)) = needle.iter().copied().enumerate().rev().find(|(_, b)| *b != first) else {
        return find_k_repeating_mapped_inner::<N, S, B>(haystack, first, needle.len(), simd_transform, byte_transform);
    };

    let n1 = Simd::from_array([first; N]);
    let n2 = Simd::from_array([last; N]);
    let candidate_start_limit = haystack.len() - needle.len() + 1;
    let matches_candidate = |candidate_index| {
        std::iter::zip(&haystack[candidate_index..candidate_index + needle.len()], needle)
            .all(|(h, n)| *n == byte_transform(*h))
    };

    let chunks1 = haystack[..candidate_start_limit]
        .as_chunks::<N>()
        .0
        .iter()
        .copied()
        .map(Simd::from_array)
        .map(&simd_transform);
    let chunks2 = haystack[n2_offset..]
        .as_chunks::<N>()
        .0
        .iter()
        .copied()
        .map(Simd::from_array)
        .map(&simd_transform);
    let z = std::iter::zip(chunks1, chunks2);

    let mut i = 0;
    for (c1, c2) in z {
        let f1 = n1.simd_eq(c1);
        let f2 = n2.simd_eq(c2);

        let mut m = (f1 & f2).to_bitmask();

        while m > 0 {
            let bit_position = m.trailing_zeros() as usize;
            let candidate_index = i + bit_position;

            if matches_candidate(candidate_index) {
                return Some(candidate_index);
            }
            m &= m - 1;
        }
        i += N;
    }

    let remaining_candidates = candidate_start_limit - i;
    if remaining_candidates == 0 {
        return None;
    }

    // Pack the leftover candidate starts into one padded SIMD block so the
    // tail still uses the same fast prefilter as the main loop.
    let mut tail1 = [0u8; N];
    tail1[..remaining_candidates].copy_from_slice(&haystack[i..i + remaining_candidates]);
    let mut tail2 = [0u8; N];
    tail2[..remaining_candidates].copy_from_slice(&haystack[i + n2_offset..i + n2_offset + remaining_candidates]);

    let active_lanes = Mask::from_array(std::array::from_fn(|lane| lane < remaining_candidates));
    let f1 = n1.simd_eq(simd_transform(Simd::from_array(tail1)));
    let f2 = n2.simd_eq(simd_transform(Simd::from_array(tail2)));
    let mut m = (f1 & f2 & active_lanes).to_bitmask();

    while m > 0 {
        let bit_position = m.trailing_zeros() as usize;
        let candidate_index = i + bit_position;

        if matches_candidate(candidate_index) {
            return Some(candidate_index);
        }
        m &= m - 1;
    }

    None
}

#[cfg(test)]
mod test {
    use super::*;
    use crate::{search::ToRangeSearch, simd::SimdByteFunctions};

    static PAD: &[u8; 150] = &[b'a'; 150];
    static NEEDLE: &[u8; 5] = b"hello";

    #[test]
    fn substring_match_units() {
        let mut haystack = *PAD;
        assert_eq!(None, substring_match(&haystack, NEEDLE));
        for start in 0..haystack.len() - NEEDLE.len() {
            haystack = *PAD;
            haystack[start..start + NEEDLE.len()].copy_from_slice(NEEDLE);
            assert_eq!(Some(start), substring_match(&haystack, NEEDLE));
        }
    }

    #[test]
    fn substring_match_simd_units() {
        let mut haystack = *PAD;
        assert_eq!(None, substring_match_simd::<8>(&haystack, NEEDLE));
        for start in 0..haystack.len() - NEEDLE.len() {
            haystack = *PAD;
            haystack[start..start + NEEDLE.len()].copy_from_slice(NEEDLE);
            assert_eq!(Some(start), substring_match_simd::<32>(&haystack, NEEDLE));
        }
    }

    #[test]
    fn substring_match_regressions() {
        let data = [
            (b"aaaaabaadaa".to_vec(), b"baabbbb".to_vec()),
            (b"dcxxxaxxxx".to_vec(), b"axx".to_vec()),
        ];

        for (haystack, needle) in data {
            assert_eq!(
                substring_match(&haystack, &needle),
                substring_match_simd::<16>(&haystack, &needle)
            );
        }
    }

    #[test]
    fn mapped_match() {
        fn simd_map<const N: usize>(s: Simd<u8, N>) -> Simd<u8, N> {
            let mut xformed = s.to_ascii_uppercase();
            xformed.if_value_then_replace(b'U', b'T');
            xformed
        }

        fn byte_map(b: u8) -> u8 {
            let mut xformed = b.to_ascii_uppercase();
            if xformed == b'U' {
                xformed = b'T';
            }
            xformed
        }

        let haystack = b"AAAAAAAAAAAAAAAAAAAAAAAAaugAAAAAAAAAA";
        assert_eq!(
            Some(24),
            find_mapped_match_simd::<8, _, _>(haystack, b"ATG", simd_map, byte_map)
        );

        let haystack = b"CCCCCCCCCCCCCCCCCCCCCCCCuCCCCCCCCCC";
        assert_eq!(
            Some(24),
            find_mapped_match_simd::<8, _, _>(haystack, b"T", simd_map, byte_map)
        );
    }

    #[test]
    fn starts_ends_with_repeating() {
        let b = b"GGGGGGGGGGGGAGCAAGCACAAAACAAAAATCCATGTAAGGAATAGGGGGGGGGGGGGG";
        assert_eq!(b.find_repeating_at_start(b'G', 10), Some(0..12));
        assert_eq!(b.find_repeating_at_end(b'G', 10), Some(46..60));

        let b = b"GGGCAAGGGGGAATAGGGGG";
        assert_eq!(b.find_repeating_at_start(b'G', 3), Some(0..3));
        assert_eq!(b.find_repeating_at_end(b'G', 5), Some(15..20));
        assert_eq!(b.find_repeating_at_start(b'G', 2), Some(0..3));
        assert_eq!(b.find_repeating_at_end(b'G', 4), Some(15..20));
        assert_eq!(b.find_repeating_at_start(b'G', 4), None);
        assert_eq!(b.find_repeating_at_end(b'G', 6), None);

        let b = b"GGGGGGGGGGGGGGG";
        assert_eq!(b.find_repeating_at_start(b'G', 3), Some(0..15));
        assert_eq!(b.find_repeating_at_end(b'G', 5), Some(0..15));

        let b = b"AGGGGGGGGGGGGGA";
        assert_eq!(b.find_repeating_at_start(b'G', 3), None);
        assert_eq!(b.find_repeating_at_end(b'G', 5), None);

        let b = b"G";
        assert_eq!(b.find_repeating_at_end(b'G', 2), None);
        assert_eq!(b.find_repeating_at_start(b'G', 2), None);
    }

    #[test]
    fn search_byte() {
        let seq = b"GGGAAGCATCACGTATCGA";
        assert_eq!(seq.find_byte(b'G'), Some(0));
        assert_eq!(seq.find_byte(b'A'), Some(3));
        assert_eq!(seq.find_byte(b'C'), Some(6));
        assert_eq!(seq.find_byte(b'T'), Some(8));
        assert_eq!(seq.find_byte(b'N'), None);
    }

    #[test]
    fn range_search_byte() {
        let seq = b"GGGAAGCATCACGTATCGA";

        assert_eq!(seq.search_in(0..).find_byte(b'G'), Some(0));
        assert_eq!(seq.search_in(1..).find_byte(b'G'), Some(1));
        assert_eq!(seq.search_in(2..).find_byte(b'G'), Some(2));
        for i in 3..=5 {
            assert_eq!(seq.search_in(i..).find_byte(b'G'), Some(5));
        }
        for i in 6..=12 {
            assert_eq!(seq.search_in(i..).find_byte(b'G'), Some(12));
        }
        for i in 13..=17 {
            assert_eq!(seq.search_in(i..).find_byte(b'G'), Some(17));
        }
        assert_eq!(seq.search_in(18..).find_byte(b'G'), None);
        assert_eq!(seq.search_in(19..).find_byte(b'G'), None);

        for i in 0..=3 {
            assert_eq!(seq.search_in_first(i).find_byte(b'A'), None);
        }
        for i in 4..=19 {
            assert_eq!(seq.search_in_first(i).find_byte(b'A'), Some(3));
        }

        for i in 0..=3 {
            assert_eq!(seq.search_in_last(i).find_byte(b'T'), None);
        }
        for i in 4..=5 {
            assert_eq!(seq.search_in_last(i).find_byte(b'T'), Some(15));
        }
        for i in 6..=10 {
            assert_eq!(seq.search_in_last(i).find_byte(b'T'), Some(13));
        }
        for i in 11..=19 {
            assert_eq!(seq.search_in_last(i).find_byte(b'T'), Some(8));
        }
    }
}