{-# LANGUAGE DeriveFunctor #-}
{-# LANGUAGE GeneralizedNewtypeDeriving #-}
{-# LANGUAGE TupleSections #-}

module Control.Monad.CollectionState
    ( CollectionState
    , runCollectionState
    , execCollectionState
    , collect
    )
where

import Control.Monad.State (State, runState, state)
import Data.Map (Map)
import qualified Data.Map as Map (empty)
import Data.Map.Utils (insert')

-- | A State monad that wraps a map that can only be added ("collected") to.
newtype CollectionState k v a = CollectionState (State (Map k v) a)
    deriving ((forall a b.
 (a -> b) -> CollectionState k v a -> CollectionState k v b)
-> (forall a b.
    a -> CollectionState k v b -> CollectionState k v a)
-> Functor (CollectionState k v)
forall a b. a -> CollectionState k v b -> CollectionState k v a
forall a b.
(a -> b) -> CollectionState k v a -> CollectionState k v b
forall k v a b. a -> CollectionState k v b -> CollectionState k v a
forall k v a b.
(a -> b) -> CollectionState k v a -> CollectionState k v b
forall (f :: * -> *).
(forall a b. (a -> b) -> f a -> f b)
-> (forall a b. a -> f b -> f a) -> Functor f
$cfmap :: forall k v a b.
(a -> b) -> CollectionState k v a -> CollectionState k v b
fmap :: forall a b.
(a -> b) -> CollectionState k v a -> CollectionState k v b
$c<$ :: forall k v a b. a -> CollectionState k v b -> CollectionState k v a
<$ :: forall a b. a -> CollectionState k v b -> CollectionState k v a
Functor, Functor (CollectionState k v)
Functor (CollectionState k v) =>
(forall a. a -> CollectionState k v a)
-> (forall a b.
    CollectionState k v (a -> b)
    -> CollectionState k v a -> CollectionState k v b)
-> (forall a b c.
    (a -> b -> c)
    -> CollectionState k v a
    -> CollectionState k v b
    -> CollectionState k v c)
-> (forall a b.
    CollectionState k v a
    -> CollectionState k v b -> CollectionState k v b)
-> (forall a b.
    CollectionState k v a
    -> CollectionState k v b -> CollectionState k v a)
-> Applicative (CollectionState k v)
forall a. a -> CollectionState k v a
forall k v. Functor (CollectionState k v)
forall a b.
CollectionState k v a
-> CollectionState k v b -> CollectionState k v a
forall a b.
CollectionState k v a
-> CollectionState k v b -> CollectionState k v b
forall a b.
CollectionState k v (a -> b)
-> CollectionState k v a -> CollectionState k v b
forall k v a. a -> CollectionState k v a
forall a b c.
(a -> b -> c)
-> CollectionState k v a
-> CollectionState k v b
-> CollectionState k v c
forall k v a b.
CollectionState k v a
-> CollectionState k v b -> CollectionState k v a
forall k v a b.
CollectionState k v a
-> CollectionState k v b -> CollectionState k v b
forall k v a b.
CollectionState k v (a -> b)
-> CollectionState k v a -> CollectionState k v b
forall k v a b c.
(a -> b -> c)
-> CollectionState k v a
-> CollectionState k v b
-> CollectionState k v c
forall (f :: * -> *).
Functor f =>
(forall a. a -> f a)
-> (forall a b. f (a -> b) -> f a -> f b)
-> (forall a b c. (a -> b -> c) -> f a -> f b -> f c)
-> (forall a b. f a -> f b -> f b)
-> (forall a b. f a -> f b -> f a)
-> Applicative f
$cpure :: forall k v a. a -> CollectionState k v a
pure :: forall a. a -> CollectionState k v a
$c<*> :: forall k v a b.
CollectionState k v (a -> b)
-> CollectionState k v a -> CollectionState k v b
<*> :: forall a b.
CollectionState k v (a -> b)
-> CollectionState k v a -> CollectionState k v b
$cliftA2 :: forall k v a b c.
(a -> b -> c)
-> CollectionState k v a
-> CollectionState k v b
-> CollectionState k v c
liftA2 :: forall a b c.
(a -> b -> c)
-> CollectionState k v a
-> CollectionState k v b
-> CollectionState k v c
$c*> :: forall k v a b.
CollectionState k v a
-> CollectionState k v b -> CollectionState k v b
*> :: forall a b.
CollectionState k v a
-> CollectionState k v b -> CollectionState k v b
$c<* :: forall k v a b.
CollectionState k v a
-> CollectionState k v b -> CollectionState k v a
<* :: forall a b.
CollectionState k v a
-> CollectionState k v b -> CollectionState k v a
Applicative, Applicative (CollectionState k v)
Applicative (CollectionState k v) =>
(forall a b.
 CollectionState k v a
 -> (a -> CollectionState k v b) -> CollectionState k v b)
-> (forall a b.
    CollectionState k v a
    -> CollectionState k v b -> CollectionState k v b)
-> (forall a. a -> CollectionState k v a)
-> Monad (CollectionState k v)
forall a. a -> CollectionState k v a
forall k v. Applicative (CollectionState k v)
forall a b.
CollectionState k v a
-> CollectionState k v b -> CollectionState k v b
forall a b.
CollectionState k v a
-> (a -> CollectionState k v b) -> CollectionState k v b
forall k v a. a -> CollectionState k v a
forall k v a b.
CollectionState k v a
-> CollectionState k v b -> CollectionState k v b
forall k v a b.
CollectionState k v a
-> (a -> CollectionState k v b) -> CollectionState k v b
forall (m :: * -> *).
Applicative m =>
(forall a b. m a -> (a -> m b) -> m b)
-> (forall a b. m a -> m b -> m b)
-> (forall a. a -> m a)
-> Monad m
$c>>= :: forall k v a b.
CollectionState k v a
-> (a -> CollectionState k v b) -> CollectionState k v b
>>= :: forall a b.
CollectionState k v a
-> (a -> CollectionState k v b) -> CollectionState k v b
$c>> :: forall k v a b.
CollectionState k v a
-> CollectionState k v b -> CollectionState k v b
>> :: forall a b.
CollectionState k v a
-> CollectionState k v b -> CollectionState k v b
$creturn :: forall k v a. a -> CollectionState k v a
return :: forall a. a -> CollectionState k v a
Monad)

runCollectionState :: CollectionState k v a -> (a, Map k v)
runCollectionState :: forall k v a. CollectionState k v a -> (a, Map k v)
runCollectionState (CollectionState State (Map k v) a
m) = State (Map k v) a -> Map k v -> (a, Map k v)
forall s a. State s a -> s -> (a, s)
runState State (Map k v) a
m Map k v
forall k a. Map k a
Map.empty

execCollectionState :: CollectionState k v () -> Map k v
execCollectionState :: forall k v. CollectionState k v () -> Map k v
execCollectionState = ((), Map k v) -> Map k v
forall a b. (a, b) -> b
snd (((), Map k v) -> Map k v)
-> (CollectionState k v () -> ((), Map k v))
-> CollectionState k v ()
-> Map k v
forall b c a. (b -> c) -> (a -> b) -> a -> c
. CollectionState k v () -> ((), Map k v)
forall k v a. CollectionState k v a -> (a, Map k v)
runCollectionState

-- | Add an entry to the collection iff the key is not yet in the collection.
--   The argument state transformation, which returns a value, is only run if
--   the value is to be added.
collect :: (Ord k) => k -> CollectionState k v v -> CollectionState k v ()
collect :: forall k v.
Ord k =>
k -> CollectionState k v v -> CollectionState k v ()
collect k
k (CollectionState State (Map k v) v
mv) = State (Map k v) () -> CollectionState k v ()
forall k v a. State (Map k v) a -> CollectionState k v a
CollectionState (State (Map k v) () -> CollectionState k v ())
-> State (Map k v) () -> CollectionState k v ()
forall a b. (a -> b) -> a -> b
$ (Map k v -> ((), Map k v)) -> State (Map k v) ()
forall a. (Map k v -> (a, Map k v)) -> StateT (Map k v) Identity a
forall s (m :: * -> *) a. MonadState s m => (s -> (a, s)) -> m a
state ((Map k v -> ((), Map k v)) -> State (Map k v) ())
-> (Map k v -> ((), Map k v)) -> State (Map k v) ()
forall a b. (a -> b) -> a -> b
$ \Map k v
s ->
    let (Map k v
s', Map k v
s'') = (Map k v, Map k v)
-> (Map k v -> (Map k v, Map k v))
-> Maybe (Map k v)
-> (Map k v, Map k v)
forall b a. b -> (a -> b) -> Maybe a -> b
maybe (Map k v
s, Map k v
s) (,Map k v
t'') (Maybe (Map k v) -> (Map k v, Map k v))
-> Maybe (Map k v) -> (Map k v, Map k v)
forall a b. (a -> b) -> a -> b
$ k -> v -> Map k v -> Maybe (Map k v)
forall k v. Ord k => k -> v -> Map k v -> Maybe (Map k v)
insert' k
k v
v Map k v
s
        (v
v, Map k v
t'') = State (Map k v) v -> Map k v -> (v, Map k v)
forall s a. State s a -> s -> (a, s)
runState State (Map k v) v
mv Map k v
s'
     in ((), Map k v
s'')