/*!
This is the Teddy searcher, but ported to AVX2.

See the module comments in the SSSE3 Teddy searcher for a more in depth
explanation of how this algorithm works. For the most part, this port is
basically the same as the SSSE3 version, but using 256-bit vectors instead of
128-bit vectors, which increases throughput.
*/

use std::cmp;

use aho_corasick::{Automaton, AcAutomaton, FullAcAutomaton};
use syntax::hir::literal::Literals;

use vector::avx2::{AVX2VectorBuilder, u8x32};

/// Corresponds to the number of bytes read at a time in the haystack.
const BLOCK_SIZE: usize = 32;

/// Match reports match information.
#[derive(Debug, Clone)]
pub struct Match {
    /// The index of the pattern that matched. The index is in correspondence
    /// with the order of the patterns given at construction.
    pub pat: usize,
    /// The start byte offset of the match.
    pub start: usize,
    /// The end byte offset of the match. This is always `start + pat.len()`.
    pub end: usize,
}

/// A SIMD accelerated multi substring searcher.
#[derive(Debug, Clone)]
pub struct Teddy {
    /// A builder for AVX2 empowered vectors.
    vb: AVX2VectorBuilder,
    /// A list of substrings to match.
    pats: Vec<Vec<u8>>,
    /// An Aho-Corasick automaton of the patterns. We use this when we need to
    /// search pieces smaller than the Teddy block size.
    ac: FullAcAutomaton<Vec<u8>>,
    /// A set of 8 buckets. Each bucket corresponds to a single member of a
    /// bitset. A bucket contains zero or more substrings. This is useful
    /// when the number of substrings exceeds 8, since our bitsets cannot have
    /// more than 8 members.
    buckets: Vec<Vec<usize>>,
    /// Our set of masks. There's one mask for each byte in the fingerprint.
    masks: Masks,
}

impl Teddy {
    /// Returns true if and only if Teddy is supported on this platform.
    ///
    /// If this returns `false`, then `Teddy::new(...)` is guaranteed to
    /// return `None`.
    pub fn available() -> bool {
        AVX2VectorBuilder::new().is_some()
    }

    /// Create a new `Teddy` multi substring matcher.
    ///
    /// If a `Teddy` matcher could not be created (e.g., `pats` is empty or has
    /// an empty substring), then `None` is returned.
    pub fn new(pats: &Literals) -> Option<Teddy> {
        let vb = match AVX2VectorBuilder::new() {
            None => return None,
            Some(vb) => vb,
        };
        if !Teddy::available() {
            return None;
        }

        let pats: Vec<_> = pats.literals().iter().map(|p|p.to_vec()).collect();
        let min_len = pats.iter().map(|p| p.len()).min().unwrap_or(0);
        // Don't allow any empty patterns and require that we have at
        // least one pattern.
        if min_len < 1 {
            return None;
        }
        // Pick the largest mask possible, but no larger than 3.
        let nmasks = cmp::min(3, min_len);
        let mut masks = Masks::new(vb, nmasks);
        let mut buckets = vec![vec![]; 8];
        // Assign a substring to each bucket, and add the bucket's bitfield to
        // the appropriate position in the mask.
        for (pati, pat) in pats.iter().enumerate() {
            let bucket = pati % 8;
            buckets[bucket].push(pati);
            masks.add(bucket as u8, pat);
        }
        Some(Teddy {
            vb: vb,
            pats: pats.to_vec(),
            ac: AcAutomaton::new(pats.to_vec()).into_full(),
            buckets: buckets,
            masks: masks,
        })
    }

    /// Returns all of the substrings matched by this `Teddy`.
    pub fn patterns(&self) -> &[Vec<u8>] {
        &self.pats
    }

    /// Returns the number of substrings in this matcher.
    pub fn len(&self) -> usize {
        self.pats.len()
    }

    /// Returns the approximate size on the heap used by this matcher.
    pub fn approximate_size(&self) -> usize {
        self.pats.iter().fold(0, |a, b| a + b.len())
    }

    /// Searches `haystack` for the substrings in this `Teddy`. If a match was
    /// found, then it is returned. Otherwise, `None` is returned.
    pub fn find(&self, haystack: &[u8]) -> Option<Match> {
        // This is safe because the only way we can construct a Teddy type
        // is if AVX2 is available.
        unsafe { self.find_impl(haystack) }
    }

