module Data.SouSiT.Trans (
map,
mapM,
mapWithState,
zipWithIndex,
take,
takeUntil,
takeUntilEq,
takeWhile,
drop,
dropUntil,
dropWhile,
filter,
filterMap,
flatMap,
accumulate,
buffer,
count,
disperse,
andThen,
loop,
loopN,
sequence,
eitherRight,
eitherLeft,
serialize,
deserialize,
mapSinkStatus,
TransFun,
applyTransFun,
mapSinkTransFun,
applyMapping,
mapSinkMapping,
toDoneTrans,
debug
) where
import Prelude hiding (map, mapM, take, takeWhile, drop, dropWhile, sequence, filter)
import qualified Prelude as P
import Data.SouSiT.Sink
import Data.SouSiT.Transform
import Data.ByteString (ByteString)
import qualified Data.Serialize as S
import Control.Monad (liftM)
import Control.Monad.IO.Class
mapSinkStatus :: Monad m => (SinkStatus a m r -> SinkStatus b m r) -> Sink a m r -> Sink b m r
mapSinkStatus f = Sink . liftM f . sinkStatus
type TransFun a b m r = (a -> m (Sink a m r)) -> m r -> b -> m (Sink b m r)
applyTransFun :: Monad m => TransFun a b m r -> SinkStatus a m r -> SinkStatus b m r
applyTransFun _ (Done r) = Done r
applyTransFun f (Cont nf cf) = Cont (f nf cf) cf
mapSinkTransFun f = mapSinkStatus (applyTransFun f)
applyMapping :: Monad m => (Sink a m r -> Sink b m r) -> (b -> a) -> SinkStatus a m r -> SinkStatus b m r
applyMapping _ _ (Done r) = Done r
applyMapping ms mi (Cont nf cf) = Cont nf' cf
where nf' i = nf (mi i) >>= return . ms
mapSinkMapping ms mi = mapSinkStatus $ applyMapping ms mi
toDoneTrans :: Monad m => Sink a m r -> Sink a m r
toDoneTrans = mapSinkStatus fun
where fun (Done r) = Done r
fun (Cont _ r) = Done r
map :: (a -> b) -> Transform a b
map f = mapSinkMapping (map f) f
mapM :: Monad m => (b -> m a) -> Sink a m r -> Sink b m r
mapM action sink = Sink (liftM f (sinkStatus sink))
where f (Done r) = Done r
f (Cont nf cf) = Cont nf' cf
where nf' i = liftM (mapM action) (action i >>= nf)
mapWithState :: (s -> a -> (b,s)) -> s -> Transform a b
mapWithState f !s = mapSinkTransFun fun
where fun nf _ i = let (i', s') = f s i in liftM (mapWithState f s') (nf i')
zipWithIndex :: Transform a (a, Int)
zipWithIndex = mapWithState fun 0
where fun s i = ((i,s), s+1)
take :: (Num n, Ord n) => n -> Transform a a
take n | n > 0 = mapSinkMapping prev id
| otherwise = toDoneTrans
where prev = take (n 1)
takeUntil :: (a -> Bool) -> Transform a a
takeUntil p = mapSinkStatus fun
where fun s@(Done _) = s
fun (Cont nf cf) = Cont nf' cf
where nf' i = if p i then return (doneSink cf)
else liftM (takeUntil p) (nf i)
takeUntilEq :: Eq a => a -> Transform a a
takeUntilEq e = takeUntil (e ==)
takeWhile :: (a -> Bool) -> Transform a a
takeWhile f = takeUntil (not . f)
accumulate :: b -> (b -> a -> b) -> Transform a b
accumulate initAcc f = mapSinkStatus fun
where fun (Done r) = Done r
fun (Cont nf _) = step initAcc
where step !acc = Cont (return . Sink . return . step . f acc) (nf acc >>= closeSink)
count :: Num n => Transform a n
count = accumulate 0 step
where step i _ = i + 1
buffer :: Int -> b -> (b -> a -> b) -> Transform a b
buffer initN initAcc f = if initN > 0 then mapSinkStatus fun
else error $ "Cannot buffer " ++ show initN ++ " elements"
where fun (Done r) = Done r
fun (Cont nf _) = step initN initAcc
where step 1 !acc = Cont (return . nf') (nf acc >>= closeSink)
where nf' i = Sink $ nf (f acc i) >>= sinkStatus >>= return . fun
step n !acc = Cont (return . Sink . return . step (n1) . f acc) (nf acc >>= closeSink)
disperse :: Transform [a] a
disperse sink = Sink $ liftM fun (sinkStatus sink)
where fun (Done r) = Done r
fun (Cont _ cf) = Cont (liftM disperse . flip feedList sink) cf
drop :: (Num n, Ord n) => n -> Transform a a
drop n0 = mapSinkStatus fun
where fun (Done r) = Done r
fun (Cont nf cf) = Cont (step n0) cf
where step n i | n > 0 = return $ contSink (step $ n1) cf
| otherwise = nf i
dropUntil :: (a -> Bool) -> Transform a a
dropUntil p = mapSinkTransFun fun
where fun nf cf i | p i = nf i
| otherwise = return $ dropUntil p $ contSink nf cf
dropWhile :: (a -> Bool) -> Transform a a
dropWhile f = dropUntil (not . f)
filter :: (a -> Bool) -> Transform a a
filter p = mapSinkTransFun fun
where fun nf cf i | p i = liftM (filter p) (nf i)
| otherwise = return $ filter p $ contSink nf cf
filterMap :: (a -> Maybe b) -> Transform a b
filterMap f = mapSinkTransFun fun
where fun nf cf i = case f i of
(Just i') -> liftM (filterMap f) (nf i')
Nothing -> return $ filterMap f $ contSink nf cf
andThen :: Transform a b -> Transform a b -> Transform a b
andThen t1 t2 sink = Sink $ (>>= f) $ sinkStatus sink
where f (Done r) = return $ Done r
f (Cont _ _) = sinkStatus $ sinkUnwrap t2 $ t1 $ sinkWrap sink
data WrapRes i m r = SinkIsDone (m r)
| SinkIsCont (i -> m (Sink i m r)) (m r)
sinkUnwrap :: Monad m => Transform a b -> Sink a m (WrapRes b m r) -> Sink a m r
sinkUnwrap t = Sink . (>>= handle) . sinkStatus
where handle (Cont nf cf) = return $ Cont (liftM (sinkUnwrap t) . nf) (cf >>= unwrapRes)
handle (Done r) = liftM (t . recSink) r >>= sinkStatus
sinkWrap :: Monad m => Sink i m r -> Sink i m (WrapRes i m r)
sinkWrap = Sink . liftM f . sinkStatus
where f (Done r) = Done $ return $ SinkIsDone r
f (Cont nf cf) = Cont (liftM sinkWrap . nf) (return $ SinkIsCont nf cf)
recSink :: Monad m => WrapRes i m r -> Sink i m r
recSink (SinkIsDone r) = doneSink r
recSink (SinkIsCont nf cf) = contSink nf cf
unwrapRes :: Monad m => WrapRes i m r -> m r
unwrapRes (SinkIsDone r) = r
unwrapRes (SinkIsCont _ r) = r
sequence :: [Transform a b] -> Transform a b
sequence [] = error "No Transform for T.sequence"
sequence (t:[]) = t
sequence (t:ts) = andThen t (sequence ts)
loop :: Transform a b -> Transform a b
loop t = andThen t (loop t)
loopN :: Int -> Transform a b -> Transform a b
loopN n t | n > 1 = andThen t $ loopN (n 1) t
| n == 1 = t
| otherwise = error $ "Invalid n=" ++ show n ++ " in T.loopN"
flatMap :: (a -> [b]) -> Transform a b
flatMap f = map f =$= disperse
eitherRight :: Transform (Either a b) b
eitherRight = mapSinkStatus f
where f (Done r) = Done r
f (Cont nf cf) = Cont (liftM eitherRight . handle) cf
where handle = either ignore nf
ignore _ = return $ contSink nf cf
eitherLeft :: Transform (Either a b) a
eitherLeft = mapSinkStatus f
where f (Done r) = Done r
f (Cont nf cf) = Cont (liftM eitherLeft . handle) cf
where handle = either nf ignore
ignore _ = return $ contSink nf cf
debug :: (Show a, MonadIO m) => String -> Sink a m r -> Sink a m r
debug label sink = mapM f sink
where f i = (liftIO . putStrLn) (label ++ ": " ++ show i) >> return i
serialize :: S.Serialize a => Transform a ByteString
serialize = map S.encode
deserialize :: S.Serialize b => Transform ByteString b
deserialize = deserialize' []
deserialize' :: S.Serialize b => [ByteString -> S.Result b] -> Transform ByteString b
deserialize' ips = mapSinkTransFun fun
where fun nf cf i = liftM (deserialize' ps') $ feedList bs (contSink nf cf)
where (ps', bs) = tryToParse (S.runGetPartial S.get:ips) i
tryToParse [] _ = ([],[])
tryToParse (p:ps) i = case p i of
(S.Done b _) -> (ps', b:bs)
(S.Fail _) -> (ps', bs)
(S.Partial p') -> (p':ps,bs)
where (ps', bs) = tryToParse ps i