{-# LANGUAGE TemplateHaskell #-}
module Polysemy.Output
(
Output (..)
, output
, runOutputList
, runOutputMonoid
, runOutputMonoidAssocR
, runOutputMonoidIORef
, runOutputMonoidTVar
, ignoreOutput
, runOutputBatched
, runOutputSem
) where
import Data.IORef
import Control.Concurrent.STM
import Data.Semigroup (Endo(..))
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 #-}
runOutputMonoidAssocR
:: forall o m r a
. Monoid m
=> (o -> m)
-> Sem (Output o ': r) a
-> Sem r (m, a)
runOutputMonoidAssocR f =
fmap (first (`appEndo` mempty))
. runOutputMonoid (\a -> Endo (f a <>))
{-# INLINE runOutputMonoidAssocR #-}
runOutputMonoidIORef
:: forall o m r a
. (Monoid m, Member (Embed IO) r)
=> IORef m
-> (o -> m)
-> Sem (Output o ': r) a
-> Sem r a
runOutputMonoidIORef ref f = interpret $ \case
Output o -> embed $ atomicModifyIORef' ref (\s -> (s <> f o, ()))
{-# INLINE runOutputMonoidIORef #-}
runOutputMonoidTVar
:: forall o m r a
. (Monoid m, Member (Embed IO) r)
=> TVar m
-> (o -> m)
-> Sem (Output o ': r) a
-> Sem r a
runOutputMonoidTVar tvar f = interpret $ \case
Output o -> embed $ atomically $ do
s <- readTVar tvar
writeTVar tvar $! s <> f o
{-# INLINE runOutputMonoidTVar #-}
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
runOutputSem :: (o -> Sem r ()) -> Sem (Output o ': r) a -> Sem r a
runOutputSem act = interpret $ \case
Output o -> act o
{-# INLINE runOutputSem #-}