    #[allow(unused_attributes)]
    #[target_feature(enable = "avx2")]
    unsafe fn find_impl(&self, haystack: &[u8]) -> Option<Match> {
        // If our haystack is smaller than the block size, then fall back to
        // a naive brute force search.
        if haystack.is_empty() || haystack.len() < (BLOCK_SIZE + 2) {
            return self.slow(haystack, 0);
        }
        match self.masks.len() {
            0 => None,
            1 => self.find1(haystack),
            2 => self.find2(haystack),
            3 => self.find3(haystack),
            _ => unreachable!(),
        }
    }

    /// `find1` is used when there is only 1 mask. This is the easy case and is
    /// pretty much as described in the module documentation.
    #[inline(always)]
    fn find1(&self, haystack: &[u8]) -> Option<Match> {
        let mut pos = 0;
        let zero = self.vb.u8x32_splat(0);
        let len = haystack.len();
        debug_assert!(len >= BLOCK_SIZE);
        while pos <= len - BLOCK_SIZE {
            let h = unsafe {
                // I tried and failed to eliminate bounds checks in safe code.
                // This is safe because of our loop invariant: pos is always
                // <= len-32.
                let p = haystack.get_unchecked(pos..);
                self.vb.u8x32_load_unchecked_unaligned(p)
            };
            // N.B. `res0` is our `C` in the module documentation.
            let res0 = self.masks.members1(h);
            // Only do expensive verification if there are any non-zero bits.
            let bitfield = res0.ne(zero).movemask();
            if bitfield != 0 {
                if let Some(m) = self.verify(haystack, pos, res0, bitfield) {
                    return Some(m);
                }
            }
            pos += BLOCK_SIZE;
        }
        self.slow(haystack, pos)
    }

    /// `find2` is used when there are 2 masks, e.g., the fingerprint is 2 bytes
    /// long.
    #[inline(always)]
    fn find2(&self, haystack: &[u8]) -> Option<Match> {
        // This is an exotic way to right shift a SIMD vector across lanes.
        // See below at use for more details.
        let zero = self.vb.u8x32_splat(0);
        let len = haystack.len();
        // The previous value of `C` (from the module documentation) for the
        // *first* byte in the fingerprint. On subsequent iterations, we take
        // the last bitset from the previous `C` and insert it into the first
        // position of the current `C`, shifting all other bitsets to the right
        // one lane. This causes `C` for the first byte to line up with `C` for
        // the second byte, so that they can be `AND`'d together.
        let mut prev0 = self.vb.u8x32_splat(0xFF);
        let mut pos = 1;
        debug_assert!(len >= BLOCK_SIZE);
        while pos <= len - BLOCK_SIZE {
            let h = unsafe {
                // I tried and failed to eliminate bounds checks in safe code.
                // This is safe because of our loop invariant: pos is always
                // <= len-32.
                let p = haystack.get_unchecked(pos..);
                self.vb.u8x32_load_unchecked_unaligned(p)
            };
            let (res0, res1) = self.masks.members2(h);

            // Do this:
            //
            //     (prev0 << 15) | (res0 >> 1)
            //
            // This lets us line up our C values for each byte.
            let res0prev0 = res0.alignr_15(prev0);

            // `AND`'s our `C` values together.
            let res = res0prev0.and(res1);
            prev0 = res0;

            let bitfield = res.ne(zero).movemask();
            if bitfield != 0 {
                let pos = pos.checked_sub(1).unwrap();
                if let Some(m) = self.verify(haystack, pos, res, bitfield) {
                    return Some(m);
                }
            }
            pos += BLOCK_SIZE;
        }
        // The windowing above doesn't check the last byte in the last
        // window, so start the slow search at the last byte of the last
        // window.
        self.slow(haystack, pos.checked_sub(1).unwrap())
    }

