{- |
Module      :  Control.Monad.Memo
Copyright   :  (c) Eduard Sergeev 2011
License     :  BSD-style (see the file LICENSE)

Maintainer  :  eduard.sergeev@gmail.com
Stability   :  experimental
Portability :  non-portable (multi-param classes, flexible instances)

Defines "MemoStateT" - generalized (to any "Data.MapLike" content) memoization monad transformer

-}

{-# LANGUAGE NoImplicitPrelude, MultiParamTypeClasses,
  FlexibleInstances, TypeSynonymInstances #-}

module Control.Monad.Trans.Memo.State
(
 
    -- * MemoStateT monad transformer
    MemoStateT(..),
    runMemoStateT,
    evalMemoStateT,
    -- * MemoState monad
    MemoState,
    runMemoState,
    evalMemoState,
    -- * Internal
    Container(..)

) where


import Data.Tuple
import Data.Function
import Data.Functor.Identity
import Control.Applicative
import Control.Monad
import Control.Monad.Trans.Class

import qualified Data.MapLike as M
import Control.Monad.Memo.Class
import Control.Monad.Trans.Memo.StateCache


newtype Container s = Container { Container s -> s
toState :: s }

-- | Memoization monad transformer based on `StateCache`
-- to be used with pure cache containers which support `M.MapLike` interface
type MemoStateT s k v = StateCache (Container s)


-- | Returns the pair of the result of `MonadMemo` computation
-- along with the final state of the internal pure container wrapped in monad
runMemoStateT :: Monad m => MemoStateT s k v m a -> s -> m (a, s)
runMemoStateT :: MemoStateT s k v m a -> s -> m (a, s)
runMemoStateT MemoStateT s k v m a
m s
s = do
  (a
a, Container s
c) <- MemoStateT s k v m a -> Container s -> m (a, Container s)
forall s (m :: * -> *) a. StateCache s m a -> s -> m (a, s)
runStateCache MemoStateT s k v m a
m (s -> Container s
forall s. s -> Container s
Container s
s)
  (a, s) -> m (a, s)
forall (m :: * -> *) a. Monad m => a -> m a
return (a
a, Container s -> s
forall s. Container s -> s
toState Container s
c)

-- | Returns the result of `MonadMemo` computation wrapped in monad.
-- This function discards the cache
evalMemoStateT :: Monad m => MemoStateT c k v m a -> c -> m a
evalMemoStateT :: MemoStateT c k v m a -> c -> m a
evalMemoStateT MemoStateT c k v m a
m c
s = MemoStateT c k v m a -> c -> m (a, c)
forall (m :: * -> *) s k v a.
Monad m =>
MemoStateT s k v m a -> s -> m (a, s)
runMemoStateT MemoStateT c k v m a
m c
s m (a, c) -> ((a, c) -> m a) -> m a
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= a -> m a
forall (m :: * -> *) a. Monad m => a -> m a
return (a -> m a) -> ((a, c) -> a) -> (a, c) -> m a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (a, c) -> a
forall a b. (a, b) -> a
fst


-- | Memoization monad based on `StateCache`
-- to be used with pure cache containers which support `M.MapLike` interface
type MemoState c k v = MemoStateT c k v Identity

-- | Returns the pair of the result of `MonadMemo` computation
-- along with the final state of the internal pure container
runMemoState :: MemoState c k v a -> c -> (a, c)
runMemoState :: MemoState c k v a -> c -> (a, c)
runMemoState MemoState c k v a
m = Identity (a, c) -> (a, c)
forall a. Identity a -> a
runIdentity (Identity (a, c) -> (a, c))
-> (c -> Identity (a, c)) -> c -> (a, c)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. MemoState c k v a -> c -> Identity (a, c)
forall (m :: * -> *) s k v a.
Monad m =>
MemoStateT s k v m a -> s -> m (a, s)
runMemoStateT MemoState c k v a
m

-- | Returns the result of `MonadMemo` computation discarding the cache
evalMemoState :: MemoState c k v a -> c -> a
evalMemoState :: MemoState c k v a -> c -> a
evalMemoState MemoState c k v a
m = Identity a -> a
forall a. Identity a -> a
runIdentity (Identity a -> a) -> (c -> Identity a) -> c -> a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. MemoState c k v a -> c -> Identity a
forall (m :: * -> *) c k v a.
Monad m =>
MemoStateT c k v m a -> c -> m a
evalMemoStateT MemoState c k v a
m


instance (Monad m, M.MapLike c k v) => MonadCache k v (MemoStateT c k v m) where
    {-# INLINE lookup #-}
    lookup :: k -> MemoStateT c k v m (Maybe v)
lookup k
k = StateCache (Container c) m (Container c)
forall (m :: * -> *) c. Monad m => StateCache c m c
container StateCache (Container c) m (Container c)
-> (Container c -> MemoStateT c k v m (Maybe v))
-> MemoStateT c k v m (Maybe v)
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= Maybe v -> MemoStateT c k v m (Maybe v)
forall (m :: * -> *) a. Monad m => a -> m a
return (Maybe v -> MemoStateT c k v m (Maybe v))
-> (Container c -> Maybe v)
-> Container c
-> MemoStateT c k v m (Maybe v)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. k -> c -> Maybe v
forall c k v. MapLike c k v => k -> c -> Maybe v
M.lookup k
k (c -> Maybe v) -> (Container c -> c) -> Container c -> Maybe v
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Container c -> c
forall s. Container s -> s
toState
    {-# INLINE add #-}
    add :: k -> v -> MemoStateT c k v m ()
add k
k v
v  = StateCache (Container c) m (Container c)
forall (m :: * -> *) c. Monad m => StateCache c m c
container StateCache (Container c) m (Container c)
-> (Container c -> MemoStateT c k v m ()) -> MemoStateT c k v m ()
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= Container c -> MemoStateT c k v m ()
forall (m :: * -> *) c. Monad m => c -> StateCache c m ()
setContainer (Container c -> MemoStateT c k v m ())
-> (Container c -> Container c)
-> Container c
-> MemoStateT c k v m ()
forall b c a. (b -> c) -> (a -> b) -> a -> c
. c -> Container c
forall s. s -> Container s
Container (c -> Container c)
-> (Container c -> c) -> Container c -> Container c
forall b c a. (b -> c) -> (a -> b) -> a -> c
. k -> v -> c -> c
forall c k v. MapLike c k v => k -> v -> c -> c
M.add k
k v
v (c -> c) -> (Container c -> c) -> Container c -> c
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Container c -> c
forall s. Container s -> s
toState

instance (Monad m, M.MapLike c k v) => MonadMemo k v (MemoStateT c k v m) where
    {-# INLINE memo #-}
    memo :: (k -> MemoStateT c k v m v) -> k -> MemoStateT c k v m v
memo = (k -> MemoStateT c k v m v) -> k -> MemoStateT c k v m v
forall k v (m :: * -> *).
(MonadCache k v m, Monad m) =>
(k -> m v) -> k -> m v
memol0