{-# LANGUAGE TemplateHaskell #-}
module Polysemy.Output
(
Output (..)
, output
, runOutputList
, runOutputMonoid
, ignoreOutput
, runOutputBatched
) 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
runOutputList
:: forall o r a
. Sem (Output o ': r) a
-> Sem r ([o], a)
runOutputList = fmap (first reverse) . runState [] . reinterpret
(\case
Output o -> modify (o :)
)
{-# INLINE runOutputList #-}
runOutputMonoid
:: forall o m r a
. Monoid m
=> (o -> m)
-> Sem (Output o ': r) a
-> Sem r (m, a)
runOutputMonoid f = runState mempty . reinterpret
(\case
Output o -> modify (`mappend` f o)
)
{-# INLINE runOutputMonoid #-}
ignoreOutput :: Sem (Output o ': r) a -> Sem r a
ignoreOutput = interpret $ \case
Output _ -> pure ()
{-# INLINE ignoreOutput #-}
runOutputBatched
:: forall o r a
. Member (Output [o]) r
=> Int
-> Sem (Output o ': r) a
-> Sem r a
runOutputBatched 0 m = ignoreOutput m
runOutputBatched size m = do
((c, res), a) <-
runState (0 :: Int, [] :: [o]) $ reinterpret (\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 @[o] (reverse res)
pure a