    /// `find3` is used when there are 3 masks, e.g., the fingerprint is 3 bytes
    /// long.
    ///
    /// N.B. This is a straight-forward extrapolation of `find2`. The only
    /// difference is that we need to keep track of two previous values of `C`,
    /// since we now need to align for three bytes.
    #[inline(always)]
    fn find3(&self, haystack: &[u8]) -> Option<Match> {
        let zero = self.vb.u8x32_splat(0);
        let len = haystack.len();
        let mut prev0 = self.vb.u8x32_splat(0xFF);
        let mut prev1 = self.vb.u8x32_splat(0xFF);
        let mut pos = 2;

        while pos <= len - BLOCK_SIZE {
            let h = unsafe {
                // I tried and failed to eliminate bounds checks in safe code.
                // This is safe because of our loop invariant: pos is always
                // <= len-32.
                let p = haystack.get_unchecked(pos..);
                self.vb.u8x32_load_unchecked_unaligned(p)
            };
            let (res0, res1, res2) = self.masks.members3(h);

            let res0prev0 = res0.alignr_14(prev0);
            let res1prev1 = res1.alignr_15(prev1);
            let res = res0prev0.and(res1prev1).and(res2);

            prev0 = res0;
            prev1 = res1;

            let bitfield = res.ne(zero).movemask();
            if bitfield != 0 {
                let pos = pos.checked_sub(2).unwrap();
                if let Some(m) = self.verify(haystack, pos, res, bitfield) {
                    return Some(m);
                }
            }
            pos += BLOCK_SIZE;
        }
        // The windowing above doesn't check the last two bytes in the last
        // window, so start the slow search at the penultimate byte of the
        // last window.
        // self.slow(haystack, pos.saturating_sub(2))
        self.slow(haystack, pos.checked_sub(2).unwrap())
    }

    /// Runs the verification procedure on `res` (i.e., `C` from the module
    /// documentation), where the haystack block starts at `pos` in
    /// `haystack`. `bitfield` has ones in the bit positions that `res` has
    /// non-zero bytes.
    ///
    /// If a match exists, it returns the first one.
    #[inline(always)]
    fn verify(
        &self,
        haystack: &[u8],
        pos: usize,
        res: u8x32,
        mut bitfield: u32,
    ) -> Option<Match> {
        while bitfield != 0 {
            // The next offset, relative to pos, where some fingerprint
            // matched.
            let byte_pos = bitfield.trailing_zeros() as usize;
            bitfield &= !(1 << byte_pos);

            // Offset relative to the beginning of the haystack.
            let start = pos + byte_pos;

            // The bitfield telling us which patterns had fingerprints that
            // match at this starting position.
            let mut patterns = res.extract(byte_pos);
            while patterns != 0 {
                let bucket = patterns.trailing_zeros() as usize;
                patterns &= !(1 << bucket);

                // Actual substring search verification.
                if let Some(m) = self.verify_bucket(haystack, bucket, start) {
                    return Some(m);
                }
            }
        }

        None
    }

    /// Verifies whether any substring in the given bucket matches in haystack
    /// at the given starting position.
    #[inline(always)]
    fn verify_bucket(
        &self,
        haystack: &[u8],
        bucket: usize,
        start: usize,
    ) -> Option<Match> {
        // This cycles through the patterns in the bucket in the order that
        // the patterns were given. Therefore, we guarantee leftmost-first
        // semantics.
        for &pati in &self.buckets[bucket] {
            let pat = &*self.pats[pati];
            if start + pat.len() > haystack.len() {
                continue;
            }
            if pat == &haystack[start..start + pat.len()] {
                return Some(Match {
                    pat: pati,
                    start: start,
                    end: start + pat.len(),
                });
            }
        }
        None
    }

    /// Slow substring search through all patterns in this matcher.
    ///
    /// This is used when we don't have enough bytes in the haystack for our
    /// block based approach.
    #[inline(never)]
    fn slow(&self, haystack: &[u8], pos: usize) -> Option<Match> {
        self.ac.find(&haystack[pos..]).next().map(|m| {
            Match {
                pat: m.pati,
                start: pos + m.start,
                end: pos + m.end,
            }
        })
    }
}

/// A list of masks. This has length equal to the length of the fingerprint.
/// The length of the fingerprint is always `min(3, len(smallest_substring))`.
#[derive(Debug, Clone)]
struct Masks {
    vb: AVX2VectorBuilder,
    masks: [Mask; 3],
    size: usize,
}

