{-# LANGUAGE BangPatterns #-}
{-# LANGUAGE DeriveDataTypeable #-}
module System.IO.Streams.Combinators
(
inputFoldM
, outputFoldM
, fold
, foldM
, any
, all
, maximum
, minimum
, unfoldM
, map
, mapM
, mapM_
, contramap
, contramapM
, contramapM_
, filter
, filterM
, filterOutput
, filterOutputM
, give
, take
, drop
, ignore
, zip
, zipWith
, zipWithM
, unzip
, intersperse
, skipToEof
, ignoreEof
, atEndOfInput
, atEndOfOutput
) where
import Control.Concurrent.MVar (newMVar, withMVar)
import Control.Monad (liftM, void, when)
import Control.Monad.IO.Class (liftIO)
import Data.Int (Int64)
import Data.IORef (atomicModifyIORef, modifyIORef, newIORef, readIORef, writeIORef)
import Data.Maybe (isJust)
import Prelude hiding (all, any, drop, filter, map, mapM, mapM_, maximum, minimum, read, take, unzip, zip, zipWith)
import System.IO.Streams.Internal (InputStream (..), OutputStream, fromGenerator, makeInputStream, makeOutputStream, read, unRead, write, yield)
outputFoldM :: (a -> b -> IO a)
-> a
-> OutputStream b
-> IO (OutputStream b, IO a)
outputFoldM f initial stream = do
ref <- newIORef initial
os <- makeOutputStream (wr ref)
return (os, fetch ref)
where
wr _ Nothing = write Nothing stream
wr ref mb@(Just x) = do
!z <- readIORef ref
!z' <- f z x
writeIORef ref z'
write mb stream
fetch ref = atomicModifyIORef ref $ \x -> (initial, x)
inputFoldM :: (a -> b -> IO a)
-> a
-> InputStream b
-> IO (InputStream b, IO a)
inputFoldM f initial stream = do
ref <- newIORef initial
is <- makeInputStream (rd ref)
return (is, fetch ref)
where
twiddle _ Nothing = return Nothing
twiddle ref mb@(Just x) = do
!z <- readIORef ref
!z' <- f z x
writeIORef ref z'
return mb
rd ref = read stream >>= twiddle ref
fetch ref = atomicModifyIORef ref $ \x -> (initial, x)
fold :: (s -> a -> s)
-> s
-> InputStream a
-> IO s
fold f seed stream = go seed
where
go !s = read stream >>= maybe (return s) (go . f s)
foldM :: (s -> a -> IO s)
-> s
-> InputStream a
-> IO s
foldM f seed stream = go seed
where
go !s = read stream >>= maybe (return s) ((go =<<) . f s)
any :: (a -> Bool) -> InputStream a -> IO Bool
any predicate stream = go
where
go = do
mElem <- read stream
case mElem of
Nothing -> return False
Just e -> if predicate e then return True else go
all :: (a -> Bool) -> InputStream a -> IO Bool
all predicate stream = go
where
go = do
mElem <- read stream
case mElem of
Nothing -> return True
Just e -> if predicate e then go else return False
maximum :: (Ord a) => InputStream a -> IO (Maybe a)
maximum stream = do
mElem0 <- read stream
case mElem0 of
Nothing -> return Nothing
Just e -> go e
where
go oldElem = do
mElem <- read stream
case mElem of
Nothing -> return (Just oldElem)
Just newElem -> go (max oldElem newElem)
minimum :: (Ord a) => InputStream a -> IO (Maybe a)
minimum stream = do
mElem0 <- read stream
case mElem0 of
Nothing -> return Nothing
Just e -> go e
where
go oldElem = do
mElem <- read stream
case mElem of
Nothing -> return (Just oldElem)
Just newElem -> go (min oldElem newElem)
unfoldM :: (b -> IO (Maybe (a, b))) -> b -> IO (InputStream a)
unfoldM f seed = fromGenerator (go seed)
where
go oldSeed = do
m <- liftIO (f oldSeed)
case m of
Nothing -> return $! ()
Just (a, newSeed) -> do
yield a
go newSeed
map :: (a -> b) -> InputStream a -> IO (InputStream b)
map f s = makeInputStream g
where
g = read s >>= return . fmap f
mapM :: (a -> IO b) -> InputStream a -> IO (InputStream b)
mapM f s = makeInputStream g
where
g = do
mb <- read s >>= maybe (return Nothing)
(\x -> liftM Just $ f x)
return mb
mapM_ :: (a -> IO b) -> InputStream a -> IO (InputStream a)
mapM_ f s = makeInputStream $ do
mb <- read s
_ <- maybe (return $! ()) (void . f) mb
return mb
contramap :: (a -> b) -> OutputStream b -> IO (OutputStream a)
contramap f s = makeOutputStream $ flip write s . fmap f
contramapM :: (a -> IO b) -> OutputStream b -> IO (OutputStream a)
contramapM f s = makeOutputStream g
where
g Nothing = write Nothing s
g (Just x) = do
!y <- f x
write (Just y) s
contramapM_ :: (a -> IO b) -> OutputStream a -> IO (OutputStream a)
contramapM_ f s = makeOutputStream $ \mb -> do
_ <- maybe (return $! ()) (void . f) mb
write mb s
skipToEof :: InputStream a -> IO ()
skipToEof str = go
where
go = read str >>= maybe (return $! ()) (const go)
{-# INLINE skipToEof #-}
filterM :: (a -> IO Bool)
-> InputStream a
-> IO (InputStream a)
filterM p src = return $! InputStream prod pb
where
prod = read src >>= maybe eof chunk
chunk s = do
b <- p s
if b then return $! Just s
else prod
eof = return Nothing
pb s = unRead s src
filter :: (a -> Bool)
-> InputStream a
-> IO (InputStream a)
filter p src = return $! InputStream prod pb
where
prod = read src >>= maybe eof chunk
chunk s = do
let b = p s
if b then return $! Just s
else prod
eof = return Nothing
pb s = unRead s src
intersperse :: a -> OutputStream a -> IO (OutputStream a)
intersperse sep os = newIORef False >>= makeOutputStream . f
where
f _ Nothing = write Nothing os
f sendRef s = do
b <- readIORef sendRef
writeIORef sendRef True
when b $ write (Just sep) os
write s os
zip :: InputStream a -> InputStream b -> IO (InputStream (a, b))
zip src1 src2 = makeInputStream src
where
src = read src1 >>= (maybe (return Nothing) $ \a ->
read src2 >>= (maybe (unRead a src1 >> return Nothing) $ \b ->
return $! Just $! (a, b)))
zipWith :: (a -> b -> c)
-> InputStream a
-> InputStream b
-> IO (InputStream c)
zipWith f src1 src2 = makeInputStream src
where
src = read src1 >>= (maybe (return Nothing) $ \a ->
read src2 >>= (maybe (unRead a src1 >> return Nothing) $ \b ->
return $! Just $! f a b ) )
zipWithM :: (a -> b -> IO c)
-> InputStream a
-> InputStream b
-> IO (InputStream c)
zipWithM f src1 src2 = makeInputStream src
where
src = read src1 >>= (maybe (return Nothing) $ \a ->
read src2 >>= (maybe (unRead a src1 >> return Nothing) $ \b ->
f a b >>= \c -> return $! Just $! c ) )
filterOutput :: (a -> Bool) -> OutputStream a -> IO (OutputStream a)
filterOutput p output = makeOutputStream chunk
where
chunk Nothing = write Nothing output
chunk ch@(Just x) = when (p x) $ write ch output
filterOutputM :: (a -> IO Bool) -> OutputStream a -> IO (OutputStream a)
filterOutputM p output = makeOutputStream chunk
where
chunk Nothing = write Nothing output
chunk ch@(Just x) = do
b <- p x
if b then write ch output else return $! ()
unzip :: InputStream (a, b) -> IO (InputStream a, InputStream b)
unzip os = do
lock <- newMVar $! ()
buf1 <- newIORef id
buf2 <- newIORef id
is1 <- makeInputStream $ src lock id buf1 buf2
is2 <- makeInputStream $ src lock twist buf2 buf1
return (is1, is2)
where
twist (a, b) = (b, a)
src lock proj myBuf theirBuf = withMVar lock $ const $ do
dl <- readIORef myBuf
case dl [] of
[] -> more
(x:xs) -> writeIORef myBuf (xs++) >> (return $! Just x)
where
more = read os >>=
maybe (return Nothing)
(\x -> do
let (a, b) = proj x
modifyIORef theirBuf (. (b:))
return $! Just a)
take :: Int64 -> InputStream a -> IO (InputStream a)
take k0 input = do
kref <- newIORef k0
return $! InputStream (prod kref) (pb kref)
where
prod kref = do
!k <- readIORef kref
if k <= 0
then return Nothing
else do
m <- read input
when (isJust m) $ modifyIORef kref $ \x -> x - 1
return m
pb kref !s = do
unRead s input
modifyIORef kref (+1)
drop :: Int64 -> InputStream a -> IO (InputStream a)
drop k0 input = do
kref <- newIORef k0
return $! InputStream (prod kref) (pb kref)
where
prod kref = do
!k <- readIORef kref
if k <= 0
then getInput kref
else discard kref
getInput kref = do
read input >>= maybe (return Nothing) (\c -> do
modifyIORef kref (\x -> x - 1)
return $! Just c)
discard kref = getInput kref >>= maybe (return Nothing) (const $ prod kref)
pb kref s = do
unRead s input
modifyIORef kref (+1)
give :: Int64 -> OutputStream a -> IO (OutputStream a)
give k output = newIORef k >>= makeOutputStream . chunk
where
chunk ref = maybe (return $! ()) $ \x -> do
!n <- readIORef ref
if n <= 0
then return $! ()
else do
writeIORef ref $! n - 1
write (Just x) output
ignore :: Int64 -> OutputStream a -> IO (OutputStream a)
ignore k output = newIORef k >>= makeOutputStream . chunk
where
chunk ref = maybe (return $! ()) $ \x -> do
!n <- readIORef ref
if n > 0
then writeIORef ref $! n - 1
else write (Just x) output
ignoreEof :: OutputStream a -> IO (OutputStream a)
ignoreEof s = makeOutputStream f
where
f Nothing = return $! ()
f x = write x s
atEndOfInput :: IO b -> InputStream a -> IO (InputStream a)
atEndOfInput m is = return $! InputStream prod pb
where
prod = read is >>= maybe eof (return . Just)
eof = void m >> return Nothing
pb s = unRead s is
atEndOfOutput :: IO b -> OutputStream a -> IO (OutputStream a)
atEndOfOutput m os = makeOutputStream f
where
f Nothing = write Nothing os >> void m
f x = write x os