{-# LANGUAGE TemplateHaskell #-}
module Polysemy.Output
(
Output (..)
, output
, runOutputAsList
, runFoldMapOutput
, runIgnoringOutput
, runBatchOutput
) where
import Data.Bifunctor (first)
import Polysemy
import Polysemy.State
import Control.Monad (when)
data Output o m a where
Output :: o -> Output o m ()
makeSem ''Output
runOutputAsList
:: forall o r a
. Sem (Output o ': r) a
-> Sem r ([o], a)
runOutputAsList = fmap (first reverse) . runState [] . reinterpret
(\case
Output o -> modify (o :)
)
{-# INLINE runOutputAsList #-}
runFoldMapOutput
:: forall o m r a
. Monoid m
=> (o -> m)
-> Sem (Output o ': r) a
-> Sem r (m, a)
runFoldMapOutput f = runState mempty . reinterpret
(\case
Output o -> modify (`mappend` f o)
)
{-# INLINE runFoldMapOutput #-}
runIgnoringOutput :: Sem (Output o ': r) a -> Sem r a
runIgnoringOutput = interpret $ \case
Output _ -> pure ()
{-# INLINE runIgnoringOutput #-}
runBatchOutput
:: forall o r a
. Int
-> Sem (Output o ': r) a
-> Sem (Output [o] ': r) a
runBatchOutput 0 m = raise $ runIgnoringOutput m
runBatchOutput size m = do
((c, res), a) <-
runState (0 :: Int, [] :: [o]) $ reinterpret2 (\case
Output o -> do
(count, acc) <- get
let newCount = 1 + count
newAcc = o : acc
if newCount < size
then put (newCount, newAcc)
else do
output (reverse newAcc)
put (0 :: Int, [] :: [o])
) m
when (c > 0) $ output (reverse res)
pure a