module Snap.Iteratee
(
Stream
, IterV
, Iteratee
, Enumerator
, module Data.Iteratee
, enumBS
, enumLBS
, enumFile
, fromWrap
, toWrap
, drop'
, takeExactly
, takeNoMoreThan
, countBytes
, bufferIteratee
, mkIterateeBuffer
, unsafeBufferIterateeWithBuffer
, unsafeBufferIteratee
) where
import Control.Monad
import Control.Monad.CatchIO
import Data.ByteString (ByteString)
import qualified Data.ByteString as S
import qualified Data.ByteString.Unsafe as S
import qualified Data.ByteString.Lazy as L
import Data.Int
import Data.IORef
import Data.Iteratee
#ifdef PORTABLE
import Data.Iteratee.IO (enumHandle)
#endif
import qualified Data.Iteratee.Base.StreamChunk as SC
import Data.Iteratee.WrappedByteString
import qualified Data.ListLike as LL
import Data.Monoid (mappend)
import Foreign
import Foreign.C.Types
import GHC.ForeignPtr
import Prelude hiding (catch,drop)
import qualified Data.DList as D
#ifdef PORTABLE
import Control.Monad.Trans (liftIO)
import System.IO
#else
import Control.Exception (SomeException)
import System.IO.Posix.MMap
#endif
type Stream = StreamG WrappedByteString Word8
type IterV m = IterGV WrappedByteString Word8 m
type Iteratee m = IterateeG WrappedByteString Word8 m
type Enumerator m a = Iteratee m a -> m (Iteratee m a)
instance (Functor m, MonadCatchIO m) =>
MonadCatchIO (IterateeG s el m) where
catch m handler = IterateeG $ \str -> do
ee <- try $ runIter m str
case ee of
(Left e) -> runIter (handler e) str
(Right v) -> return v
block m = IterateeG $ \str -> block $ runIter m str
unblock m = IterateeG $ \str -> unblock $ runIter m str
countBytes :: (Monad m) => Iteratee m a -> Iteratee m (a, Int64)
countBytes = go 0
where
go !n iter = IterateeG $ f n iter
f !n !iter ch@(Chunk ws) = do
iterv <- runIter iter ch
case iterv of
Done x rest -> let !n' = n + m len rest
in return $! Done (x, n') rest
Cont i err -> return $ Cont ((go $! n + m) i) err
where
m = fromIntegral $ S.length (unWrap ws)
len (EOF _) = 0
len (Chunk s) = fromIntegral $ S.length (unWrap s)
f !n !iter stream = do
iterv <- runIter iter stream
case iterv of
Done x rest -> return $ Done (x, n) rest
Cont i err -> return $ Cont (go n i) err
bufferIteratee :: Iteratee IO a -> IO (Iteratee IO a, IORef Bool)
bufferIteratee iteratee = do
esc <- newIORef False
return $ (start esc iteratee, esc)
where
blocksize = 8192
start esc iter = IterateeG $! checkRef esc iter
checkRef esc iter ch = do
quit <- readIORef esc
if quit
then runIter iter ch
else f (D.empty,0) iter ch
go (!dl,!n) iter = IterateeG $! f (dl,n) iter
f _ !iter ch@(EOF (Just _)) = runIter iter ch
f (!dl,_) !iter ch@(EOF Nothing) = do
iter' <- if S.null str
then return iter
else liftM liftI $ runIter iter $ Chunk big
runIter iter' ch
where
str = S.concat $ D.toList dl
big = WrapBS str
f (!dl,!n) iter (Chunk (WrapBS s)) =
if n' >= blocksize
then do
iterv <- runIter iter (Chunk big)
case iterv of
Done x rest -> return $ Done x rest
Cont i (Just e) -> return $ Cont i (Just e)
Cont i Nothing -> return $ Cont (go (D.empty,0) i) Nothing
else return $ Cont (go (dl',n') iter) Nothing
where
m = S.length s
n' = n+m
dl' = D.snoc dl s
big = WrapBS $ S.concat $ D.toList dl'
bUFSIZ :: Int
bUFSIZ = 8192
mkIterateeBuffer :: IO (ForeignPtr CChar)
mkIterateeBuffer = mallocPlainForeignPtrBytes bUFSIZ
unsafeBufferIteratee :: Iteratee IO a -> IO (Iteratee IO a, IORef Bool)
unsafeBufferIteratee iter = do
buf <- mkIterateeBuffer
unsafeBufferIterateeWithBuffer buf iter
unsafeBufferIterateeWithBuffer :: ForeignPtr CChar
-> Iteratee IO a
-> IO (Iteratee IO a, IORef Bool)
unsafeBufferIterateeWithBuffer buf iteratee = do
esc <- newIORef False
return $! (start esc iteratee, esc)
where
start esc iter = IterateeG $! checkRef esc iter
go bytesSoFar iter =
IterateeG $! f bytesSoFar iter
checkRef esc iter ch = do
quit <- readIORef esc
if quit
then runIter iter ch
else f 0 iter ch
sendBuf n iter =
withForeignPtr buf $ \ptr -> do
s <- S.unsafePackCStringLen (ptr, n)
runIter iter $ Chunk $ WrapBS s
copy c@(EOF _) = c
copy (Chunk (WrapBS s)) = Chunk $ WrapBS $ S.copy s
f _ iter ch@(EOF (Just _)) = runIter iter ch
f !n iter ch@(EOF Nothing) =
if n == 0
then runIter iter ch
else do
iter' <- liftM liftI $ sendBuf n iter
runIter iter' ch
f !n iter (Chunk (WrapBS s)) = do
let m = S.length s
if m+n > bUFSIZ
then overflow n iter s m
else copyAndCont n iter s m
copyAndCont n iter s m =
do
S.unsafeUseAsCStringLen s $ \(p,sz) ->
withForeignPtr buf $ \bufp -> do
let b' = plusPtr bufp n
copyBytes b' p sz
return $ Cont (go (n+m) iter) Nothing
overflow n iter s m =
do
let rest = bUFSIZ n
let m2 = m rest
let (s1,s2) = S.splitAt rest s
S.unsafeUseAsCStringLen s1 $ \(p,_) ->
withForeignPtr buf $ \bufp -> do
let b' = plusPtr bufp n
copyBytes b' p rest
iv <- sendBuf bUFSIZ iter
case iv of
Done x r -> return $
Done x (copy r `mappend` (Chunk $ WrapBS s2))
Cont i (Just e) -> return $ Cont i (Just e)
Cont i Nothing -> do
if m2 >= bUFSIZ
then do
iv' <- runIter i (Chunk $ WrapBS s2)
case iv' of
Done x r -> return $ Done x (copy r)
Cont i' (Just e) -> return $ Cont i' (Just e)
Cont i' Nothing -> return $ Cont (go 0 i') Nothing
else copyAndCont 0 i s2 m2
enumBS :: (Monad m) => ByteString -> Enumerator m a
enumBS bs = enumPure1Chunk $ WrapBS bs
enumLBS :: (Monad m) => L.ByteString -> Enumerator m a
enumLBS lbs = el chunks
where
el [] i = liftM liftI $ runIter i (EOF Nothing)
el (x:xs) i = do
i' <- liftM liftI $ runIter i (Chunk $ WrapBS x)
el xs i'
chunks = L.toChunks lbs
toWrap :: L.ByteString -> WrappedByteString Word8
toWrap = WrapBS . S.concat . L.toChunks
fromWrap :: WrappedByteString Word8 -> L.ByteString
fromWrap = L.fromChunks . (:[]) . unWrap
drop' :: (SC.StreamChunk s el, Monad m)
=> Int64
-> IterateeG s el m ()
drop' 0 = return ()
drop' n = IterateeG step
where
step (Chunk str)
| strlen <= n = return $ Cont (drop' (n strlen)) Nothing
where
strlen = fromIntegral $ SC.length str
step (Chunk str) = return $ Done () (Chunk (LL.drop (fromIntegral n) str))
step stream = return $ Done () stream
takeExactly :: (SC.StreamChunk s el, Monad m)
=> Int64
-> EnumeratorN s el s el m a
takeExactly 0 iter = return iter
takeExactly n' iter =
if n' < 0
then takeExactly 0 iter
else IterateeG (step n')
where
step n chk@(Chunk str)
| SC.null str = return $ Cont (takeExactly n iter) Nothing
| strlen < n = liftM (flip Cont Nothing) inner
| otherwise = done (Chunk s1) (Chunk s2)
where
strlen = fromIntegral $ SC.length str
inner = liftM (check (n strlen)) (runIter iter chk)
(s1, s2) = SC.splitAt (fromIntegral n) str
step _n (EOF (Just e)) = return $ Cont undefined (Just e)
step _n (EOF Nothing) = return $ Cont undefined (Just (Err "short write"))
check n (Done x _) = drop' n >> return (return x)
check n (Cont x Nothing) = takeExactly n x
check n (Cont _ (Just e)) = drop' n >> throwErr e
done s1 s2 = liftM (flip Done s2) (runIter iter s1 >>= checkIfDone return)
takeNoMoreThan :: (SC.StreamChunk s el, Monad m)
=> Int64
-> EnumeratorN s el s el m a
takeNoMoreThan n' iter =
if n' < 0
then takeNoMoreThan 0 iter
else IterateeG (step n')
where
step n chk@(Chunk str)
| SC.null str = return $ Cont (takeNoMoreThan n iter) Nothing
| strlen < n = liftM (flip Cont Nothing) inner
| otherwise = done (Chunk s1) (Chunk s2)
where
strlen = fromIntegral $ SC.length str
inner = liftM (check (n strlen)) (runIter iter chk)
(s1, s2) = SC.splitAt (fromIntegral n) str
step _n (EOF (Just e)) = return $ Cont undefined (Just e)
step _n chk@(EOF Nothing) = do
v <- runIter iter chk
case v of
(Done x s) -> return $ Done (return x) s
(Cont _ (Just e)) -> return $ Cont undefined (Just e)
(Cont _ Nothing) -> return $ Cont (throwErr $ Err "premature EOF") Nothing
check _ v@(Done _ _) = return $ liftI v
check n (Cont x Nothing) = takeNoMoreThan n x
check _ (Cont _ (Just e)) = throwErr e
done _ (EOF _) = error "impossible"
done s1 s2@(Chunk s2') = do
v <- runIter iter s1
case v of
(Done x s') -> return $ Done (return x) (s' `mappend` s2)
(Cont _ (Just e)) -> return $ Cont undefined (Just e)
(Cont i Nothing) ->
if SC.null s2'
then return $ Cont (takeNoMoreThan 0 i) Nothing
else return $ Cont undefined (Just $ Err "too many bytes")
enumFile :: FilePath -> Iteratee IO a -> IO (Iteratee IO a)
#ifdef PORTABLE
enumFile fp iter = do
h <- liftIO $ openBinaryFile fp ReadMode
i' <- enumHandle h iter
return $ do
x <- i'
liftIO (hClose h)
return x
#else
enumFile fp iter = do
es <- (try $
liftM WrapBS $
unsafeMMapFile fp) :: IO (Either SomeException (WrappedByteString Word8))
case es of
(Left e) -> return $ throwErr $ Err $ "IO error" ++ show e
(Right s) -> liftM liftI $ runIter iter $ Chunk s
#endif