clojure transducers in haskell

transducers.hs

{-# LANGUAGE GeneralizedNewtypeDeriving #-}
{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE ExistentialQuantification #-}
{-# LANGUAGE MultiWayIf #-}
module Transducer (GenericTDAction, TDAction, StateTDAction,
                   Transducer, stop, transducer, stateTransducer, writeState,
                   readState, returnStream, Transducer.map, Transducer.filter,
                   Transducer.take, Transducer.cat, Transducer.mapcat,
                   Transducer.remove, Transducer.takeWhile, Transducer.takeNth,
                   Transducer.drop, Transducer.dropWhile) where
import Streamly.Prelude
import Streamly
import qualified Control.Category as C
import Control.Monad.Trans.Maybe
import Control.Monad.State
import Control.Monad.Trans
import Control.Monad.Reader
import Data.IORef
import Control.Monad ((>=>))
import Control.Applicative

newtype GenericTDAction stateRef m a = GenericTDAction
  { getTDAction :: ReaderT stateRef (SerialT (MaybeT m)) a }
  deriving (Monad, MonadIO, Functor, Applicative)
type TDAction m a = GenericTDAction () m a
type StateTDAction s m a = GenericTDAction (IORef s) m a

data Transducer m a b =
  forall s . Transducer (m s) (a -> GenericTDAction s m b)

stop :: Monad m => GenericTDAction s m a
stop = GenericTDAction $ lift $ lift $ MaybeT $ pure Nothing

instance Monad m => C.Category (Transducer m) where
  id = Transducer (pure ()) pure
  Transducer initS2 g . Transducer initS1 f =
    Transducer (liftA2 (,) initS1 initS2) $
    \x -> GenericTDAction $ ReaderT $
    \(s1, s2) -> do y <- runReaderT (getTDAction (f x)) s1
                    runReaderT (getTDAction (g y)) s2
    
transducer :: Monad m => (a -> TDAction m b) -> Transducer m a b
transducer f = Transducer (pure ()) f

stateTransducer :: (Monad m, MonadIO m)
                => s 
                -> (a -> StateTDAction s m b)
                -> Transducer m a b
stateTransducer initState f =
  Transducer (liftIO $ newIORef initState) f

writeState :: MonadIO m => s -> StateTDAction s m ()
writeState a = GenericTDAction $ do
  ref <- ask
  liftIO $ writeIORef ref a

readState :: MonadIO m => StateTDAction s m s
readState = GenericTDAction $ ask >>= liftIO . readIORef

returnStream :: Monad m => SerialT (MaybeT m) a -> GenericTDAction s m a
returnStream = GenericTDAction . lift

map :: Monad m => (a -> b) -> Transducer m a b
map f = transducer $ pure . f

filter :: Monad m => (a -> Bool) -> Transducer m a a
filter p = transducer $ \x ->
  if p x then pure x else returnStream nil

take :: (Monad m, MonadIO m) => Int -> Transducer m a a
take initN = stateTransducer initN $ \x -> do
  n <- readState
  if n == (0 :: Int)
    then stop
    else do writeState (n-1)
            pure x

cat :: Monad m => Transducer m [a] a
cat = transducer $ returnStream . fromList

mapcat :: Monad m => (a -> [b]) -> Transducer m a b
mapcat f = Transducer.map f C.>>> cat

remove :: Monad m => (a -> Bool) -> Transducer m a a
remove p = Transducer.filter $ not . p

takeWhile :: Monad m => (a -> Bool) -> Transducer m a a
takeWhile p = transducer $ \x ->
  if p x then pure x else stop

takeNth :: MonadIO m => Int -> Transducer m a a
takeNth initN = stateTransducer initN $ \x ->
  do n <- readState
     if | n > 0 -> do writeState (n-1); returnStream nil
        | n == 0 -> do writeState (n-1); pure x
        | otherwise -> stop

drop :: MonadIO m => Int -> Transducer m a a
drop initN = stateTransducer initN $ \x ->
  do n <- readState
     if n > 0
       then do writeState (n-1)
               pure x
       else stop

dropWhile p = stateTransducer False $ \x ->
  do dropped <- readState
     if | dropped -> pure x
        | not (p x) -> do writeState True
                          pure x
        | otherwise -> returnStream nil