{-# LANGUAGE NoImplicitPrelude,
MultiParamTypeClasses, FunctionalDependencies,
FlexibleInstances, FlexibleContexts, TypeSynonymInstances,
UndecidableInstances, TypeFamilies #-}
module Control.Monad.Memo.Vector.Expandable
(
VectorCache,
VectorMemo,
startEvalVectorMemo,
startRunVectorMemo,
UVectorCache,
UVectorMemo,
startEvalUVectorMemo,
startRunUVectorMemo,
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 }
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)
type Vector = M.MVector
type VectorCache s e = Cache Vector s e
class MaybeLike e v => VectorMemo v e | v -> e
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
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
type UVector = UM.MVector
type UVectorCache s e = Cache UVector s e
class MaybeLike e v => UVectorMemo v e | v -> e
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
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)