Await.hs

{-# LANGUAGE DeriveFunctor, ExistentialQuantification, FlexibleContexts, FlexibleInstances, GeneralizedNewtypeDeriving,
             TypeApplications, KindSignatures, ScopedTypeVariables, LambdaCase, MultiParamTypeClasses, StandaloneDeriving, TypeOperators, UndecidableInstances #-}
module Control.Effect.Await
  ( Await (..)
  , await
  , awaits
  , peek
  , runAwaitList
  ) where

import Control.Effect.Carrier
import Control.Effect.State
import Control.Effect.Sum
import Data.Coerce

data Await i (m :: * -> *) k
  = Awaits (Maybe i -> k)
  | Peek (Maybe i -> k)
    deriving Functor

instance HFunctor (Await i) where hmap _ = coerce

instance Effect (Await i) where
  handle state handler (Awaits k) = Awaits (handler . (<$ state) . k)
  handle state handler (Peek k)   = Peek (handler . (<$ state) . k)

awaits :: (Member (Await i) sig, Carrier sig m) => m (Maybe i)
awaits = send (Awaits pure)

await :: (Member (Await i) sig, Carrier sig m) => m i
await = awaits >>= maybe await pure

peek :: (Member (Await i) sig, Carrier sig m) => m (Maybe i)
peek = send (Peek pure)

newtype AwaitListC i m a = AwaitListC { runAwaitListC :: StateC [i] m a }
  deriving (Applicative, Functor, Monad)

instance forall i sig m . (Effect sig, Carrier sig m) => Carrier (Await i :+: sig) (AwaitListC i m) where
  eff (R other) = AwaitListC (eff (R (handleCoercible other)))
  eff (L act) = do
    st <- AwaitListC $ get @[i]
    case (act, st) of
      (Awaits k, [])     -> k Nothing
      (Peek k,   [])     -> k Nothing
      (Awaits k, (a:as)) -> AwaitListC (put as) *> k (Just a)
      (Peek k,   (a:_))  -> k (Just a)

runAwaitList :: forall i a m . Functor m => [i] -> AwaitListC i m a -> m a
runAwaitList is = evalState is . runAwaitListC

Main.hs

{-# LANGUAGE FlexibleContexts, RankNTypes, TypeApplications, LambdaCase, ScopedTypeVariables #-}

module Main where

import           Control.Effect
import           Control.Effect.Await
import           Control.Effect.Yield
import           Test.Hspec

main :: IO ()
main = hspec $
  describe "await + yield" $ do
    it "should filter out odd numbers, returning fast for a reversed list" $ do
      let (inp :: [Int]) = [1..50]
      let (res :: [Int]) = run . execYieldCons @Int @[Int] [] . runAwaitList inp $ go
            where go = do
                    mx <- awaits @Int
                    case mx of
                      Nothing -> pure ()
                      Just x | even x -> yield x *> go
                      _               -> go
      res `shouldSatisfy` all even
      res `shouldSatisfy` not . null

Yield.hs

{-# LANGUAGE DeriveFunctor, ExistentialQuantification, FlexibleContexts, FlexibleInstances, GeneralizedNewtypeDeriving,
             KindSignatures, LambdaCase, MultiParamTypeClasses, StandaloneDeriving, TypeOperators,
             UndecidableInstances, ScopedTypeVariables, TypeApplications #-}

module Control.Effect.Yield where

import           Control.Effect.Carrier
import           Control.Effect.Reader
import           Control.Effect.State
import           Control.Effect.Sum
import           Data.Coerce
import           Data.Semigroup.Reducer (Reducer)
import qualified Data.Semigroup.Reducer as Reducer

data Yield o (m :: * -> *) k
  = Yield o k
    deriving Functor

instance HFunctor (Yield o) where hmap _ = coerce

instance Effect (Yield o) where
  handle state handler (Yield o k) = Yield o (handler . (<$ state) $ k)

yield :: (Member (Yield o) sig, Carrier sig m) => o -> m ()
yield o = send (Yield o (pure ()))

newtype YieldC o r (m :: * -> *) a = YieldC { unYieldC :: ReaderC (o -> r -> r) (StateC r m) a }
  deriving (Functor, Applicative, Monad)

instance forall o r sig m . (Effect sig, Carrier sig m) => Carrier (Yield o :+: sig) (YieldC o r m) where
  eff (R other)       = YieldC (eff (R (R (handleCoercible other))))
  eff (L (Yield o k)) = YieldC (ask @(o -> r -> r) >>= modify . ($ o)) *> k

-- | @runYieldAppend mempty mappend@ is equivalent to 'runWriter'.
runYieldAppend :: r -> (o -> r -> r) -> YieldC o r m a -> m (r, a)
runYieldAppend zero app = runState zero . runReader app . unYieldC

runYieldCons :: Reducer o r => r -> YieldC o r m a -> m (r, a)
runYieldCons start = runYieldAppend start Reducer.cons

execYieldCons :: forall o r m a . (Functor m, Reducer o r) => r -> YieldC o r m a -> m r
execYieldCons start = fmap fst . runYieldCons start

runYieldSnoc :: Reducer o r => r -> YieldC o r m a -> m (r, a)
runYieldSnoc start = runYieldAppend start (flip Reducer.snoc)

runYieldFold :: Monoid r => (o -> r) -> YieldC o r m a -> m (r, a)
runYieldFold fm = runYieldAppend mempty (\o r -> mappend (fm o) r)

newtype YieldIgnoreC o (m :: * -> *) a = YieldIgnoreC { runYieldIgnoreC :: m a }
  deriving (Applicative, Functor, Monad)

instance Carrier sig m => Carrier (Yield o :+: sig) (YieldIgnoreC o m) where
  eff (R other)       = YieldIgnoreC (eff (handleCoercible other))
  eff (L (Yield _ k)) = k

runYieldIgnoring :: YieldIgnoreC o m a -> m a
runYieldIgnoring = runYieldIgnoreC