io-streams-1.5.2.2: src/System/IO/Streams/Internal/Search.hs
{-# LANGUAGE BangPatterns #-}
{-# LANGUAGE OverloadedStrings #-}
{-# LANGUAGE RankNTypes #-}
module System.IO.Streams.Internal.Search
( search
, MatchInfo(..)
) where
------------------------------------------------------------------------------
import Control.Monad (when)
import Control.Monad.IO.Class (liftIO)
import Control.Monad.ST (ST)
import Data.ByteString.Char8 (ByteString)
import qualified Data.ByteString.Char8 as S
import qualified Data.ByteString.Unsafe as S
import qualified Data.Vector.Unboxed as V
import qualified Data.Vector.Unboxed.Mutable as MV
import Prelude (Bool (..), Either (..), Enum (..), Eq (..), IO, Int, Monad (..), Num (..), Ord (..), Show, either, id, maybe, not, otherwise, ($), ($!), (&&), (.), (||))
------------------------------------------------------------------------------
import System.IO.Streams.Internal (InputStream)
import qualified System.IO.Streams.Internal as Streams
------------------------------------------------------------------------------
-- | 'MatchInfo' provides match information when performing string search.
data MatchInfo = Match {-# UNPACK #-} !ByteString
| NoMatch {-# UNPACK #-} !ByteString
deriving (Show, Eq)
------------------------------------------------------------------------------
-- | Does the given needle match the haystack over the given ranges of indices?
matches :: ByteString -- ^ needle
-> Int -- ^ needle start
-> Int -- ^ needle end (inclusive)
-> ByteString -- ^ haystack
-> Int -- ^ haystack start
-> Int -- ^ haystack end (inclusive)
-> Bool
matches !needle !nstart !nend' !haystack !hstart !hend' =
go nend' hend'
where
go !nend !hend =
if nend < nstart || hend < hstart
then True
else let !nc = S.unsafeIndex needle nend
!hc = S.unsafeIndex haystack hend
in if nc /= hc
then False
else go (nend-1) (hend-1)
{-# INLINE matches #-}
------------------------------------------------------------------------------
-- | Given a 'ByteString' to look for (the \"needle\") and an 'InputStream',
-- produces a new 'InputStream' which yields data of type 'MatchInfo'.
--
-- Example:
--
-- @
-- ghci> 'System.IO.Streams.fromList' [\"food\", \"oof\", \"oodles\", \"ok\"] >>=
-- 'search' \"foo\" >>= 'System.IO.Streams.toList'
-- ['Match' \"foo\",'NoMatch' \"d\",'NoMatch' \"oo\",'Match' \"foo\",'NoMatch' \"dlesok\"]
-- @
--
-- Uses the Boyer-Moore-Horspool algorithm
-- (<http://en.wikipedia.org/wiki/Boyer%E2%80%93Moore%E2%80%93Horspool_algorithm>).
search :: ByteString -- ^ \"needle\" to look for
-> InputStream ByteString -- ^ input stream to wrap
-> IO (InputStream MatchInfo)
search needle stream = Streams.fromGenerator $
lookahead nlen >>= either finishAndEOF startSearch
where
--------------------------------------------------------------------------
finishAndEOF x = if S.null x
then return $! ()
else Streams.yield $! NoMatch x
--------------------------------------------------------------------------
startSearch !haystack =
if S.null haystack
then lookahead nlen >>= either finishAndEOF startSearch
else go 0
where
----------------------------------------------------------------------
!hlen = S.length haystack
----------------------------------------------------------------------
go !hidx
| hend >= hlen = crossBound hidx
| otherwise = do
let match = matches needle 0 lastIdx haystack hidx hend
if match
then do
let !nomatch = S.take hidx haystack
let !aftermatch = S.drop (hend + 1) haystack
produceMatch nomatch aftermatch
else do
-- skip ahead
let c = S.unsafeIndex haystack hend
let !skip = V.unsafeIndex table $ fromEnum c
go (hidx + skip)
where
!hend = hidx + nlen - 1
----------------------------------------------------------------------
mkCoeff hidx = let !ll = hlen - hidx
!nm = nlen - ll
in (ll, nm)
----------------------------------------------------------------------
crossBound !hidx0 = do
let (!leftLen, needMore) = mkCoeff hidx0
lookahead needMore >>=
either (\s -> finishAndEOF $ S.append haystack s)
(runNext hidx0 leftLen needMore)
where
runNext !hidx !leftLen !needMore !nextHaystack = do
let match1 = matches needle leftLen lastIdx nextHaystack 0
(needMore-1)
let match2 = matches needle 0 (leftLen-1) haystack hidx
(hlen-1)
if match1 && match2
then do
let !nomatch = S.take hidx haystack
let !aftermatch = S.drop needMore nextHaystack
produceMatch nomatch aftermatch
else do
let c = S.unsafeIndex nextHaystack $ needMore - 1
let p = V.unsafeIndex table (fromEnum c)
if p < leftLen
then do
let !hidx' = hidx + p
let (!leftLen', needMore') = mkCoeff hidx'
let !nextlen = S.length nextHaystack
if nextlen < needMore'
then
-- this should be impossibly rare
lookahead (needMore' - nextlen) >>=
either (\s -> finishAndEOF $
S.concat [ haystack
, nextHaystack
, s ])
(\s -> runNext hidx' leftLen' needMore' $
S.append nextHaystack s)
else runNext hidx' leftLen' needMore' nextHaystack
else do
let sidx = p - leftLen
let (!crumb, rest) = S.splitAt sidx nextHaystack
Streams.yield $! NoMatch $! S.append haystack crumb
startSearch rest
--------------------------------------------------------------------------
produceMatch nomatch aftermatch = do
when (not $ S.null nomatch) $ Streams.yield $! NoMatch nomatch
Streams.yield $! Match needle
startSearch aftermatch
--------------------------------------------------------------------------
!nlen = S.length needle
!lastIdx = nlen - 1
--------------------------------------------------------------------------
!table = V.create $ do
t <- MV.replicate 256 nlen
go t
where
go :: forall s . MV.MVector s Int -> ST s (MV.MVector s Int)
go !t = go' 0
where
go' !i | i >= lastIdx = return t
| otherwise = do
let c = fromEnum $ S.unsafeIndex needle i
MV.unsafeWrite t c (lastIdx - i)
go' $! i+1
--------------------------------------------------------------------------
lookahead n = go id n
where
go dlist !k = liftIO (Streams.read stream) >>= maybe eof chunk
where
eof = return $! Left $! S.concat $ dlist []
chunk x = if r <= 0
then return $! Right $! S.concat $ d' []
else go d' r
where
l = S.length x
r = k - l
d' = dlist . (x:)