Another Pipe experiment

pipes.hs

{-# LANGUAGE CPP, GADTs, KindSignatures, TypeFamilies, DataKinds #-}

import Control.Monad

type family If (t :: Bool) a b
type instance If True a b = a
type instance If False a b = b

data Pipe a b (t :: Bool) m r where
  Pure :: r -> Pipe a b t m r
  Await :: (Maybe a -> Pipe a b t m r) -> Pipe a b t m r
  Yield :: b -> Pipe a b t m r -> Pipe a b t m r
  Do :: m (Pipe a b t m r) -> Pipe a b t m r
  -- can defer upstream only if t == True
  Defer :: Pipe a b True m r

instance Monad m => Monad (Pipe a b t m) where
  return = Pure
  Pure r >>= f = f r
  Await k >>= f = Await (\a -> k a >>= f)
  Yield x p >>= f = Yield x (p >>= f)
  Do m >>= f = Do (m >>= \p -> return (p >>= f))
  Defer >>= _ = Defer

tryAwait :: Pipe a b t m (Maybe a)
tryAwait = Await Pure

await :: Pipe a b True m a
await = Await $ maybe Defer Pure

yield :: b -> Pipe a b t m ()
yield x = Yield x (Pure ())

lift :: Monad m => m r -> Pipe a b t m r
lift m = Do (liftM Pure m)

compose :: (Monad m, s ~ If t r s)
        => Pipe a b u m r
        -> Pipe b c t m s
        -> Pipe a c u m s
compose p1 p2 = case (p1, p2) of
  (_, Pure r) -> Pure r
  (Pure r, Defer) -> Pure r

  (_, Do m) -> Do (m >>= \p2' -> return (compose p1 p2'))
  (_, Yield x p2') -> Yield x (compose p1 p2')

  (Pure r, Await k) -> compose p1 (k Nothing)
  (Do m, _) -> Do (m >>= \p1' -> return (compose p1' p2))
  (Await k, _) -> Await (\x -> compose (k x) p2)

  (Yield x p1', Await k) -> compose p1' (k (Just x))
  (Yield x p1', Defer) -> compose p1' p2

  (Defer, _) -> Defer

(>+>) :: Monad m
      => Pipe a b u m r
      -> Pipe b c True m r
      -> Pipe a c u m r
(>+>) = compose

($$) :: Monad m
     => Pipe a b u m r
     -> Pipe b c False m s
     -> Pipe a c u m s
($$) = compose

runPipe :: Monad m => Pipe () x False m r -> m r
runPipe (Pure r) = return r
runPipe (Await k) = runPipe $ k (Just ())
runPipe (Yield x p) = runPipe p
runPipe (Do m) = m >>= runPipe

pipe :: Monad m => (a -> b) -> Pipe a b t m r
pipe f = forever $ tryAwait >>= maybe (pipe f) (yield . f)

idP :: Monad m => Pipe a a t m r
idP = pipe id

--

fold :: Monad m => (a -> b -> b) -> b -> Pipe a b t m b
fold f z = tryAwait >>= \x -> case x of
  Nothing -> return z
  Just a -> let z' = f a z
            in z' `seq` fold f z'

main :: IO ()
main = do
  r <- runPipe $ mapM_ yield [1 .. 10 :: Int] $$ (idP $$ fold (+) 0)
  print r