{-# LANGUAGE NoImplicitPrelude #-}
{-# LANGUAGE LambdaCase #-}
{-# LANGUAGE ScopedTypeVariables #-}

module DSV.Pipes
  ( Pipe, Producer, Consumer, Effect
  , yield, await, (>->), lift, runEffect
  , count
  ) where

import DSV.Numbers
import DSV.Prelude

-- pipes
import Pipes (Proxy, lift, next, await, yield, runEffect)
import qualified Pipes

-- | Use @>->@ to build an 'Effect': a pipeline consisting of a 'Producer' at the beginning, any number of 'Pipe's in the middle, and a 'Consumer' at the end.
(>->) :: Monad m => Proxy a' a () b m r -> Proxy () b c' c m r -> Proxy a' a c' c m r
>-> :: Proxy a' a () b m r -> Proxy () b c' c m r -> Proxy a' a c' c m r
(>->) = Proxy a' a () b m r -> Proxy () b c' c m r -> Proxy a' a c' c m r
forall (m :: * -> *) a' a b r c' c.
Functor m =>
Proxy a' a () b m r -> Proxy () b c' c m r -> Proxy a' a c' c m r
(Pipes.>->)

type Producer b m r = Pipes.Producer b m r

type Pipe a b m r = Pipes.Pipe a b m r

type Consumer a m r = Pipes.Consumer a m r

type Effect m r = Pipes.Effect m r

count ::
    forall a m r .
    Monad m
    => Producer a m r
    -> Producer a m (Natural, r)

count :: Producer a m r -> Producer a m (Natural, r)
count = Natural -> Producer a m r -> Producer a m (Natural, r)
forall (m :: * -> *) t a b x' x.
(Monad m, Num t) =>
t -> Producer a m b -> Proxy x' x () a m (t, b)
go Natural
0
  where
    go :: t -> Producer a m b -> Proxy x' x () a m (t, b)
go t
n Producer a m b
p =
        m (Either b (a, Producer a m b))
-> Proxy x' x () a m (Either b (a, Producer a m b))
forall (t :: (* -> *) -> * -> *) (m :: * -> *) a.
(MonadTrans t, Monad m) =>
m a -> t m a
lift (Producer a m b -> m (Either b (a, Producer a m b))
forall (m :: * -> *) a r.
Monad m =>
Producer a m r -> m (Either r (a, Producer a m r))
next Producer a m b
p) Proxy x' x () a m (Either b (a, Producer a m b))
-> (Either b (a, Producer a m b) -> Proxy x' x () a m (t, b))
-> Proxy x' x () a m (t, b)
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>=
        \case
            Left b
r -> (t, b) -> Proxy x' x () a m (t, b)
forall (m :: * -> *) a. Monad m => a -> m a
return (t
n, b
r)
            Right (a
a, Producer a m b
p') -> a -> Proxy x' x () a m ()
forall (m :: * -> *) a x' x. Functor m => a -> Proxy x' x () a m ()
yield a
a Proxy x' x () a m ()
-> Proxy x' x () a m (t, b) -> Proxy x' x () a m (t, b)
forall (f :: * -> *) a b. Applicative f => f a -> f b -> f b
*> t -> Producer a m b -> Proxy x' x () a m (t, b)
go (t
n t -> t -> t
forall a. Num a => a -> a -> a
+ t
1) Producer a m b
p'