    {- |
    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 { forall s. 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 :: forall (m :: * -> *) s k v a.
    Monad m =>
    MemoStateT s k v m a -> s -> m (a, s)
    runMemoStateT MemoStateT s k v m a
    m s
    s = do
      (a, 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)
      return (a, toState 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 :: forall (m :: * -> *) c k v a.
    Monad m =>
    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 a b. m a -> (a -> m b) -> m b
    forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
    >>= a -> m a
    forall a. 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 :: forall c k v a. 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 :: forall c k v a. 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 a b.
    StateCache (Container c) m a
    -> (a -> StateCache (Container c) m b)
    -> StateCache (Container c) m b
    forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
    >>= Maybe v -> MemoStateT c k v m (Maybe v)
    forall a. a -> StateCache (Container c) m a
    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 a b.
    StateCache (Container c) m a
    -> (a -> StateCache (Container c) m b)
    -> StateCache (Container c) m b
    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