module Data.PrimitiveArray.Class where
import Control.Applicative (Applicative, pure, (<$>), (<*>))
import Control.Exception (assert)
import Control.Monad.Except
import Control.Monad (forM_)
import Control.Monad.Primitive (PrimMonad, liftPrim)
import Control.Monad.ST (runST)
import Data.Proxy
import Data.Vector.Fusion.Util
import Debug.Trace
import GHC.Generics (Generic)
import Prelude as P
import qualified Data.Vector.Fusion.Stream.Monadic as SM
import Data.PrimitiveArray.Index.Class
data family MutArr (m :: * -> *) (arr :: *) :: *
class (Index sh) => MPrimArrayOps arr sh elm where
upperBoundM :: MutArr m (arr sh elm) -> LimitType sh
fromListM :: PrimMonad m => LimitType sh -> [elm] -> m (MutArr m (arr sh elm))
newM :: PrimMonad m => LimitType sh -> m (MutArr m (arr sh elm))
newWithM :: PrimMonad m => LimitType sh -> elm -> m (MutArr m (arr sh elm))
readM :: PrimMonad m => MutArr m (arr sh elm) -> sh -> m elm
writeM :: PrimMonad m => MutArr m (arr sh elm) -> sh -> elm -> m ()
class (Index sh) => PrimArrayOps arr sh elm where
upperBound :: arr sh elm -> LimitType sh
unsafeFreeze :: PrimMonad m => MutArr m (arr sh elm) -> m (arr sh elm)
unsafeThaw :: PrimMonad m => arr sh elm -> m (MutArr m (arr sh elm))
unsafeIndex :: arr sh elm -> sh -> elm
transformShape :: (Index sh') => (LimitType sh -> LimitType sh') -> arr sh elm -> arr sh' elm
class (Index sh) => PrimArrayMap arr sh e e' where
map :: (e -> e') -> arr sh e -> arr sh e'
data PAErrors
= PAEUpperBound
deriving (Eq,Generic)
instance Show PAErrors where
show (PAEUpperBound) = "Upper bound is too large for @Int@ size!"
(!) :: PrimArrayOps arr sh elm => arr sh elm -> sh -> elm
(!) arr idx = assert (inBounds (upperBound arr) idx) $ unsafeIndex arr idx
{-# INLINE (!) #-}
inBoundsM :: (Monad m, MPrimArrayOps arr sh elm) => MutArr m (arr sh elm) -> sh -> Bool
inBoundsM marr idx = inBounds (upperBoundM marr) idx
{-# INLINE inBoundsM #-}
fromAssocsM
:: (PrimMonad m, MPrimArrayOps arr sh elm)
=> LimitType sh -> elm -> [(sh,elm)] -> m (MutArr m (arr sh elm))
fromAssocsM ub def xs = do
ma <- newWithM ub def
forM_ xs $ \(k,v) -> writeM ma k v
return ma
{-# INLINE fromAssocsM #-}
newWithPA
∷ (PrimMonad m, MPrimArrayOps arr sh elm, PrimArrayOps arr sh elm)
⇒ LimitType sh
→ elm
→ m (arr sh elm)
newWithPA ub def = do
ma ← newWithM ub def
unsafeFreeze ma
{-# Inlinable newWithPA #-}
safeNewWithPA
∷ forall m arr sh elm
. (PrimMonad m, MonadError PAErrors m, MPrimArrayOps arr sh elm, PrimArrayOps arr sh elm)
⇒ LimitType sh
→ elm
→ m (arr sh elm)
safeNewWithPA ub def = do
case runExcept $ sizeIsValid maxBound [totalSize ub] of
Left (SizeError _) → throwError PAEUpperBound
Right (CellSize _) → newWithPA ub def
{-# Inlinable safeNewWithPA #-}
assocs :: forall arr sh elm . (IndexStream sh, PrimArrayOps arr sh elm) => arr sh elm -> [(sh,elm)]
assocs arr = P.map (\k -> (k,unsafeIndex arr k)) . unId . SM.toList $ streamUp zeroBound' (upperBound arr) where
{-# INLINE assocs #-}
fromList :: (PrimArrayOps arr sh elm, MPrimArrayOps arr sh elm) => LimitType sh -> [elm] -> arr sh elm
fromList ub xs = runST $ fromListM ub xs >>= unsafeFreeze
{-# INLINE fromList #-}
fromAssocs :: (PrimArrayOps arr sh elm, MPrimArrayOps arr sh elm) => LimitType sh -> elm -> [(sh,elm)] -> arr sh elm
fromAssocs ub def xs = runST $ fromAssocsM ub def xs >>= unsafeFreeze
{-# INLINE fromAssocs #-}
toList :: forall arr sh elm . (IndexStream sh, PrimArrayOps arr sh elm) => arr sh elm -> [elm]
toList arr = let ub = upperBound arr in P.map ((!) arr) . unId . SM.toList $ streamUp zeroBound' ub
{-# INLINE toList #-}
class FreezeTables m t where
type Frozen t :: *
freezeTables :: t -> m (Frozen t)
instance Applicative m => FreezeTables m Z where
type Frozen Z = Z
freezeTables Z = pure Z
{-# INLINE freezeTables #-}
instance (Functor m, Applicative m, Monad m, PrimMonad m, FreezeTables m ts, PrimArrayOps arr sh elm) => FreezeTables m (ts:.MutArr m (arr sh elm)) where
type Frozen (ts:.MutArr m (arr sh elm)) = Frozen ts :. arr sh elm
freezeTables (ts:.t) = (:.) <$> freezeTables ts <*> unsafeFreeze t
{-# INLINE freezeTables #-}