(original) (raw)

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

module Control.Monad.Memo.Vector.Unsafe (

VectorCache, VectorMemo, unsafeEvalVectorMemo, unsafeRunVectorMemo,

UVectorCache, UVectorMemo, unsafeEvalUVectorMemo, unsafeRunUVectorMemo,

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 }

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)

type Vector = M.MVector

type VectorCache s e = Cache Vector s e

class MaybeLike e v => VectorMemo v e | v -> e

unsafeEvalVectorMemo :: (PrimMonad m, VectorMemo v e) => VectorCache (PrimState m) e m a -> Int
-> m a
{-# 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

unsafeRunVectorMemo :: (PrimMonad m, VectorMemo v e) => VectorCache (PrimState m) e m a -> Int
-> m (a, Vector (PrimState m) e)
{-# 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

type UVector = UM.MVector

type UVectorCache s e = Cache UVector s e

class MaybeLike e v => UVectorMemo v e | v -> e

unsafeEvalUVectorMemo :: (PrimMonad m, UVectorMemo v e, MVector UVector e) => UVectorCache (PrimState m) e m a -> Int
-> m a
{-# 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

unsafeRunUVectorMemo :: (PrimMonad m, UVectorMemo v e, MVector UVector e) => UVectorCache (PrimState m) e m a -> Int
-> m (a, UVector (PrimState m) e)
{-# 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)