    {- |
    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 `MonadCache` with unsafe operations.

    This is a version of "Control.Monad.Memo.Mutable.Vector" but implemented using /unsafe*/ vector operations.
    Faster than default implementation but you must be sure that your code doesn't try to read/write outside vector boundaries.

    -}

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

    module Control.Monad.Memo.Vector.Unsafe
     (

       -- * VectorCache for boxed types
       VectorCache,
       VectorMemo,
       unsafeEvalVectorMemo,
       unsafeRunVectorMemo,
       -- * UVectorCache for unboxed types
       UVectorCache,
       UVectorMemo,
       unsafeEvalUVectorMemo,
       unsafeRunUVectorMemo,
       -- * Generic functions for VectorCache
       Container(..),
       Cache,
       genericUnsafeEvalVectorMemo,
       genericUnsafeRunVectorMemo

    ) where

    import Data.Function
    import Data.Int
    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 k e = ReaderCache (Container (vec k 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 $ unsafeRead (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 $ unsafeWrite (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 $ unsafeRead (toVector c) k
              if isNothing e
                 then do
                   v <- f k
                   lift $ unsafeWrite (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 and unsafe operations
    --
    -- Vector length must covers all possible keys used in computation
    -- otherwise the behaviour is undefined (i.e. segfault)
    unsafeEvalVectorMemo :: (PrimMonad m, VectorMemo v e) =>
                            VectorCache (PrimState m) e m a -- ^memoized computation
                         -> Int                             -- ^vector length
                         -> m a                             -- ^result
    {-# INLINE unsafeEvalVectorMemo #-}
    unsafeEvalVectorMemo :: forall (m :: * -> *) v e a.
    (PrimMonad m, VectorMemo v e) =>
    VectorCache (PrimState m) e m a -> Int -> m a
    unsafeEvalVectorMemo = 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
    genericUnsafeEvalVectorMemo

    -- | Evaluate computation using mutable boxed vector and unsafe operations.
    -- It also returns the final content of the vector cache
    --
    -- Vector length must covers all possible keys used in computation
    -- otherwise the behaviour is undefined (i.e. segfault)
    unsafeRunVectorMemo :: (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 unsafeRunVectorMemo #-}
    unsafeRunVectorMemo :: forall (m :: * -> *) v e a.
    (PrimMonad m, VectorMemo v e) =>
    VectorCache (PrimState m) e m a
    -> Int -> m (a, Vector (PrimState m) e)
    unsafeRunVectorMemo = 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)
    genericUnsafeRunVectorMemo

    -- 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 and unsafe operations
    --
    -- Vector length must covers all possible keys used in computation
    -- otherwise the behaviour is undefined (i.e. segfault)
    unsafeEvalUVectorMemo :: (PrimMonad m, UVectorMemo v e, MVector UVector e) =>
                             UVectorCache (PrimState m) e m a -- ^memoized computation
                          -> Int                              -- ^vector length
                          -> m a                              -- ^result
    {-# INLINE unsafeEvalUVectorMemo #-}
    unsafeEvalUVectorMemo :: forall (m :: * -> *) v e a.
    (PrimMonad m, UVectorMemo v e, MVector UVector e) =>
    UVectorCache (PrimState m) e m a -> Int -> m a
    unsafeEvalUVectorMemo = 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
    genericUnsafeEvalVectorMemo

    -- | Evaluate computation using mutable boxed vector and unsafe operations.
    -- It also returns the final content of the vector cache
    --
    -- Vector length must covers all possible keys used in computation
    -- otherwise the behaviour is undefined (i.e. segfault)
    unsafeRunUVectorMemo :: (PrimMonad m, UVectorMemo v e, MVector UVector e) =>
                            UVectorCache (PrimState m) e m a -- ^memoized computation
                         -> Int                              -- ^vector length
                         -> m (a, UVector (PrimState m) e)   -- ^result and final vector cache
    {-# INLINE unsafeRunUVectorMemo #-}
    unsafeRunUVectorMemo :: forall (m :: * -> *) v e a.
    (PrimMonad m, UVectorMemo v e, MVector UVector e) =>
    UVectorCache (PrimState m) e m a
    -> Int -> m (a, UVector (PrimState m) e)
    unsafeRunUVectorMemo = 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)
    genericUnsafeRunVectorMemo

    genericUnsafeEvalVectorMemo :: (MaybeLike e v, PrimMonad m, MVector c e) =>
                                   Cache c (PrimState m) e m a -> Int -> m a
    {-# INLINE genericUnsafeEvalVectorMemo #-}
    genericUnsafeEvalVectorMemo :: forall e v (m :: * -> *) (c :: * -> * -> *) a.
    (MaybeLike e v, PrimMonad m, MVector c e) =>
    Cache c (PrimState m) e m a -> Int -> m a
    genericUnsafeEvalVectorMemo Cache c (PrimState m) e m a
    m Int
    n = do
      vec <- 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 vec)

    genericUnsafeRunVectorMemo :: (MaybeLike e v, PrimMonad m, MVector c e) =>
                                  Cache c (PrimState m) e m a -> Int -> m (a, c (PrimState m) e)
    {-# INLINE genericUnsafeRunVectorMemo #-}
    genericUnsafeRunVectorMemo :: 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)
    genericUnsafeRunVectorMemo Cache c (PrimState m) e m a
    m Int
    n = do
      vec <- 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 vec)
      return (a, vec)