add ListT monad transformer and relevant instances

Author: chessai

src/Streaming/Prelude.hs

@@ -53,9 +53,12 @@
 {-# LANGUAGE DeriveFoldable #-}
 {-# LANGUAGE DeriveFunctor #-}
 {-# LANGUAGE DeriveTraversable #-}
+{-# LANGUAGE FlexibleInstances #-}
+{-# LANGUAGE MultiParamTypeClasses #-}
 {-# LANGUAGE RankNTypes #-}
 {-# LANGUAGE ScopedTypeVariables #-}
 {-# LANGUAGE TypeFamilies #-}
+{-# LANGUAGE UndecidableInstances #-}
 
 {-# OPTIONS_GHC -Wall #-}
 
@@ -255,15 +258,26 @@ module Streaming.Prelude (
 
     -- * Basic Type
     , Stream
+
+    -- * ListT
+    , ListT(..)
+    , runListT
   ) where
 import Streaming.Internal
 
 import Control.Monad hiding (filterM, mapM, mapM_, foldM, foldM_, replicateM, sequence)
 import Data.Functor.Identity
 import Data.Functor.Sum
 import Control.Monad.Trans
-import Control.Applicative (Applicative (..))
+import Control.Applicative (Applicative (..), Alternative (..))
+import Control.Monad.Morph
+import Control.Monad.Error.Class
+import Control.Monad.Reader.Class
+import Control.Monad.State.Class
+import Control.Monad.Writer.Class
+import Control.Monad.Zip
 import Data.Functor (Functor (..), (<$))
+import Data.Semigroup (Semigroup ((<>)))
 
 import qualified Prelude as Prelude
 import qualified Data.Foldable as Foldable
@@ -2921,3 +2935,173 @@ mapMaybeM phi = loop where
         Nothing -> loop snext
         Just b -> Step (b :> loop snext)
 {-#INLINABLE mapMaybeM #-}
+
+{-| The list monad transformer.
+    'pure' and 'return' correspond to 'yield', yielding a single value.
+    ('>>=') corresponds to 'for', calling the second computation once for
+    each time the first computation 'yield's.
+-}
+newtype ListT m a = Select { enumerate :: Stream (Of a) m () }
+
+instance Monad m => Functor (ListT m) where
+  fmap f p = Select (for (enumerate p) (\a -> yield (f a)))
+  {-# INLINE fmap #-}
+
+instance Monad m => Applicative (ListT m) where
+  pure a = Select (yield a)
+  {-# INLINE pure #-}
+  mf <*> mx = Select (
+    for (enumerate mf) (\f ->
+    for (enumerate mx) (\x ->
+    yield (f x) ) ) )
+
+instance Monad m => Monad (ListT m) where
+  return = pure
+  {-# INLINE return #-}
+  m >>= f = Select (for (enumerate m) (\a -> enumerate (f a)))
+  {-# INLINE (>>=) #-}
+
+instance Foldable m => Foldable (ListT m) where
+  foldMap f = go . enumerate
+    where
+      go p = case p of
+        Return () -> mempty
+        Effect m -> Foldable.foldMap go m
+        Step (a :> rest) -> f a `mappend` go rest
+  {-# INLINE foldMap #-}
+
+instance (Monad m, Traversable m) => Traversable (ListT m) where
+  traverse k (Select p) = fmap Select (traverse_ p)
+    where
+      traverse_ (Return ()) = pure (Return ())
+      traverse_ (Effect m)  = fmap Effect (traverse traverse_ m)
+      traverse_ (Step (a :> rest)) = (\a_ rest_ -> Step (a_ :> rest_)) <$> k a <*> traverse_ rest
+
+instance MonadTrans ListT where
+  lift m = Select (do
+    a <- lift m
+    yield a )
+
+instance MonadIO m => MonadIO (ListT m) where
+  liftIO m = lift (liftIO m)
+  {-# INLINE liftIO #-}
+
+instance Monad m => Alternative (ListT m) where
+  empty = Select (pure ())
+  {-# INLINE empty #-}
+  p1 <|> p2 = Select (do
+    enumerate p1
+    enumerate p2 )
+
+instance Monad m => MonadPlus (ListT m) where
+  mzero = empty
+  {-# INLINE mzero #-}
+  mplus = (<|>)
+  {-# INLINE mplus #-}
+
+instance MFunctor ListT where
+  hoist morph = Select . hoist morph . enumerate
+  {-# INLINE hoist #-}
+
+instance MMonad ListT where
+  embed f (Select p0) = Select (loop p0)
+    where
+      loop (Return ()) = Return ()
+      loop (Effect m)  = for (enumerate (fmap loop (f m))) id
+      loop (Step (a :> rest)) = Step (a :> loop rest)
+  {-# INLINE embed #-}
+
+instance Monad m => Semigroup (ListT m a) where
+  (<>) = (<|>)
+  {-# INLINE (<>) #-}
+
+instance Monad m => Monoid (ListT m a) where
+  mempty = empty
+  {-# INLINE mempty #-}
+#if !(MIN_VERSION_base(4,11,0))
+  mappend = (<|>)
+  {-# INLINE mappend #-}
+#endif
+
+instance (MonadState s m) => MonadState s (ListT m) where
+  get = lift get
+  {-# INLINE get #-}
+
+  put s = lift (put s)
+  {-# INLINE put #-}
+
+  state f = lift (state f)
+  {-# INLINE state #-}
+
+instance (MonadWriter w m) => MonadWriter w (ListT m) where
+  writer = lift . writer
+  {-# INLINE writer #-}
+
+  tell w = lift (tell w)
+  {-# INLINE tell #-}
+
+  --listen :: ListT m a -> ListT m (a, w)
+  listen l = Select (go (enumerate l) mempty)
+    where
+      go p w = case p of
+        Return () -> Return ()
+        Effect m  -> Effect (do
+          (p', w') <- listen m
+          pure (go p' $! mappend w w') )
+        Step (a :> rest) -> Step ( (a,w) :> go rest w)
+
+  pass l = Select (go (enumerate l) mempty)
+    where
+      --go :: forall m a w. Stream (Of (w, a)) m () -> (w -> w) -> Stream (Of a) m ()
+      go p w = case p of
+        Return () -> Return ()
+        Effect m  -> Effect (do
+          (p', w') <- listen m
+          pure (go p' $! mappend w w'))
+        Step ((b, f) :> rest) -> Effect (pass (return (Step (b :> (go rest (f w))), \_ -> f w) ))
+
+instance (MonadReader i m) => MonadReader i (ListT m) where
+  ask = lift ask
+  {-# INLINE ask #-}
+
+  local f l = Select (local f (enumerate l))
+  {-# INLINE local #-}
+
+  reader f = lift (reader f)
+  {-# INLINE reader #-}
+
+instance (MonadError e m) => MonadError e (ListT m) where
+  throwError e = lift (throwError e)
+  {-# INLINE throwError #-}
+
+  catchError l k = Select (catchError (enumerate l) (\e -> enumerate (k e)))
+  {-# INLINE catchError #-}
+
+{- These instances require a dependency on `exceptions`.
+instance MonadThrow m => MonadThrow (ListT m) where
+  throwM = Select . throwM
+  {-# INLINE throwM #-}
+instance MonadCatch m => MonadCatch (ListT m) where
+  catch l k = Select (Control.Monad.Catch.catch (enumerate l) (\e -> enumerate (k e)))
+  {-# INLINE catch #-}
+-}
+
+instance Monad m => MonadZip (ListT m) where
+  mzipWith f = go
+    where
+      go xs ys = Select $ do
+        xres <- lift $ next (enumerate xs)
+        case xres of
+          Left () -> pure ()
+          Right (x, xrest) -> do
+            yres <- lift $ next (enumerate ys)
+            case yres of
+              Left () -> pure ()
+              Right (y, yrest) -> do
+                yield (f x y)
+                enumerate (go (Select xrest) (Select yrest))
+
+-- | Run a self-contained 'ListT' computation
+runListT :: Monad m => ListT m a -> m ()
+runListT l = effects (enumerate (l >> mzero))
+{-# INLINABLE runListT #-}