    {- |
    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)

    Vector-based `MonadCache` implementation which dynamically expands the vector during the computation to accomodate all requested keys.
    This implementation does not require to specify the length of the vector up front, but may be slower than "Control.Monad.Memo.Vector".

    -}

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

    module Control.Monad.Memo.Vector.Expandable
     (

       -- * VectorCache for boxed types
       VectorCache,
       VectorMemo,
       startEvalVectorMemo,
       startRunVectorMemo,
       -- * UVectorCache for unboxed types
       UVectorCache,
       UVectorMemo,
       startEvalUVectorMemo,
       startRunUVectorMemo,
       -- * Generic functions for VectorCache
       Container(..),
       Cache,
       genericStartEvalVectorMemo,
       genericStartRunVectorMemo

    ) where

    import Data.Int
    import Data.Eq
    import Data.Ord
    import Data.Function
    import Prelude (Num(..))
    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.Monad
    import Control.Monad.Trans.Class
    import Control.Monad.Primitive

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

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

    -- | Generic Vector-based memo cache
    type Cache vec k e = StateCache (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 <- StateCache (Container (c s e)) m (Container (c s e))
    forall (m :: * -> *) c. Monad m => StateCache c m c
    container
              e <- lift $ cacheRead (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 <- StateCache (Container (c s e)) m (Container (c s e))
    forall (m :: * -> *) c. Monad m => StateCache c m c
    container
              v' <- lift $ cacheWrite (toVector c) k (just v)
              setContainer (Container 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
              Container vec <- StateCache (Container (c s e)) m (Container (c s e))
    forall (m :: * -> *) c. Monad m => StateCache c m c
    container
              let l = c s e -> Int
    forall (v :: * -> * -> *) a s. MVector v a => v s a -> Int
    length c s e
    vec
                  d = Int
    k Int -> Int -> Int
    forall a. Num a => a -> a -> a
    + Int
    1 Int -> Int -> Int
    forall a. Num a => a -> a -> a
    - Int
    l
              if d > 0
                then do
                  vec' <- lift $ expand vec l d
                  setContainer (Container vec')
                  v <- f k
                  Container vec'' <- container
                  lift $ unsafeWrite vec'' k (just v)
                  return v
                else do
                  e <- lift $ cacheRead vec k
                  if isNothing e
                     then do
                       v <- f k
                       Container vec' <- container
                       lift $ unsafeWrite vec' 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 which dynamically grows to accomodate all requested keys
    startEvalVectorMemo :: (PrimMonad m, VectorMemo v e) =>
                           VectorCache (PrimState m) e m a -> m a
    {-# INLINE startEvalVectorMemo #-}
    startEvalVectorMemo :: forall (m :: * -> *) v e a.
    (PrimMonad m, VectorMemo v e) =>
    VectorCache (PrimState m) e m a -> m a
    startEvalVectorMemo = Cache MVector (PrimState m) e m a -> m a
    forall e v (m :: * -> *) (vec :: * -> * -> *) a.
    (MaybeLike e v, PrimMonad m, MVector vec e) =>
    Cache vec (PrimState m) e m a -> m a
    genericStartEvalVectorMemo

    -- | Evaluate computation using mutable boxed vector
    -- which dynamically grows to accomodate all requested keys.
    -- This function also returns the final content of the vector cache
    startRunVectorMemo :: (PrimMonad m, VectorMemo v e) =>
                          VectorCache (PrimState m) e m a -> m (a, Vector (PrimState m) e)
    {-# INLINE startRunVectorMemo #-}
    startRunVectorMemo :: forall (m :: * -> *) v e a.
    (PrimMonad m, VectorMemo v e) =>
    VectorCache (PrimState m) e m a -> m (a, Vector (PrimState m) e)
    startRunVectorMemo = Cache MVector (PrimState m) e m a -> m (a, MVector (PrimState m) e)
    forall e v (m :: * -> *) (vec :: * -> * -> *) a.
    (MaybeLike e v, PrimMonad m, MVector vec e) =>
    Cache vec (PrimState m) e m a -> m (a, vec (PrimState m) e)
    genericStartRunVectorMemo

    -- 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
    -- which dynamically grows to accomodate all requested keys
    startEvalUVectorMemo :: (PrimMonad m, UVectorMemo v e, MVector UVector e) =>
                            UVectorCache (PrimState m) e m a -> m a
    {-# INLINE startEvalUVectorMemo #-}
    startEvalUVectorMemo :: forall (m :: * -> *) v e a.
    (PrimMonad m, UVectorMemo v e, MVector UVector e) =>
    UVectorCache (PrimState m) e m a -> m a
    startEvalUVectorMemo = Cache UVector (PrimState m) e m a -> m a
    forall e v (m :: * -> *) (vec :: * -> * -> *) a.
    (MaybeLike e v, PrimMonad m, MVector vec e) =>
    Cache vec (PrimState m) e m a -> m a
    genericStartEvalVectorMemo

    -- | Evaluate computation using mutable unboxed vector
    -- which dynamically grows to accomodate all requested keys.
    -- This function also returns the final content of the vector cache
    startRunUVectorMemo :: (PrimMonad m, UVectorMemo v e, MVector UVector e) =>
                           UVectorCache (PrimState m) e m a -> m (a, UVector (PrimState m) e)
    {-# INLINE startRunUVectorMemo #-}
    startRunUVectorMemo :: forall (m :: * -> *) v e a.
    (PrimMonad m, UVectorMemo v e, MVector UVector e) =>
    UVectorCache (PrimState m) e m a -> m (a, UVector (PrimState m) e)
    startRunUVectorMemo = Cache UVector (PrimState m) e m a -> m (a, MVector (PrimState m) e)
    forall e v (m :: * -> *) (vec :: * -> * -> *) a.
    (MaybeLike e v, PrimMonad m, MVector vec e) =>
    Cache vec (PrimState m) e m a -> m (a, vec (PrimState m) e)
    genericStartRunVectorMemo

    genericStartEvalVectorMemo :: (MaybeLike e v, PrimMonad m, MVector vec e) =>
                                  Cache vec (PrimState m) e m a -> m a
    {-# INLINE genericStartEvalVectorMemo #-}
    genericStartEvalVectorMemo :: forall e v (m :: * -> *) (vec :: * -> * -> *) a.
    (MaybeLike e v, PrimMonad m, MVector vec e) =>
    Cache vec (PrimState m) e m a -> m a
    genericStartEvalVectorMemo Cache vec (PrimState m) e m a
    m = do
      (a,_) <- Cache vec (PrimState m) e m a -> m (a, vec (PrimState m) e)
    forall e v (m :: * -> *) (vec :: * -> * -> *) a.
    (MaybeLike e v, PrimMonad m, MVector vec e) =>
    Cache vec (PrimState m) e m a -> m (a, vec (PrimState m) e)
    genericStartRunVectorMemo Cache vec (PrimState m) e m a
    m
      return a

    genericStartRunVectorMemo :: (MaybeLike e v, PrimMonad m, MVector vec e) =>
                                 Cache vec (PrimState m) e m a -> m (a, vec (PrimState m) e)
    {-# INLINE genericStartRunVectorMemo #-}
    genericStartRunVectorMemo :: forall e v (m :: * -> *) (vec :: * -> * -> *) a.
    (MaybeLike e v, PrimMonad m, MVector vec e) =>
    Cache vec (PrimState m) e m a -> m (a, vec (PrimState m) e)
    genericStartRunVectorMemo Cache vec (PrimState m) e m a
    m = do
      vec <- Int -> e -> m (vec (PrimState m) e)
    forall (m :: * -> *) (v :: * -> * -> *) a.
    (PrimMonad m, MVector v a) =>
    Int -> a -> m (v (PrimState m) a)
    replicate Int
    0 e
    forall a v. MaybeLike a v => a
    nothing
      (a, c) <- runStateCache m (Container vec)
      return (a, toVector c)

    {-# INLINE cacheRead #-}
    cacheRead :: v (PrimState m) a -> Int -> m a
    cacheRead v (PrimState m) a
    c Int
    k =
        if Int
    k Int -> Int -> Bool
    forall a. Ord a => a -> a -> Bool
    >= v (PrimState m) a -> Int
    forall (v :: * -> * -> *) a s. MVector v a => v s a -> Int
    length v (PrimState m) a
    c
            then a -> m a
    forall a. a -> m a
    forall (m :: * -> *) a. Monad m => a -> m a
    return a
    forall a v. MaybeLike a v => a
    nothing
            else v (PrimState m) a -> Int -> m a
    forall (m :: * -> *) (v :: * -> * -> *) a.
    (PrimMonad m, MVector v a) =>
    v (PrimState m) a -> Int -> m a
    unsafeRead v (PrimState m) a
    c Int
    k

    {-# INLINE cacheWrite #-}
    cacheWrite :: v (PrimState m) a -> Int -> a -> m (v (PrimState m) a)
    cacheWrite v (PrimState m) a
    c Int
    k a
    e = do
      c' <- if Int
    d Int -> Int -> Bool
    forall a. Ord a => a -> a -> Bool
    > Int
    0
               then v (PrimState m) a -> Int -> Int -> m (v (PrimState m) a)
    forall {m :: * -> *} {v :: * -> * -> *} {a} {v}.
    (PrimMonad m, MVector v a, MaybeLike a v) =>
    v (PrimState m) a -> Int -> Int -> m (v (PrimState m) a)
    expand v (PrimState m) a
    c Int
    l Int
    d
               else v (PrimState m) a -> m (v (PrimState m) a)
    forall a. a -> m a
    forall (m :: * -> *) a. Monad m => a -> m a
    return v (PrimState m) a
    c
      unsafeWrite c' k e
      return c'
        where
          l :: Int
    l = v (PrimState m) a -> Int
    forall (v :: * -> * -> *) a s. MVector v a => v s a -> Int
    length v (PrimState m) a
    c
          d :: Int
    d = Int
    k Int -> Int -> Int
    forall a. Num a => a -> a -> a
    + Int
    1 Int -> Int -> Int
    forall a. Num a => a -> a -> a
    - Int
    l

    {-# INLINE expand #-}
    expand :: v (PrimState m) a -> Int -> Int -> m (v (PrimState m) a)
    expand v (PrimState m) a
    c Int
    l Int
    d = do
      uc <- v (PrimState m) a -> Int -> m (v (PrimState m) a)
    forall (m :: * -> *) (v :: * -> * -> *) a.
    (PrimMonad m, MVector v a) =>
    v (PrimState m) a -> Int -> m (v (PrimState m) a)
    unsafeGrow v (PrimState m) a
    c Int
    toGrow
      unsafeWrite uc l nothing
      initialise uc 1
          where
            toGrow :: Int
    toGrow = Int
    d Int -> Int -> Int
    forall a. Ord a => a -> a -> a
    `max` (Int
    l Int -> Int -> Int
    forall a. Num a => a -> a -> a
    * Int
    2)
            {-# INLINE initialise #-}
            initialise :: v (PrimState m) a -> Int -> m (v (PrimState m) a)
    initialise v (PrimState m) a
    c Int
    i | Int
    i Int -> Int -> Bool
    forall a. Eq a => a -> a -> Bool
    == Int
    toGrow = v (PrimState m) a -> m (v (PrimState m) a)
    forall a. a -> m a
    forall (m :: * -> *) a. Monad m => a -> m a
    return v (PrimState m) a
    c
            initialise v (PrimState m) a
    c Int
    i = do
                     let n :: Int
    n = Int
    i Int -> Int -> Int
    forall a. Ord a => a -> a -> a
    `min` (Int
    toGrowInt -> Int -> Int
    forall a. Num a => a -> a -> a
    -Int
    i)
                         t :: v (PrimState m) a
    t = Int -> Int -> v (PrimState m) a -> v (PrimState m) a
    forall (v :: * -> * -> *) a s.
    MVector v a =>
    Int -> Int -> v s a -> v s a
    unsafeSlice (Int
    lInt -> Int -> Int
    forall a. Num a => a -> a -> a
    +Int
    i) Int
    n v (PrimState m) a
    c
                         s :: v (PrimState m) a
    s = Int -> Int -> v (PrimState m) a -> v (PrimState m) a
    forall (v :: * -> * -> *) a s.
    MVector v a =>
    Int -> Int -> v s a -> v s a
    unsafeSlice Int
    l Int
    n v (PrimState m) a
    c
                     v (PrimState m) a -> v (PrimState m) a -> m ()
    forall (m :: * -> *) (v :: * -> * -> *) a.
    (PrimMonad m, MVector v a) =>
    v (PrimState m) a -> v (PrimState m) a -> m ()
    unsafeCopy v (PrimState m) a
    t v (PrimState m) a
    s
                     v (PrimState m) a -> Int -> m (v (PrimState m) a)
    initialise v (PrimState m) a
    c (Int
    iInt -> Int -> Int
    forall a. Num a => a -> a -> a
    +Int
    n)