impl Masks {
    /// Create a new set of masks of size `n`, where `n` corresponds to the
    /// number of bytes in a fingerprint.
    fn new(vb: AVX2VectorBuilder, n: usize) -> Masks {
        Masks {
            vb: vb,
            masks: [Mask::new(vb), Mask::new(vb), Mask::new(vb)],
            size: n,
        }
    }

    /// Returns the number of masks.
    fn len(&self) -> usize {
        self.size
    }

    /// Adds the given pattern to the given bucket. The bucket should be a
    /// power of `2 <= 2^7`.
    fn add(&mut self, bucket: u8, pat: &[u8]) {
        for i in 0..self.len() {
            self.masks[i].add(bucket, pat[i]);
        }
    }

    /// Finds the fingerprints that are in the given haystack block. i.e., this
    /// returns `C` as described in the module documentation.
    ///
    /// More specifically, `for i in 0..16` and `j in 0..8, C[i][j] == 1` if and
    /// only if `haystack_block[i]` corresponds to a fingerprint that is part
    /// of a pattern in bucket `j`.
    #[inline(always)]
    fn members1(&self, haystack_block: u8x32) -> u8x32 {
        let masklo = self.vb.u8x32_splat(0xF);
        let hlo = haystack_block.and(masklo);
        let hhi = haystack_block.bit_shift_right_4().and(masklo);

        self.masks[0].lo.shuffle(hlo).and(self.masks[0].hi.shuffle(hhi))
    }

    /// Like members1, but computes C for the first and second bytes in the
    /// fingerprint.
    #[inline(always)]
    fn members2(&self, haystack_block: u8x32) -> (u8x32, u8x32) {
        let masklo = self.vb.u8x32_splat(0xF);
        let hlo = haystack_block.and(masklo);
        let hhi = haystack_block.bit_shift_right_4().and(masklo);

        let res0 =
            self.masks[0].lo.shuffle(hlo).and(self.masks[0].hi.shuffle(hhi));
        let res1 =
            self.masks[1].lo.shuffle(hlo).and(self.masks[1].hi.shuffle(hhi));
        (res0, res1)
    }

    /// Like `members1`, but computes `C` for the first, second and third bytes
    /// in the fingerprint.
    #[inline(always)]
    fn members3(&self, haystack_block: u8x32) -> (u8x32, u8x32, u8x32) {
        let masklo = self.vb.u8x32_splat(0xF);
        let hlo = haystack_block.and(masklo);
        let hhi = haystack_block.bit_shift_right_4().and(masklo);

        let res0 =
            self.masks[0].lo.shuffle(hlo).and(self.masks[0].hi.shuffle(hhi));
        let res1 =
            self.masks[1].lo.shuffle(hlo).and(self.masks[1].hi.shuffle(hhi));
        let res2 =
            self.masks[2].lo.shuffle(hlo).and(self.masks[2].hi.shuffle(hhi));
        (res0, res1, res2)
    }
}

/// A single mask.
#[derive(Debug, Clone, Copy)]
struct Mask {
    /// Bitsets for the low nybbles in a fingerprint.
    lo: u8x32,
    /// Bitsets for the high nybbles in a fingerprint.
    hi: u8x32,
}

impl Mask {
    /// Create a new mask with no members.
    fn new(vb: AVX2VectorBuilder) -> Mask {
        Mask {
            lo: vb.u8x32_splat(0),
            hi: vb.u8x32_splat(0),
        }
    }

    /// Adds the given byte to the given bucket.
    fn add(&mut self, bucket: u8, byte: u8) {
        // Split our byte into two nybbles, and add each nybble to our
        // mask.
        let byte_lo = (byte & 0xF) as usize;
        let byte_hi = (byte >> 4) as usize;

        let lo = self.lo.extract(byte_lo) | ((1 << bucket) as u8);
        self.lo.replace(byte_lo, lo);
        self.lo.replace(byte_lo + 16, lo);

        let hi = self.hi.extract(byte_hi) | ((1 << bucket) as u8);
        self.hi.replace(byte_hi, hi);
        self.hi.replace(byte_hi + 16, hi);
    }
}