{-# LANGUAGE BangPatterns #-}
{-# LANGUAGE CPP #-}
{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE FlexibleInstances #-}
{-# LANGUAGE MultiParamTypeClasses #-}
{-# LANGUAGE NamedFieldPuns #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE TypeFamilies #-}
{-# LANGUAGE UndecidableInstances #-}
module Data.Massiv.Array.Manifest.Unboxed
( U (..)
, VU.Unbox
, Array(..)
, toUnboxedVector
, toUnboxedMVector
, fromUnboxedVector
, fromUnboxedMVector
) where
import Control.DeepSeq (NFData(..), deepseq)
import Data.Massiv.Array.Delayed.Pull (eq, ord)
import Data.Massiv.Array.Manifest.Internal (M, toManifest)
import Data.Massiv.Array.Manifest.List as A
import Data.Massiv.Vector.Stream as S (steps, isteps)
import Data.Massiv.Array.Mutable
import Data.Massiv.Core.Common
import Data.Massiv.Core.List
import qualified Data.Vector.Generic.Mutable as VGM
import qualified Data.Vector.Unboxed as VU
import qualified Data.Vector.Unboxed.Mutable as MVU
import GHC.Exts as GHC (IsList(..))
import Prelude hiding (mapM)
import System.IO.Unsafe (unsafePerformIO)
#include "massiv.h"
data U = U deriving Int -> U -> ShowS
[U] -> ShowS
U -> String
(Int -> U -> ShowS) -> (U -> String) -> ([U] -> ShowS) -> Show U
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
showList :: [U] -> ShowS
$cshowList :: [U] -> ShowS
show :: U -> String
$cshow :: U -> String
showsPrec :: Int -> U -> ShowS
$cshowsPrec :: Int -> U -> ShowS
Show
data instance Array U ix e = UArray { Array U ix e -> Comp
uComp :: !Comp
, Array U ix e -> Sz ix
uSize :: !(Sz ix)
, Array U ix e -> Vector e
uData :: !(VU.Vector e)
}
instance (Ragged L ix e, Show e, VU.Unbox e) => Show (Array U ix e) where
showsPrec :: Int -> Array U ix e -> ShowS
showsPrec = (Array U ix e -> Array U ix e) -> Int -> Array U ix e -> ShowS
forall r r' ix ix' e.
(Ragged L ix' e, Load r ix e, Source r' ix' e, Show e) =>
(Array r ix e -> Array r' ix' e) -> Int -> Array r ix e -> ShowS
showsArrayPrec Array U ix e -> Array U ix e
forall a. a -> a
id
showList :: [Array U ix e] -> ShowS
showList = [Array U ix e] -> ShowS
forall arr. Show arr => [arr] -> ShowS
showArrayList
instance NFData ix => NFData (Array U ix e) where
rnf :: Array U ix e -> ()
rnf (UArray c sz v) = Comp
c Comp -> Sz ix -> Sz ix
forall a b. NFData a => a -> b -> b
`deepseq` Sz ix
sz Sz ix -> Vector e -> Vector e
forall a b. NFData a => a -> b -> b
`deepseq` Vector e
v Vector e -> () -> ()
forall a b. NFData a => a -> b -> b
`deepseq` ()
{-# INLINE rnf #-}
instance (VU.Unbox e, Index ix) => Construct U ix e where
setComp :: Comp -> Array U ix e -> Array U ix e
setComp Comp
c Array U ix e
arr = Array U ix e
R:ArrayUixe ix e
arr { uComp :: Comp
uComp = Comp
c }
{-# INLINE setComp #-}
makeArray :: Comp -> Sz ix -> (ix -> e) -> Array U ix e
makeArray !Comp
comp !Sz ix
sz ix -> e
f = IO (Array U ix e) -> Array U ix e
forall a. IO a -> a
unsafePerformIO (IO (Array U ix e) -> Array U ix e)
-> IO (Array U ix e) -> Array U ix e
forall a b. (a -> b) -> a -> b
$ Comp -> Sz ix -> (ix -> IO e) -> IO (Array U ix e)
forall r ix e (m :: * -> *).
(MonadUnliftIO m, PrimMonad m, Mutable r ix e) =>
Comp -> Sz ix -> (ix -> m e) -> m (Array r ix e)
generateArray Comp
comp Sz ix
sz (e -> IO e
forall (m :: * -> *) a. Monad m => a -> m a
return (e -> IO e) -> (ix -> e) -> ix -> IO e
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ix -> e
f)
{-# INLINE makeArray #-}
instance (VU.Unbox e, Eq e, Index ix) => Eq (Array U ix e) where
== :: Array U ix e -> Array U ix e -> Bool
(==) = (e -> e -> Bool) -> Array U ix e -> Array U ix e -> Bool
forall r1 ix e1 r2 e2.
(Source r1 ix e1, Source r2 ix e2) =>
(e1 -> e2 -> Bool) -> Array r1 ix e1 -> Array r2 ix e2 -> Bool
eq e -> e -> Bool
forall a. Eq a => a -> a -> Bool
(==)
{-# INLINE (==) #-}
instance (VU.Unbox e, Ord e, Index ix) => Ord (Array U ix e) where
compare :: Array U ix e -> Array U ix e -> Ordering
compare = (e -> e -> Ordering) -> Array U ix e -> Array U ix e -> Ordering
forall r1 ix e1 r2 e2.
(Source r1 ix e1, Source r2 ix e2) =>
(e1 -> e2 -> Ordering)
-> Array r1 ix e1 -> Array r2 ix e2 -> Ordering
ord e -> e -> Ordering
forall a. Ord a => a -> a -> Ordering
compare
{-# INLINE compare #-}
instance (VU.Unbox e, Index ix) => Source U ix e where
unsafeLinearIndex :: Array U ix e -> Int -> e
unsafeLinearIndex (UArray _ _ v) =
INDEX_CHECK("(Source U ix e).unsafeLinearIndex", Sz . VU.length, VU.unsafeIndex) v
{-# INLINE unsafeLinearIndex #-}
unsafeLinearSlice :: Int -> Sz1 -> Array U ix e -> Array U Int e
unsafeLinearSlice Int
i Sz1
k (UArray c _ v) = Comp -> Sz1 -> Vector e -> Array U Int e
forall ix e. Comp -> Sz ix -> Vector e -> Array U ix e
UArray Comp
c Sz1
k (Vector e -> Array U Int e) -> Vector e -> Array U Int e
forall a b. (a -> b) -> a -> b
$ Int -> Int -> Vector e -> Vector e
forall a. Unbox a => Int -> Int -> Vector a -> Vector a
VU.unsafeSlice Int
i (Sz1 -> Int
forall ix. Sz ix -> ix
unSz Sz1
k) Vector e
v
{-# INLINE unsafeLinearSlice #-}
instance Index ix => Resize U ix where
unsafeResize :: Sz ix' -> Array U ix e -> Array U ix' e
unsafeResize !Sz ix'
sz !Array U ix e
arr = Array U ix e
R:ArrayUixe ix e
arr { uSize :: Sz ix'
uSize = Sz ix'
sz }
{-# INLINE unsafeResize #-}
instance (VU.Unbox e, Index ix) => Extract U ix e where
unsafeExtract :: ix -> Sz ix -> Array U ix e -> Array (R U) ix e
unsafeExtract !ix
sIx !Sz ix
newSz !Array U ix e
arr = ix -> Sz ix -> Array M ix e -> Array (R M) ix e
forall r ix e.
Extract r ix e =>
ix -> Sz ix -> Array r ix e -> Array (R r) ix e
unsafeExtract ix
sIx Sz ix
newSz (Array U ix e -> Array M ix e
forall r ix e. Manifest r ix e => Array r ix e -> Array M ix e
toManifest Array U ix e
arr)
{-# INLINE unsafeExtract #-}
instance (VU.Unbox e, Index ix) => Load U ix e where
type R U = M
size :: Array U ix e -> Sz ix
size = Array U ix e -> Sz ix
forall ix e. Array U ix e -> Sz ix
uSize
{-# INLINE size #-}
getComp :: Array U ix e -> Comp
getComp = Array U ix e -> Comp
forall ix e. Array U ix e -> Comp
uComp
{-# INLINE getComp #-}
loadArrayM :: Scheduler m () -> Array U ix e -> (Int -> e -> m ()) -> m ()
loadArrayM !Scheduler m ()
scheduler !Array U ix e
arr = Scheduler m () -> Int -> (Int -> e) -> (Int -> e -> m ()) -> m ()
forall (m :: * -> *) b.
Monad m =>
Scheduler m () -> Int -> (Int -> b) -> (Int -> b -> m ()) -> m ()
splitLinearlyWith_ Scheduler m ()
scheduler (Array U ix e -> Int
forall r ix e. Load r ix e => Array r ix e -> Int
elemsCount Array U ix e
arr) (Array U ix e -> Int -> e
forall r ix e. Source r ix e => Array r ix e -> Int -> e
unsafeLinearIndex Array U ix e
arr)
{-# INLINE loadArrayM #-}
instance (VU.Unbox e, Index ix) => StrideLoad U ix e
instance {-# OVERLAPPING #-} VU.Unbox e => Slice U Ix1 e where
unsafeSlice :: Array U Int e -> Int -> Sz1 -> Dim -> m (Elt U Int e)
unsafeSlice Array U Int e
arr Int
i Sz1
_ Dim
_ = e -> m e
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Array U Int e -> Int -> e
forall r ix e. Source r ix e => Array r ix e -> Int -> e
unsafeLinearIndex Array U Int e
arr Int
i)
{-# INLINE unsafeSlice #-}
instance ( VU.Unbox e
, Index ix
, Index (Lower ix)
, Elt U ix e ~ Elt M ix e
, Elt M ix e ~ Array M (Lower ix) e
) =>
Slice U ix e where
unsafeSlice :: Array U ix e -> ix -> Sz ix -> Dim -> m (Elt U ix e)
unsafeSlice Array U ix e
arr = Array M ix e -> ix -> Sz ix -> Dim -> m (Elt M ix e)
forall r ix e (m :: * -> *).
(Slice r ix e, MonadThrow m) =>
Array r ix e -> ix -> Sz ix -> Dim -> m (Elt r ix e)
unsafeSlice (Array U ix e -> Array M ix e
forall r ix e. Manifest r ix e => Array r ix e -> Array M ix e
toManifest Array U ix e
arr)
{-# INLINE unsafeSlice #-}
instance {-# OVERLAPPING #-} VU.Unbox e => OuterSlice U Ix1 e where
unsafeOuterSlice :: Array U Int e -> Int -> Elt U Int e
unsafeOuterSlice = Array U Int e -> Int -> Elt U Int e
forall r ix e. Source r ix e => Array r ix e -> Int -> e
unsafeLinearIndex
{-# INLINE unsafeOuterSlice #-}
instance ( VU.Unbox e
, Index ix
, Index (Lower ix)
, Elt U ix e ~ Elt M ix e
, Elt M ix e ~ Array M (Lower ix) e
) =>
OuterSlice U ix e where
unsafeOuterSlice :: Array U ix e -> Int -> Elt U ix e
unsafeOuterSlice Array U ix e
arr = Array M ix e -> Int -> Elt M ix e
forall r ix e.
OuterSlice r ix e =>
Array r ix e -> Int -> Elt r ix e
unsafeOuterSlice (Array U ix e -> Array M ix e
forall r ix e. Manifest r ix e => Array r ix e -> Array M ix e
toManifest Array U ix e
arr)
{-# INLINE unsafeOuterSlice #-}
instance {-# OVERLAPPING #-} VU.Unbox e => InnerSlice U Ix1 e where
unsafeInnerSlice :: Array U Int e -> (Sz (Lower Int), Sz1) -> Int -> Elt U Int e
unsafeInnerSlice Array U Int e
arr (Sz (Lower Int), Sz1)
_ = Array U Int e -> Int -> e
forall r ix e. Source r ix e => Array r ix e -> Int -> e
unsafeLinearIndex Array U Int e
arr
{-# INLINE unsafeInnerSlice #-}
instance ( VU.Unbox e
, Index ix
, Index (Lower ix)
, Elt U ix e ~ Elt M ix e
, Elt M ix e ~ Array M (Lower ix) e
) =>
InnerSlice U ix e where
unsafeInnerSlice :: Array U ix e -> (Sz (Lower ix), Sz1) -> Int -> Elt U ix e
unsafeInnerSlice Array U ix e
arr = Array M ix e -> (Sz (Lower ix), Sz1) -> Int -> Elt M ix e
forall r ix e.
InnerSlice r ix e =>
Array r ix e -> (Sz (Lower ix), Sz1) -> Int -> Elt r ix e
unsafeInnerSlice (Array U ix e -> Array M ix e
forall r ix e. Manifest r ix e => Array r ix e -> Array M ix e
toManifest Array U ix e
arr)
{-# INLINE unsafeInnerSlice #-}
instance (VU.Unbox e, Index ix) => Manifest U ix e where
unsafeLinearIndexM :: Array U ix e -> Int -> e
unsafeLinearIndexM (UArray _ _ v) =
INDEX_CHECK("(Manifest U ix e).unsafeLinearIndexM", Sz . VU.length, VU.unsafeIndex) v
{-# INLINE unsafeLinearIndexM #-}
instance (VU.Unbox e, Index ix) => Mutable U ix e where
data MArray s U ix e = MUArray !(Sz ix) !(VU.MVector s e)
msize :: MArray s U ix e -> Sz ix
msize (MUArray sz _) = Sz ix
sz
{-# INLINE msize #-}
unsafeThaw :: Array U ix e -> m (MArray (PrimState m) U ix e)
unsafeThaw (UArray _ sz v) = Sz ix -> MVector (PrimState m) e -> MArray (PrimState m) U ix e
forall s ix e. Sz ix -> MVector s e -> MArray s U ix e
MUArray Sz ix
sz (MVector (PrimState m) e -> MArray (PrimState m) U ix e)
-> m (MVector (PrimState m) e) -> m (MArray (PrimState m) U ix e)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Vector e -> m (MVector (PrimState m) e)
forall a (m :: * -> *).
(Unbox a, PrimMonad m) =>
Vector a -> m (MVector (PrimState m) a)
VU.unsafeThaw Vector e
v
{-# INLINE unsafeThaw #-}
unsafeFreeze :: Comp -> MArray (PrimState m) U ix e -> m (Array U ix e)
unsafeFreeze Comp
comp (MUArray sz v) = Comp -> Sz ix -> Vector e -> Array U ix e
forall ix e. Comp -> Sz ix -> Vector e -> Array U ix e
UArray Comp
comp Sz ix
sz (Vector e -> Array U ix e) -> m (Vector e) -> m (Array U ix e)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> MVector (PrimState m) e -> m (Vector e)
forall a (m :: * -> *).
(Unbox a, PrimMonad m) =>
MVector (PrimState m) a -> m (Vector a)
VU.unsafeFreeze MVector (PrimState m) e
v
{-# INLINE unsafeFreeze #-}
unsafeNew :: Sz ix -> m (MArray (PrimState m) U ix e)
unsafeNew Sz ix
sz = Sz ix -> MVector (PrimState m) e -> MArray (PrimState m) U ix e
forall s ix e. Sz ix -> MVector s e -> MArray s U ix e
MUArray Sz ix
sz (MVector (PrimState m) e -> MArray (PrimState m) U ix e)
-> m (MVector (PrimState m) e) -> m (MArray (PrimState m) U ix e)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Int -> m (MVector (PrimState m) e)
forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
Int -> m (MVector (PrimState m) a)
MVU.unsafeNew (Sz ix -> Int
forall ix. Index ix => Sz ix -> Int
totalElem Sz ix
sz)
{-# INLINE unsafeNew #-}
initialize :: MArray (PrimState m) U ix e -> m ()
initialize (MUArray _ marr) = MVector (PrimState m) e -> m ()
forall (v :: * -> * -> *) a (m :: * -> *).
(MVector v a, PrimMonad m) =>
v (PrimState m) a -> m ()
VGM.basicInitialize MVector (PrimState m) e
marr
{-# INLINE initialize #-}
unsafeLinearCopy :: MArray (PrimState m) U ix' e
-> Int -> MArray (PrimState m) U ix e -> Int -> Sz1 -> m ()
unsafeLinearCopy (MUArray _ mvFrom) Int
iFrom (MUArray _ mvTo) Int
iTo (Sz Int
k) =
MVector (PrimState m) e -> MVector (PrimState m) e -> m ()
forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
MVector (PrimState m) a -> MVector (PrimState m) a -> m ()
MVU.unsafeCopy (Int -> Int -> MVector (PrimState m) e -> MVector (PrimState m) e
forall a s. Unbox a => Int -> Int -> MVector s a -> MVector s a
MVU.unsafeSlice Int
iTo Int
k MVector (PrimState m) e
mvTo) (Int -> Int -> MVector (PrimState m) e -> MVector (PrimState m) e
forall a s. Unbox a => Int -> Int -> MVector s a -> MVector s a
MVU.unsafeSlice Int
iFrom Int
k MVector (PrimState m) e
mvFrom)
{-# INLINE unsafeLinearCopy #-}
unsafeLinearRead :: MArray (PrimState m) U ix e -> Int -> m e
unsafeLinearRead (MUArray _ mv) =
INDEX_CHECK("(Mutable U ix e).unsafeLinearRead", Sz . MVU.length, MVU.unsafeRead) mv
{-# INLINE unsafeLinearRead #-}
unsafeLinearWrite :: MArray (PrimState m) U ix e -> Int -> e -> m ()
unsafeLinearWrite (MUArray _ mv) =
INDEX_CHECK("(Mutable U ix e).unsafeLinearWrite", Sz . MVU.length, MVU.unsafeWrite) mv
{-# INLINE unsafeLinearWrite #-}
unsafeLinearGrow :: MArray (PrimState m) U ix e
-> Sz ix -> m (MArray (PrimState m) U ix e)
unsafeLinearGrow (MUArray _ mv) Sz ix
sz = Sz ix -> MVector (PrimState m) e -> MArray (PrimState m) U ix e
forall s ix e. Sz ix -> MVector s e -> MArray s U ix e
MUArray Sz ix
sz (MVector (PrimState m) e -> MArray (PrimState m) U ix e)
-> m (MVector (PrimState m) e) -> m (MArray (PrimState m) U ix e)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> MVector (PrimState m) e -> Int -> m (MVector (PrimState m) e)
forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
MVector (PrimState m) a -> Int -> m (MVector (PrimState m) a)
MVU.unsafeGrow MVector (PrimState m) e
mv (Sz ix -> Int
forall ix. Index ix => Sz ix -> Int
totalElem Sz ix
sz)
{-# INLINE unsafeLinearGrow #-}
instance (Index ix, VU.Unbox e) => Stream U ix e where
toStream :: Array U ix e -> Steps Id e
toStream = Array U ix e -> Steps Id e
forall r ix e (m :: * -> *).
(Monad m, Source r ix e) =>
Array r ix e -> Steps m e
S.steps
{-# INLINE toStream #-}
toStreamIx :: Array U ix e -> Steps Id (ix, e)
toStreamIx = Array U ix e -> Steps Id (ix, e)
forall r ix e (m :: * -> *).
(Monad m, Source r ix e) =>
Array r ix e -> Steps m (ix, e)
S.isteps
{-# INLINE toStreamIx #-}
instance ( VU.Unbox e
, IsList (Array L ix e)
, Nested LN ix e
, Nested L ix e
, Ragged L ix e
) =>
IsList (Array U ix e) where
type Item (Array U ix e) = Item (Array L ix e)
fromList :: [Item (Array U ix e)] -> Array U ix e
fromList = Comp -> [ListItem ix e] -> Array U ix e
forall r ix e.
(Nested LN ix e, Ragged L ix e, Mutable r ix e) =>
Comp -> [ListItem ix e] -> Array r ix e
A.fromLists' Comp
Seq
{-# INLINE fromList #-}
toList :: Array U ix e -> [Item (Array U ix e)]
toList = Array L ix e -> [ListItem ix e]
forall l. IsList l => l -> [Item l]
GHC.toList (Array L ix e -> [ListItem ix e])
-> (Array U ix e -> Array L ix e)
-> Array U ix e
-> [ListItem ix e]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Array U ix e -> Array L ix e
forall ix e r.
(Construct L ix e, Source r ix e) =>
Array r ix e -> Array L ix e
toListArray
{-# INLINE toList #-}
toUnboxedVector :: Array U ix e -> VU.Vector e
toUnboxedVector :: Array U ix e -> Vector e
toUnboxedVector = Array U ix e -> Vector e
forall ix e. Array U ix e -> Vector e
uData
{-# INLINE toUnboxedVector #-}
toUnboxedMVector :: MArray s U ix e -> VU.MVector s e
toUnboxedMVector :: MArray s U ix e -> MVector s e
toUnboxedMVector (MUArray _ mv) = MVector s e
mv
{-# INLINE toUnboxedMVector #-}
fromUnboxedVector :: VU.Unbox e => VU.Vector e -> Array U Ix1 e
fromUnboxedVector :: Vector e -> Array U Int e
fromUnboxedVector Vector e
v = Comp -> Sz1 -> Vector e -> Array U Int e
forall ix e. Comp -> Sz ix -> Vector e -> Array U ix e
UArray Comp
Seq (Int -> Sz1
forall ix. ix -> Sz ix
SafeSz (Vector e -> Int
forall a. Unbox a => Vector a -> Int
VU.length Vector e
v)) Vector e
v
{-# INLINE fromUnboxedVector #-}
fromUnboxedMVector :: VU.Unbox e => VU.MVector s e -> MArray s U Ix1 e
fromUnboxedMVector :: MVector s e -> MArray s U Int e
fromUnboxedMVector MVector s e
mv = Sz1 -> MVector s e -> MArray s U Int e
forall s ix e. Sz ix -> MVector s e -> MArray s U ix e
MUArray (Int -> Sz1
forall ix. ix -> Sz ix
SafeSz (MVector s e -> Int
forall a s. Unbox a => MVector s a -> Int
MVU.length MVector s e
mv)) MVector s e
mv
{-# INLINE fromUnboxedMVector #-}