batch-mtl.hs

{-# LANGUAGE FlexibleInstances, MultiParamTypeClasses, GeneralizedNewtypeDeriving, FlexibleContexts #-}

module Lib where

import Control.Monad.State
import Control.Monad.Writer
import Control.Monad.IO.Class
import Control.Monad.Reader

someFunc :: IO ()
someFunc = print =<< runBatchT 50 writin

writin :: MonadWriter [Int] m => m ()
writin = do
    forM_ [0..1000] $ \i -> do
        tell [i]

newtype BatchT w m a = BatchT { unBatchT :: StateT (Int, [[w]]) m a }
    deriving (Functor, Applicative, Monad, MonadIO)

runBatchT :: Monad m => Int -> BatchT w m a -> m ([[w]], a)
runBatchT limit (BatchT s) = do
    (a, (_, r)) <- runStateT s (limit, mempty)
    pure (r, a)

instance (Monoid w, Monad m) => MonadWriter w (BatchT w m) where
    tell x = BatchT $ do
        (limit, logs) <- get
        case logs of
            [] ->
                put (limit, [[x]])
            (l:ls)
                | length l > limit ->
                    put (limit, [x] : (l : ls))
                | otherwise ->
                    put (limit, (x : l) : ls)