monadic-partition.hs

{-# LANGUAGE Rank2Types #-}
import qualified Data.IntMap as IM
import Data.IORef
import Control.Monad.Writer
import Data.Foldable as F
import Data.Monoid
import Control.Applicative
import Control.Lens

newtype Builder a m = Builder { unBuilder :: forall r. (a -> m r) -> m r }

builder :: Monad m => m a -> Builder a m
builder m = Builder (m >>=)

getBuilder :: Applicative m => Builder a m -> m a
getBuilder (Builder k) = k pure

instance Monoid a => Monoid (Builder a m) where
  mempty = Builder $ \c -> c mempty
  mappend (Builder m) (Builder n) = Builder $ \c -> m $ \a -> n $ \b -> c (mappend a b)

partitionRef :: Ord a => a -> IM.IntMap (IORef a)
  -> IO (IM.IntMap (IORef a), IM.IntMap (IORef a))
partitionRef a m = do
  (Endo l, Endo r) <- getBuilder $ IM.foldMapWithKey go m
  return (IM.fromAscList $ l [], IM.fromAscList $ r [])
    where
      go i r = builder $ mpartition (Endo ((i, r):)) <$> (<=a) <$> readIORef r

mpartition :: (Monoid a) => a -> Bool -> (a, a)
mpartition a False = (a, mempty)
mpartition a True = (mempty, a)

partitionRef' :: Ord a => a -> IM.IntMap (IORef a)
  -> IO (IM.IntMap (IORef a), IM.IntMap (IORef a))
partitionRef' a m = getBuilder (IM.foldMapWithKey go m) where
  go i r = builder $ mpartition (IM.singleton i r) <$> (<=a) <$> readIORef r