{-# LANGUAGE BangPatterns, FlexibleContexts #-}
{-# OPTIONS_HADDOCK hide, prune #-}
-- |
-- Module         : Data.ByteString.Search.Internal.Utils
-- Copyright      : Daniel Fischer
-- Licence        : BSD3
-- Maintainer     : Daniel Fischer <daniel.is.fischer@googlemail.com>
-- Stability      : Provisional
-- Portabiltity   : non-portable
--
-- Author         : Daniel Fischer
--
-- Utilities for several searching algorithms.

module Data.ByteString.Search.Internal.Utils ( kmpBorders
                                             , automaton
                                             , occurs
                                             , suffShifts
                                             , ldrop
                                             , ltake
                                             , lsplit
                                             , release
                                             , keep
                                             , strictify
                                             ) where

import qualified Data.ByteString as S
import qualified Data.ByteString.Lazy as L
import Data.ByteString.Unsafe (unsafeIndex)

import Data.Array.Base (unsafeRead, unsafeWrite, unsafeAt)
import Data.Array.ST
import Data.Array.Unboxed
import Control.Monad (when)

import Data.Bits
import Data.Word (Word8)

------------------------------------------------------------------------------
--                              Preprocessing                               --
------------------------------------------------------------------------------

{-# INLINE automaton #-}
automaton :: S.ByteString -> UArray Int Int
automaton !pat = runSTUArray (do
    let !patLen = S.length pat
        {-# INLINE patAt #-}
        patAt !i = fromIntegral (unsafeIndex pat i)
        !bord = kmpBorders pat
    aut <- newArray (0, (patLen + 1)*256 - 1) 0
    unsafeWrite aut (patAt 0) 1
    let loop !state = do
            let !base = state `shiftL` 8
                inner j
                    | j < 0     = if state == patLen
                                    then return aut
                                    else loop (state+1)
                    | otherwise = do
                        let !i = base + patAt j
                        s <- unsafeRead aut i
                        when (s == 0) (unsafeWrite aut i (j+1))
                        inner (unsafeAt bord j)
            if state == patLen
                then inner (unsafeAt bord state)
                else inner state
    loop 1)

-- kmpBorders calculates the width of the widest borders of the prefixes
-- of the pattern which are not extensible to borders of the next
-- longer prefix. Most entries will be 0.
{-# INLINE kmpBorders #-}
kmpBorders :: S.ByteString -> UArray Int Int
kmpBorders pat = runSTUArray (do
    let !patLen = S.length pat
        {-# INLINE patAt #-}
        patAt :: Int -> Word8
        patAt i = unsafeIndex pat i
    ar <- newArray_ (0, patLen)
    unsafeWrite ar 0 (-1)
    let dec w j
            | j < 0 || w == patAt j = return $! j+1
            | otherwise = unsafeRead ar j >>= dec w
        bordLoop !i !j
            | patLen < i    = return ar
            | otherwise     = do
                let !w = patAt (i-1)
                j' <- dec w j
                if i < patLen && patAt j' == patAt i
                    then unsafeRead ar j' >>= unsafeWrite ar i
                    else unsafeWrite ar i j'
                bordLoop (i+1) j'
    bordLoop 1 (-1))

------------------------------------------------------------------------------
--                        Boyer-Moore Preprocessing                         --
------------------------------------------------------------------------------

{- Table of last occurrences of bytes in the pattern.

For each byte we record the (negated) position of its last
occurrence in the pattern except at the last position.

Thus, if byte b gives a mismatch at pattern position patPos,
we know that we can shift the window right by at least

patPos - (last occurrence of b in init pat)

or, since we negated the positions,

patPos + (occurs pat)

If the byte doesn't occur in the pattern, we can shift the window
so that the start of the pattern is aligned with the byte after this,
hence the default value of 1.

Complexity: O(patLen + size of alphabet)

-}
{- Precondition: non-empty pattern

This invariant is guaranteed by not exporting occurs,
inside this module, we don't call it for empty patterns.

-}
{-# INLINE occurs #-}
occurs :: S.ByteString -> UArray Int Int
occurs pat = runSTUArray (do
    let !patEnd = S.length pat - 1
        {-# INLINE patAt #-}
        patAt :: Int -> Int
        patAt i = fromIntegral (unsafeIndex pat i)
    ar <- newArray (0, 255) 1
    let loop !i
            | i == patEnd   = return ar
            | otherwise     = do
                unsafeWrite ar (patAt i) (-i)
                loop (i + 1)
    loop 0)

{- Table of suffix-shifts.

When a mismatch occurs at pattern position patPos, assumed to be not the
last position in the pattern, the suffix u of length (patEnd - patPos)
has been successfully matched.
Let c be the byte in the pattern at position patPos.

If the sub-pattern u also occurs in the pattern somewhere *not* preceded
by c, let uPos be the position of the last byte in u for the last of
all such occurrences. Then there can be no match if the window is shifted
less than (patEnd - uPos) places, because either the part of the string
which matched the suffix u is not aligned with an occurrence of u in the
pattern, or it is aligned with an occurrence of u which is preceded by
the same byte c as the originally matched suffix.

If the complete sub-pattern u does not occur again in the pattern, or all
of its occurrences are preceded by the byte c, then we can align the
pattern with the string so that a suffix v of u matches a prefix of the
pattern. If v is chosen maximal, no smaller shift can give a match, so
we can shift by at least (patLen - length v).

If a complete match is encountered, we can shift by at least the same
amount as if the first byte of the pattern was a mismatch, no complete
match is possible between these positions.

For non-periodic patterns, only very short suffixes will usually occur
again in the pattern, so if a longer suffix has been matched before a
mismatch, the window can then be shifted entirely past the partial
match, so that part of the string will not be re-compared.
For periodic patterns, the suffix shifts will be shorter in general,
leading to an O(strLen * patLen) worst-case performance.

To compute the suffix-shifts, we use an array containing the lengths of
the longest common suffixes of the entire pattern and its prefix ending
with position pos.

-}
{- Precondition: non-empty pattern -}
{-# INLINE suffShifts #-}
suffShifts :: S.ByteString -> UArray Int Int
suffShifts pat = runSTUArray (do
    let !patLen = S.length pat
        !patEnd = patLen - 1
        !suff   = suffLengths pat
    ar <- newArray (0,patEnd) patLen
    let preShift !idx !j
            | idx < 0   = return ()
            | suff `unsafeAt` idx == idx + 1 = do
                let !shf = patEnd - idx
                    fillToShf !i
                        | i == shf  = return ()
                        | otherwise = do
                            unsafeWrite ar i shf
                            fillToShf (i + 1)
                fillToShf j
                preShift (idx - 1) shf
            | otherwise = preShift (idx - 1) j
        sufShift !idx
            | idx == patEnd = return ar
            | otherwise     = do
                unsafeWrite ar (patEnd - unsafeAt suff idx) (patEnd - idx)
                sufShift (idx + 1)
    preShift (patEnd - 1) 0
    sufShift 0)

{- Table of suffix-lengths.

The value of this array at place i is the length of the longest common
suffix of the entire pattern and the prefix of the pattern ending at
position i.

Usually, most of the entries will be 0. Only if the byte at position i
is the same as the last byte of the pattern can the value be positive.
In any case the value at index patEnd is patLen (since the pattern is
identical to itself) and 0 <= value at i <= (i + 1).

To keep this part of preprocessing linear in the length of the pattern,
the implementation must be non-obvious (the obvious algorithm for this
is quadratic).

When the index under consideration is inside a previously identified
common suffix, we align that suffix with the end of the pattern and
check whether the suffix ending at the position corresponding to idx
is shorter than the part of the suffix up to idx. If that is the case,
the length of the suffix ending at idx is that of the suffix at the
corresponding position. Otherwise extend the suffix as far as possible.
If the index under consideration is not inside a previously identified
common suffix, compare with the last byte of the pattern. If that gives
a suffix of length > 1, for the next index we're in the previous
situation, otherwise we're back in the same situation for the next
index.

-}
{- Precondition: non-empty pattern -}
{-# INLINE suffLengths #-}
suffLengths :: S.ByteString -> UArray Int Int
suffLengths pat = runSTUArray (do
    let !patLen = S.length pat
        !patEnd = patLen - 1
        !preEnd = patEnd - 1
        {-# INLINE patAt #-}
        patAt i = unsafeIndex pat i
        -- last byte for comparisons
        !pe     = patAt patEnd
        -- find index preceding the longest suffix
        dec !diff !j
            | j < 0 || patAt j /= patAt (j + diff) = j
            | otherwise = dec diff (j - 1)
    ar <- newArray_ (0, patEnd)
    unsafeWrite ar patEnd patLen
    let noSuff !i
            | i < 0     = return ar
            | patAt i == pe = do
                let !diff  = patEnd - i
                    !nextI = i - 1
                    !prevI = dec diff nextI
                if prevI == nextI
                    then unsafeWrite ar i 1 >> noSuff nextI
                    else do unsafeWrite ar i (i - prevI)
                            suffLoop prevI preEnd nextI
            | otherwise = do
                unsafeWrite ar i 0
                noSuff (i - 1)
        suffLoop !pre !end !idx
            | idx < 0   = return ar
            | pre < idx =
              if patAt idx /= pe
                then unsafeWrite ar idx 0 >> suffLoop pre (end - 1) (idx - 1)
                else do
                    prevS <- unsafeRead ar end
                    if pre + prevS < idx
                        then do unsafeWrite ar idx prevS
                                suffLoop pre (end - 1) (idx - 1)
                        else do let !prI = dec (patEnd - idx) pre
                                unsafeWrite ar idx (idx - prI)
                                suffLoop prI preEnd (idx - 1)
            | otherwise = noSuff idx
    noSuff preEnd)

------------------------------------------------------------------------------
--                             Helper Functions                             --
------------------------------------------------------------------------------

{-# INLINE strictify #-}
strictify :: L.ByteString -> S.ByteString
strictify = S.concat . L.toChunks

-- drop k bytes from a list of strict ByteStrings
{-# INLINE ldrop #-}
ldrop :: Int -> [S.ByteString] -> [S.ByteString]
ldrop _ [] = []
ldrop k (!h : t)
  | k < l     = S.drop k h : t
  | otherwise = ldrop (k - l) t
    where
      !l = S.length h

-- take k bytes from a list of strict ByteStrings
{-# INLINE ltake #-}
ltake :: Int -> [S.ByteString] -> [S.ByteString]
ltake _ [] = []
ltake !k (!h : t)
  | l < k     = h : ltake (k - l) t
  | otherwise = [S.take k h]
    where
      !l = S.length h

-- split a list of strict ByteStrings at byte k
{-# INLINE lsplit #-}
lsplit :: Int -> [S.ByteString] -> ([S.ByteString], [S.ByteString])
lsplit _ [] = ([],[])
lsplit !k (!h : t)
  = case compare k l of
      LT -> ([S.take k h], S.drop k h : t)
      EQ -> ([h], t)
      GT -> let (u, v) = lsplit (k - l) t in (h : u, v)
  where
    !l = S.length h


-- release is used to keep the zipper in lazySearcher from remembering
-- the leading part of the searched string.  The deep parameter is the
-- number of characters that the past needs to hold.  This ensures
-- lazy streaming consumption of the searched string.
{-# INLINE release #-}
release :: Int ->  [S.ByteString] -> [S.ByteString]
release !deep _
    | deep <= 0 = []
release !deep (!x:xs) = let !rest = release (deep-S.length x) xs in x : rest
release _ [] = error "stringsearch.release could not find enough past!"

-- keep is like release, only we mustn't forget the part of the past
-- we don't need anymore for matching but have to keep it for
-- breaking, splitting and replacing.
-- The names would be more appropriate the other way round, but that's
-- a historical accident, so what?
{-# INLINE keep #-}
keep :: Int -> [S.ByteString] -> ([S.ByteString],[S.ByteString])
keep !deep xs
    | deep < 1    = ([],xs)
keep deep (!x:xs) = let (!p,d) = keep (deep - S.length x) xs in (x:p,d)
keep _ [] = error "Forgot too much"