{- |
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 { 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
          Container (c s e)
c <- ReaderCache (Container (c s e)) m (Container (c s e))
forall (m :: * -> *) c. Monad m => ReaderCache c m c
container
          e
e <- m e -> ReaderCache (Container (c s e)) m e
forall (t :: (* -> *) -> * -> *) (m :: * -> *) a.
(MonadTrans t, Monad m) =>
m a -> t m a
lift (m e -> ReaderCache (Container (c s e)) m e)
-> m e -> ReaderCache (Container (c s e)) m e
forall a b. (a -> b) -> a -> b
$ c (PrimState m) e -> Int -> m e
forall (m :: * -> *) (v :: * -> * -> *) a.
(PrimMonad m, MVector v a) =>
v (PrimState m) a -> Int -> m a
unsafeRead (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 <- ReaderCache (Container (c s e)) m (Container (c s e))
forall (m :: * -> *) c. Monad m => ReaderCache c m c
container
          m () -> Cache c s e m ()
forall (t :: (* -> *) -> * -> *) (m :: * -> *) a.
(MonadTrans t, Monad m) =>
m a -> t m a
lift (m () -> Cache c s e m ()) -> m () -> Cache 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 (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)

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)
c <- ReaderCache (Container (c s e)) m (Container (c s e))
forall (m :: * -> *) c. Monad m => ReaderCache c m c
container
          e
e <- m e -> ReaderCache (Container (c s e)) m e
forall (t :: (* -> *) -> * -> *) (m :: * -> *) a.
(MonadTrans t, Monad m) =>
m a -> t m a
lift (m e -> ReaderCache (Container (c s e)) m e)
-> m e -> ReaderCache (Container (c s e)) m e
forall a b. (a -> b) -> a -> b
$ c (PrimState m) e -> Int -> m e
forall (m :: * -> *) (v :: * -> * -> *) a.
(PrimMonad m, MVector v a) =>
v (PrimState m) a -> Int -> m a
unsafeRead (Container (c s e) -> c s e
forall vec. Container vec -> vec
toVector Container (c s e)
c) 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
               m () -> ReaderCache (Container (c s e)) m ()
forall (t :: (* -> *) -> * -> *) (m :: * -> *) a.
(MonadTrans t, Monad m) =>
m a -> t m a
lift (m () -> ReaderCache (Container (c s e)) m ())
-> m () -> ReaderCache (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 (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)
               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 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 :: VectorCache (PrimState m) e m a -> Int -> m a
unsafeEvalVectorMemo = VectorCache (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 :: VectorCache (PrimState m) e m a
-> Int -> m (a, Vector (PrimState m) e)
unsafeRunVectorMemo = VectorCache (PrimState m) e m a
-> Int -> m (a, Vector (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 :: UVectorCache (PrimState m) e m a -> Int -> m a
unsafeEvalUVectorMemo = UVectorCache (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 :: UVectorCache (PrimState m) e m a
-> Int -> m (a, UVector (PrimState m) e)
unsafeRunUVectorMemo = UVectorCache (PrimState m) e m a
-> Int -> m (a, UVector (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 :: Cache c (PrimState m) e m a -> Int -> m a
genericUnsafeEvalVectorMemo Cache c (PrimState m) e m a
m Int
n = do
  c (PrimState m) e
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
  Cache c (PrimState m) e m a -> Container (c (PrimState m) e) -> m a
forall r (m :: * -> *) a. ReaderCache r m a -> r -> m a
evalReaderCache Cache c (PrimState m) e m a
m (c (PrimState m) e -> Container (c (PrimState m) e)
forall vec. vec -> Container vec
Container c (PrimState m) e
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 :: 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
  c (PrimState m) e
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
a <- Cache c (PrimState m) e m a -> Container (c (PrimState m) e) -> m a
forall r (m :: * -> *) a. ReaderCache r m a -> r -> m a
evalReaderCache Cache c (PrimState m) e m a
m (c (PrimState m) e -> Container (c (PrimState m) e)
forall vec. vec -> Container vec
Container c (PrimState m) e
vec)
  (a, c (PrimState m) e) -> m (a, c (PrimState m) e)
forall (m :: * -> *) a. Monad m => a -> m a
return (a
a, c (PrimState m) e
vec)