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

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

    VectorCache - mutable-vector-based (`IO` and `ST` hosted) `MonadCache`

    The fastest memoization cache, however it is even more limiting than "Control.Monad.Memo.Array" due to nature of "Data.Vector.Mutable". Still if you can use this cache please do since it will give you dramatic calculation speed up in comparison to pure `Data.Map.Map`-based cache, especially when unboxed `UVectorCache` is used.

    Limitations: Since `Data.Vector.Generic.Mutable.MVector` is used as `MonadCache` the key must be `Int` and the size of the cache's vector must be known beforehand with vector being allocated before the first call. In addition unboxed `UVectorCache` can only store `Data.Vector.Unboxed.Unbox` values (but it does it very efficiently).

    -}

    {-# LANGUAGE NoImplicitPrelude,
      MultiParamTypeClasses, FunctionalDependencies,
      FlexibleInstances, FlexibleContexts, TypeFamilies,
      UndecidableInstances, TypeSynonymInstances #-}

    module Control.Monad.Memo.Vector
     (

       -- * VectorCache for boxed types
       Vector,
       VectorCache,
       VectorMemo,
       evalVectorMemo,
       runVectorMemo,
       -- * UVectorCache for unboxed types
       UVector,
       UVectorCache,
       UVectorMemo,
       evalUVectorMemo,
       runUVectorMemo,
       -- * Generic functions for VectorCache
       Container(..),
       Cache,
       genericEvalVectorMemo,
       genericRunVectorMemo

    ) where

    import Data.Int
    import Data.Function
    import Data.Maybe (Maybe(..))
    import Data.Vector.Generic.Mutable
    import qualified Data.Vector.Mutable as M
    import qualified Data.Vector.Unboxed.Mutable as UM
    import Control.Applicative
    import Control.Monad
    import Control.Monad.Fix
    import Control.Monad.Trans.Class
    import Control.Monad.Primitive

    import Data.MaybeLike
    import Control.Monad.Memo.Class
    import Control.Monad.Trans.Memo.ReaderCache

    newtype Container vec = Container { forall vec. Container vec -> vec
    toVector :: vec }

    -- | Generic Vector-based memo cache
    type Cache vec s e = ReaderCache (Container (vec s e))

    instance (PrimMonad m, PrimState m ~ s, MaybeLike e v, MVector c e) =>
        MonadCache Int v (Cache c s e m) where
            {-# INLINE lookup #-}
            lookup :: Int -> Cache c s e m (Maybe v)
    lookup Int
    k = do
              c <- ReaderCache (Container (c s e)) m (Container (c s e))
    forall (m :: * -> *) c. Monad m => ReaderCache c m c
    container
              e <- lift $ read (toVector c) k
              return (if isNothing e then Nothing else Just (fromJust e))
            {-# INLINE add #-}
            add :: Int -> v -> Cache c s e m ()
    add Int
    k v
    v = do
              c <- ReaderCache (Container (c s e)) m (Container (c s e))
    forall (m :: * -> *) c. Monad m => ReaderCache c m c
    container
              lift $ write (toVector c) k (just v)

    instance (PrimMonad m, PrimState m ~ s, MaybeLike e v, MVector c e) =>
        MonadMemo Int v (Cache c s e m) where
            {-# INLINE memo #-}
            memo :: (Int -> Cache c s e m v) -> Int -> Cache c s e m v
    memo Int -> Cache c s e m v
    f Int
    k = do
              c <- ReaderCache (Container (c s e)) m (Container (c s e))
    forall (m :: * -> *) c. Monad m => ReaderCache c m c
    container
              e <- lift $ read (toVector c) k
              if isNothing e
                then do
                  v <- f k
                  lift $ write (toVector c) k (just v)
                  return v
                else return (fromJust e)

    -- VectorCache for boxed types
    -- --------------------------

    -- | Boxed vector
    type Vector = M.MVector

    -- | `MonadCache` based on boxed vector
    type VectorCache s e = Cache Vector s e

    -- | This is just to be able to infer the type of the `VectorCache` element.
    class MaybeLike e v => VectorMemo v e | v -> e

    -- | Evaluate computation using mutable boxed vector
    --
    -- Vector length must covers all possible keys used in computation
    -- otherwise /index out of bound/ error is generated by vector code
    evalVectorMemo :: (PrimMonad m, VectorMemo v e) =>
                      VectorCache (PrimState m) e m a  -- ^memoized computation
                   -> Int                              -- ^vector length
                   -> m a                              -- ^result
    {-# INLINE evalVectorMemo #-}
    evalVectorMemo :: forall (m :: * -> *) v e a.
    (PrimMonad m, VectorMemo v e) =>
    VectorCache (PrimState m) e m a -> Int -> m a
    evalVectorMemo = Cache MVector (PrimState m) e m a -> Int -> m a
    forall e v (m :: * -> *) (c :: * -> * -> *) a.
    (MaybeLike e v, PrimMonad m, MVector c e) =>
    Cache c (PrimState m) e m a -> Int -> m a
    genericEvalVectorMemo

    -- | Evaluate computation using mutable boxed vector.
    -- It also returns the final content of the vector cache
    --
    -- Vector length must covers all possible keys used in computation
    -- otherwise /index out of bound/ error is generated by vector code
    runVectorMemo :: (PrimMonad m, VectorMemo v e) =>
                     VectorCache (PrimState m) e m a  -- ^memoized computation
                  -> Int                              -- ^vector length
                  -> m (a, Vector (PrimState m) e)    -- ^result and final vector cache
    {-# INLINE runVectorMemo #-}
    runVectorMemo :: forall (m :: * -> *) v e a.
    (PrimMonad m, VectorMemo v e) =>
    VectorCache (PrimState m) e m a
    -> Int -> m (a, Vector (PrimState m) e)
    runVectorMemo = Cache MVector (PrimState m) e m a
    -> Int -> m (a, MVector (PrimState m) e)
    forall e v (m :: * -> *) (c :: * -> * -> *) a.
    (MaybeLike e v, PrimMonad m, MVector c e) =>
    Cache c (PrimState m) e m a -> Int -> m (a, c (PrimState m) e)
    genericRunVectorMemo

    -- VectorCache for unboxed types
    -- ----------------------------

    -- | Unboxed vector
    type UVector = UM.MVector

    -- | `MonadCache` based on unboxed vector
    type UVectorCache s e = Cache UVector s e

    -- | This is just to be able to infer the type of the `UVectorCache` element.
    class MaybeLike e v => UVectorMemo v e | v -> e

    -- | Evaluate computation using mutable unboxed vector
    --
    -- Vector length must covers all possible keys used in computation
    -- otherwise /index out of bound/ error is generated by vector code
    evalUVectorMemo :: (PrimMonad m, MVector UVector e, UVectorMemo v e) =>
                       UVectorCache (PrimState m) e m a  -- ^memoized computation
                    -> Int                               -- ^vector length
                    -> m a                               -- ^result
    {-# INLINE evalUVectorMemo #-}
    evalUVectorMemo :: forall (m :: * -> *) e v a.
    (PrimMonad m, MVector UVector e, UVectorMemo v e) =>
    UVectorCache (PrimState m) e m a -> Int -> m a
    evalUVectorMemo = Cache UVector (PrimState m) e m a -> Int -> m a
    forall e v (m :: * -> *) (c :: * -> * -> *) a.
    (MaybeLike e v, PrimMonad m, MVector c e) =>
    Cache c (PrimState m) e m a -> Int -> m a
    genericEvalVectorMemo

    -- | Evaluate computation using mutable unboxed vector.
    -- It also returns the final content of the vector cache
    --
    -- Vector length must covers all possible keys used in computation
    -- otherwise /index out of bound/ error is generated by vector code
    runUVectorMemo :: (PrimMonad m, MVector UVector e, UVectorMemo v e) =>
                      UVectorCache (PrimState m) e m a  -- ^memoized computation
                   -> Int                               -- ^vector length
                   -> m (a, UVector (PrimState m) e)    -- ^result and final vector cache
    {-# INLINE runUVectorMemo #-}
    runUVectorMemo :: forall (m :: * -> *) e v a.
    (PrimMonad m, MVector UVector e, UVectorMemo v e) =>
    UVectorCache (PrimState m) e m a
    -> Int -> m (a, UVector (PrimState m) e)
    runUVectorMemo = Cache UVector (PrimState m) e m a
    -> Int -> m (a, MVector (PrimState m) e)
    forall e v (m :: * -> *) (c :: * -> * -> *) a.
    (MaybeLike e v, PrimMonad m, MVector c e) =>
    Cache c (PrimState m) e m a -> Int -> m (a, c (PrimState m) e)
    genericRunVectorMemo

    genericEvalVectorMemo :: (MaybeLike e v, PrimMonad m, MVector c e) =>
                             Cache c (PrimState m) e m a -> Int -> m a
    {-# INLINE genericEvalVectorMemo #-}
    genericEvalVectorMemo :: forall e v (m :: * -> *) (c :: * -> * -> *) a.
    (MaybeLike e v, PrimMonad m, MVector c e) =>
    Cache c (PrimState m) e m a -> Int -> m a
    genericEvalVectorMemo Cache c (PrimState m) e m a
    m Int
    n = do
      c <- Int -> e -> m (c (PrimState m) e)
    forall (m :: * -> *) (v :: * -> * -> *) a.
    (PrimMonad m, MVector v a) =>
    Int -> a -> m (v (PrimState m) a)
    replicate Int
    n e
    forall a v. MaybeLike a v => a
    nothing
      evalReaderCache m (Container c)

    genericRunVectorMemo :: (MaybeLike e v, PrimMonad m, MVector c e) =>
                            Cache c (PrimState m) e m a -> Int -> m (a, c (PrimState m) e)
    {-# INLINE genericRunVectorMemo #-}
    genericRunVectorMemo :: forall e v (m :: * -> *) (c :: * -> * -> *) a.
    (MaybeLike e v, PrimMonad m, MVector c e) =>
    Cache c (PrimState m) e m a -> Int -> m (a, c (PrimState m) e)
    genericRunVectorMemo Cache c (PrimState m) e m a
    m Int
    n = do
      c <- Int -> e -> m (c (PrimState m) e)
    forall (m :: * -> *) (v :: * -> * -> *) a.
    (PrimMonad m, MVector v a) =>
    Int -> a -> m (v (PrimState m) a)
    replicate Int
    n e
    forall a v. MaybeLike a v => a
    nothing
      a <- evalReaderCache m (Container c)
      return (a, c)