{- |
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 { 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
          Container (c s e)
c <- StateCache (Container (c s e)) m (Container (c s e))
forall (m :: * -> *) c. Monad m => StateCache c m c
container
          e
e <- m e -> StateCache (Container (c s e)) m e
forall (t :: (* -> *) -> * -> *) (m :: * -> *) a.
(MonadTrans t, Monad m) =>
m a -> t m a
lift (m e -> StateCache (Container (c s e)) m e)
-> m e -> StateCache (Container (c s e)) m e
forall a b. (a -> b) -> a -> b
$ c (PrimState m) e -> Int -> m e
forall (v :: * -> * -> *) a (m :: * -> *) v.
(MVector v a, MaybeLike a v, PrimMonad m) =>
v (PrimState m) a -> Int -> m a
cacheRead (Container (c s e) -> c s e
forall vec. Container vec -> vec
toVector Container (c s e)
c) Int
k
          Maybe v -> Cache c s e m (Maybe v)
forall (m :: * -> *) a. Monad m => a -> m a
return (if e -> Bool
forall a v. MaybeLike a v => a -> Bool
isNothing e
e then Maybe v
forall a. Maybe a
Nothing else v -> Maybe v
forall a. a -> Maybe a
Just (e -> v
forall a v. MaybeLike a v => a -> v
fromJust e
e))
        {-# INLINE add #-}
        add :: Int -> v -> Cache c s e m ()
add Int
k v
v = do 
          Container (c s e)
c <- StateCache (Container (c s e)) m (Container (c s e))
forall (m :: * -> *) c. Monad m => StateCache c m c
container
          c s e
v' <- m (c s e) -> StateCache (Container (c s e)) m (c s e)
forall (t :: (* -> *) -> * -> *) (m :: * -> *) a.
(MonadTrans t, Monad m) =>
m a -> t m a
lift (m (c s e) -> StateCache (Container (c s e)) m (c s e))
-> m (c s e) -> StateCache (Container (c s e)) m (c s e)
forall a b. (a -> b) -> a -> b
$ c (PrimState m) e -> Int -> e -> m (c (PrimState m) e)
forall (v :: * -> * -> *) a (m :: * -> *) v.
(MVector v a, PrimMonad m, MaybeLike a v) =>
v (PrimState m) a -> Int -> a -> m (v (PrimState m) a)
cacheWrite (Container (c s e) -> c s e
forall vec. Container vec -> vec
toVector Container (c s e)
c) Int
k (v -> e
forall a v. MaybeLike a v => v -> a
just v
v)
          Container (c s e) -> Cache c s e m ()
forall (m :: * -> *) c. Monad m => c -> StateCache c m ()
setContainer (c s e -> Container (c s e)
forall vec. vec -> Container vec
Container c s e
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 c s e
vec <- StateCache (Container (c s e)) m (Container (c s e))
forall (m :: * -> *) c. Monad m => StateCache c m c
container
          let l :: Int
l = c s e -> Int
forall (v :: * -> * -> *) a s. MVector v a => v s a -> Int
length c s e
vec
              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
          if Int
d Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
> Int
0
            then do
              c s e
vec' <- m (c s e) -> StateCache (Container (c s e)) m (c s e)
forall (t :: (* -> *) -> * -> *) (m :: * -> *) a.
(MonadTrans t, Monad m) =>
m a -> t m a
lift (m (c s e) -> StateCache (Container (c s e)) m (c s e))
-> m (c s e) -> StateCache (Container (c s e)) m (c s e)
forall a b. (a -> b) -> a -> b
$ c (PrimState m) e -> Int -> Int -> m (c (PrimState m) e)
forall (m :: * -> *) (v :: * -> * -> *) a v.
(MVector v a, PrimMonad m, MaybeLike a v) =>
v (PrimState m) a -> Int -> Int -> m (v (PrimState m) a)
expand c s e
c (PrimState m) e
vec Int
l Int
d
              Container (c s e) -> StateCache (Container (c s e)) m ()
forall (m :: * -> *) c. Monad m => c -> StateCache c m ()
setContainer (c s e -> Container (c s e)
forall vec. vec -> Container vec
Container c s e
vec')
              v
v <- Int -> Cache c s e m v
f Int
k
              Container c s e
vec'' <- StateCache (Container (c s e)) m (Container (c s e))
forall (m :: * -> *) c. Monad m => StateCache c m c
container
              m () -> StateCache (Container (c s e)) m ()
forall (t :: (* -> *) -> * -> *) (m :: * -> *) a.
(MonadTrans t, Monad m) =>
m a -> t m a
lift (m () -> StateCache (Container (c s e)) m ())
-> m () -> StateCache (Container (c s e)) m ()
forall a b. (a -> b) -> a -> b
$ c (PrimState m) e -> Int -> e -> m ()
forall (m :: * -> *) (v :: * -> * -> *) a.
(PrimMonad m, MVector v a) =>
v (PrimState m) a -> Int -> a -> m ()
unsafeWrite c s e
c (PrimState m) e
vec'' Int
k (v -> e
forall a v. MaybeLike a v => v -> a
just v
v)
              v -> Cache c s e m v
forall (m :: * -> *) a. Monad m => a -> m a
return v
v
            else do
              e
e <- m e -> StateCache (Container (c s e)) m e
forall (t :: (* -> *) -> * -> *) (m :: * -> *) a.
(MonadTrans t, Monad m) =>
m a -> t m a
lift (m e -> StateCache (Container (c s e)) m e)
-> m e -> StateCache (Container (c s e)) m e
forall a b. (a -> b) -> a -> b
$ c (PrimState m) e -> Int -> m e
forall (v :: * -> * -> *) a (m :: * -> *) v.
(MVector v a, MaybeLike a v, PrimMonad m) =>
v (PrimState m) a -> Int -> m a
cacheRead c s e
c (PrimState m) e
vec Int
k
              if e -> Bool
forall a v. MaybeLike a v => a -> Bool
isNothing e
e
                 then do
                   v
v <- Int -> Cache c s e m v
f Int
k
                   Container c s e
vec' <- StateCache (Container (c s e)) m (Container (c s e))
forall (m :: * -> *) c. Monad m => StateCache c m c
container
                   m () -> StateCache (Container (c s e)) m ()
forall (t :: (* -> *) -> * -> *) (m :: * -> *) a.
(MonadTrans t, Monad m) =>
m a -> t m a
lift (m () -> StateCache (Container (c s e)) m ())
-> m () -> StateCache (Container (c s e)) m ()
forall a b. (a -> b) -> a -> b
$ c (PrimState m) e -> Int -> e -> m ()
forall (m :: * -> *) (v :: * -> * -> *) a.
(PrimMonad m, MVector v a) =>
v (PrimState m) a -> Int -> a -> m ()
unsafeWrite c s e
c (PrimState m) e
vec' Int
k (v -> e
forall a v. MaybeLike a v => v -> a
just v
v)
                   v -> Cache c s e m v
forall (m :: * -> *) a. Monad m => a -> m a
return v
v
                 else v -> Cache c s e m v
forall (m :: * -> *) a. Monad m => a -> m a
return (e -> v
forall a v. MaybeLike a v => a -> v
fromJust e
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 :: VectorCache (PrimState m) e m a -> m a
startEvalVectorMemo = VectorCache (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 :: VectorCache (PrimState m) e m a -> m (a, Vector (PrimState m) e)
startRunVectorMemo = VectorCache (PrimState m) e m a -> m (a, Vector (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 :: UVectorCache (PrimState m) e m a -> m a
startEvalUVectorMemo = UVectorCache (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 :: UVectorCache (PrimState m) e m a -> m (a, UVector (PrimState m) e)
startRunUVectorMemo = UVectorCache (PrimState m) e m a -> m (a, UVector (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 :: Cache vec (PrimState m) e m a -> m a
genericStartEvalVectorMemo Cache vec (PrimState m) e m a
m = do
  (a
a,vec (PrimState m) e
_) <- 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
  a -> m a
forall (m :: * -> *) a. Monad m => a -> m a
return a
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 :: Cache vec (PrimState m) e m a -> m (a, vec (PrimState m) e)
genericStartRunVectorMemo Cache vec (PrimState m) e m a
m = do
  vec (PrimState m) e
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
a, Container (vec (PrimState m) e)
c) <- Cache vec (PrimState m) e m a
-> Container (vec (PrimState m) e)
-> m (a, Container (vec (PrimState m) e))
forall s (m :: * -> *) a. StateCache s m a -> s -> m (a, s)
runStateCache Cache vec (PrimState m) e m a
m (vec (PrimState m) e -> Container (vec (PrimState m) e)
forall vec. vec -> Container vec
Container vec (PrimState m) e
vec)
  (a, vec (PrimState m) e) -> m (a, vec (PrimState m) e)
forall (m :: * -> *) a. Monad m => a -> m a
return (a
a, Container (vec (PrimState m) e) -> vec (PrimState m) e
forall vec. Container vec -> vec
toVector Container (vec (PrimState m) e)
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 (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
  v (PrimState m) a
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.
(MVector v a, PrimMonad m, 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 (m :: * -> *) a. Monad m => a -> m a
return v (PrimState m) a
c
  v (PrimState m) a -> Int -> a -> m ()
forall (m :: * -> *) (v :: * -> * -> *) a.
(PrimMonad m, MVector v a) =>
v (PrimState m) a -> Int -> a -> m ()
unsafeWrite v (PrimState m) a
c' Int
k a
e
  v (PrimState m) a -> m (v (PrimState m) a)
forall (m :: * -> *) a. Monad m => a -> m a
return v (PrimState m) a
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
  v (PrimState m) a
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
  v (PrimState m) a -> Int -> a -> m ()
forall (m :: * -> *) (v :: * -> * -> *) a.
(PrimMonad m, MVector v a) =>
v (PrimState m) a -> Int -> a -> m ()
unsafeWrite v (PrimState m) a
uc Int
l a
forall a v. MaybeLike a v => a
nothing
  v (PrimState m) a -> Int -> m (v (PrimState m) a)
initialise v (PrimState m) a
uc Int
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 (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)