{-# LANGUAGE BangPatterns #-}
module Data.ByteString.Substring
( breakSubstringLazy
, prepareBreakSubstring
, breakSubstringResume
, KarpRabinState
, KarpRabinResult(..)
) where
import Control.Monad.ST (RealWorld)
import Data.Bits
import Data.ByteString
import Data.ByteString.Unsafe
import Data.Int
import Data.Primitive.ByteArray
import Data.Word
import Prelude hiding (length,null)
import System.IO.Unsafe (unsafePerformIO)
import qualified Data.ByteString.Lazy as LB
import qualified Data.ByteString.Lazy.Internal as LBI
breakSubstringLazy ::
ByteString -- ^ needle, strict bytestring
-> LB.ByteString -- ^ haystack, lazy bytestring
-> (LB.ByteString,LB.ByteString)
breakSubstringLazy pat lb =
let !lp = intToInt64 (length pat)
!b0 = LB.toStrict (LB.take lp lb)
!lb' = LB.drop lp lb
!s = prepareBreakSubstring pat b0
in case go s lb' of
BreakBackwardDone pre post -> (if null b0 then pre else LBI.Chunk b0 pre,post)
BreakBackwardBy !n -> let (bpre,bpost) = unsafeSplitAt (length b0 - n) b0 in
(LB.fromStrict bpre,if null bpost then lb' else LBI.Chunk bpost lb')
where
go :: KarpRabinState -> LB.ByteString -> BreakBackward
go _ LBI.Empty = BreakBackwardDone LBI.Empty LBI.Empty
go s1 (LBI.Chunk !c cs) = case breakSubstringResume s1 c of
KarpRabinResultMore s2 -> case go s2 cs of
BreakBackwardBy n -> if n <= length c
then let (ca,cb) = unsafeSplitAt (length c - n) c in
BreakBackwardDone
(if null ca then LBI.Empty else LBI.Chunk ca LBI.Empty)
(if null cb then LBI.Empty else LBI.Chunk cb cs)
else BreakBackwardBy (n - length c)
BreakBackwardDone cs' cs'' -> BreakBackwardDone (LBI.Chunk c cs') cs''
KarpRabinResultDone ix _ -> if ix < 0
then BreakBackwardBy (negate ix)
else let (ca,cb) = unsafeSplitAt ix c in
BreakBackwardDone
(if null ca then LBI.Empty else LBI.Chunk ca LBI.Empty)
(LBI.Chunk cb cs)
prepareBreakSubstring ::
ByteString -- ^ needle
-> ByteString -- ^ first n characters in haystack, where n is length of needle
-> KarpRabinState
prepareBreakSubstring pat b0 =
let !lp = intToInt64 (length pat)
!log2BufSize = finiteBitSize (0 :: Int) - countLeadingZeros (length pat)
!bufSize = 2 ^ log2BufSize
!buf = unsafePerformIO (newBuffer bufSize b0)
!m = 2891336453 ^ length pat
!hp = rollingHash pat
!hs0 = rollingHash b0
in (KarpRabinState hp m pat hs0 log2BufSize (bufSize - 1) buf lp)
rollingHash :: ByteString -> Word32
rollingHash = foldl' (\h b -> h * 2891336453 + word8ToWord32 b) 0
data BreakBackward
= BreakBackwardDone LB.ByteString LB.ByteString
| BreakBackwardBy !Int
-- ^ The int in here should always be positive
intToInt64 :: Int -> Int64
intToInt64 = fromIntegral
int64ToInt :: Int64 -> Int
int64ToInt = fromIntegral
word8ToWord32 :: Word8 -> Word32
word8ToWord32 = fromIntegral
breakSubstringResume :: KarpRabinState -> ByteString -> KarpRabinResult
breakSubstringResume (KarpRabinState hp m pat hs0 log2BufSize bufMask buf i64Start) chunk =
unsafePerformIO (search hs0 i64Start)
where
k = 2891336453 :: Word32
readByteArrayMod :: Int -> IO Word8
readByteArrayMod ix = readByteArray buf (ix .&. bufMask)
writeByteArrayMod :: Int -> Word8 -> IO ()
writeByteArrayMod ix = writeByteArray buf (ix .&. bufMask)
search :: Word32 -> Int64 -> IO KarpRabinResult
search !hs !i64 = do
let !i = int64ToInt (i64 - i64Start)
keepGoing :: IO KarpRabinResult
keepGoing = do
oldVal <- readByteArrayMod (int64ToInt (i64 - intToInt64 (length pat)))
let w8 :: Word8
!w8 = unsafeIndex chunk i
!hs' = hs * k + fromIntegral w8 - m * fromIntegral oldVal
writeByteArrayMod (int64ToInt i64) w8
search hs' (i64 + 1)
if hp == hs
then do
b <- mutableByteArrayEqByteString
(bufMask .&. (int64ToInt i64 - length pat))
bufMask buf pat
if b
then return (KarpRabinResultDone (i - length pat) (i64 - fromIntegral (length pat)))
else keepGoing
else if length chunk <= i
then return (KarpRabinResultMore (KarpRabinState hp m pat hs log2BufSize bufMask buf i64))
else keepGoing
unsafeSplitAt :: Int -> ByteString -> (ByteString,ByteString)
unsafeSplitAt i s = (unsafeTake i s, unsafeDrop i s)
-- | You must provide a size equal to or larger than the
-- ByteString length
newBuffer :: Int -> ByteString -> IO (MutableByteArray RealWorld)
newBuffer sz bs = do
arr <- newByteArray sz
copyIntoBuffer bs arr
return arr
copyIntoBuffer :: ByteString -> MutableByteArray RealWorld -> IO ()
copyIntoBuffer bs arr = go 0
where
go :: Int -> IO ()
go ix = if ix < length bs
then do
let !w = unsafeIndex bs ix
writeByteArray arr ix w
go (ix + 1)
else return ()
mutableByteArrayEqByteString :: Int -> Int -> MutableByteArray RealWorld -> ByteString -> IO Bool
mutableByteArrayEqByteString arrIx bufMask arr bs = go 0
where
readByteArrayMod :: Int -> IO Word8
readByteArrayMod ix = readByteArray arr (ix .&. bufMask)
go :: Int -> IO Bool
go i = if i < length bs
then do
w1 <- readByteArrayMod (arrIx + i)
let w2 = unsafeIndex bs i
if w1 == w2 then go (i + 1) else return False
else return True
data KarpRabinState = KarpRabinState
!Word32 -- pattern fingerprint
!Word32 -- constant k exponentiated
!ByteString -- pattern
!Word32 -- current fingerprint
!Int -- log base 2 of buffer size
!Int -- mask
!(MutableByteArray RealWorld) -- current buffer, contains end of previous bytestring
!Int64 -- total number of bytes consumed, also works as buffer index
-- after doing some modular arithmetic
data KarpRabinResult
= KarpRabinResultDone !Int !Int64
-- ^ The first number is the index into the current chunk.
-- The second number is the total number of characters
-- that were consumed. Note that since these both refer
-- to the index of the beginning of the match, the first
-- one is allowed to be negative, but the second is not.
-- The third item is the bytes preceeding the match location.
-- This is provided to help streaming providers that may have
-- already discarded the old data.
| KarpRabinResultMore !KarpRabinState