{-# LANGUAGE BangPatterns #-}
{-# LANGUAGE CPP #-}
{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE FlexibleInstances #-}
{-# LANGUAGE MultiParamTypeClasses #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE TypeFamilies #-}
{-# LANGUAGE UndecidableInstances #-}
module Data.Massiv.Array.Manifest.Unboxed
( U (..)
, Unbox
, Array(..)
, MArray(..)
, toUnboxedVector
, toUnboxedMVector
, fromUnboxedVector
, fromUnboxedMVector
) where
import Control.DeepSeq (NFData(..), deepseq)
import Data.Massiv.Array.Delayed.Pull (eqArrays, compareArrays)
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 Data.Massiv.Core.Operations
import Data.Vector.Unboxed (Unbox)
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)
}
data instance MArray s U ix e = MUArray !(Sz ix) !(VU.MVector s e)
instance (Ragged L ix e, Show e, 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 e.
(Ragged L ix e, Load r ix e, Load r' ix e, Source r' 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 NFData ix => NFData (MArray s U ix e) where
rnf :: MArray s U ix e -> ()
rnf (MUArray sz mv) = Sz ix
sz Sz ix -> MVector s e -> MVector s e
forall a b. NFData a => a -> b -> b
`deepseq` MVector s e
mv MVector s e -> () -> ()
forall a b. NFData a => a -> b -> b
`deepseq` ()
{-# INLINE rnf #-}
instance Strategy U where
getComp :: Array U ix e -> Comp
getComp = Array U ix e -> Comp
forall ix e. Array U ix e -> Comp
uComp
{-# INLINE getComp #-}
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 #-}
instance (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 ix r1 e1 r2 e2.
(Index ix, Source r1 e1, Source r2 e2) =>
(e1 -> e2 -> Bool) -> Array r1 ix e1 -> Array r2 ix e2 -> Bool
eqArrays e -> e -> Bool
forall a. Eq a => a -> a -> Bool
(==)
{-# INLINE (==) #-}
instance (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 ix r1 e1 r2 e2.
(Index ix, Source r1 e1, Source r2 e2) =>
(e1 -> e2 -> Ordering)
-> Array r1 ix e1 -> Array r2 ix e2 -> Ordering
compareArrays e -> e -> Ordering
forall a. Ord a => a -> a -> Ordering
compare
{-# INLINE compare #-}
instance Unbox e => Source U 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 #-}
unsafeOuterSlice :: Array U ix e -> Sz (Lower ix) -> Int -> Array U (Lower ix) e
unsafeOuterSlice (UArray c _ v) Sz (Lower ix)
szL Int
i =
let k :: Int
k = Sz (Lower ix) -> Int
forall ix. Index ix => Sz ix -> Int
totalElem Sz (Lower ix)
szL
in Comp -> Sz (Lower ix) -> Vector e -> Array U (Lower ix) e
forall ix e. Comp -> Sz ix -> Vector e -> Array U ix e
UArray Comp
c Sz (Lower ix)
szL (Vector e -> Array U (Lower ix) e)
-> Vector e -> Array U (Lower ix) 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 Int -> Int -> Int
forall a. Num a => a -> a -> a
* Int
k) Int
k Vector e
v
{-# INLINE unsafeOuterSlice #-}
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 => Shape U ix where
maxLinearSize :: Array U ix e -> Maybe Sz1
maxLinearSize = Sz1 -> Maybe Sz1
forall a. a -> Maybe a
Just (Sz1 -> Maybe Sz1)
-> (Array U ix e -> Sz1) -> Array U ix e -> Maybe Sz1
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Int -> Sz1
forall ix. ix -> Sz ix
SafeSz (Int -> Sz1) -> (Array U ix e -> Int) -> Array U ix e -> Sz1
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Array U ix e -> Int
forall ix r e. (Index ix, Size r) => Array r ix e -> Int
elemsCount
{-# INLINE maxLinearSize #-}
instance Size U where
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 #-}
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 (Unbox e, Index ix) => Load U ix e where
makeArrayLinear :: Comp -> Sz ix -> (Int -> e) -> Array U ix e
makeArrayLinear !Comp
comp !Sz ix
sz Int -> 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 -> (Int -> IO e) -> IO (Array U ix e)
forall r ix e (m :: * -> *).
(MonadUnliftIO m, Manifest r e, Index ix) =>
Comp -> Sz ix -> (Int -> m e) -> m (Array r ix e)
generateArrayLinear Comp
comp Sz ix
sz (e -> IO e
forall (f :: * -> *) a. Applicative f => a -> f a
pure (e -> IO e) -> (Int -> e) -> Int -> IO e
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Int -> e
f)
{-# INLINE makeArrayLinear #-}
replicate :: Comp -> Sz ix -> e -> Array U ix e
replicate Comp
comp !Sz ix
sz !e
e = (forall s. ST s (Array U ix e)) -> Array U ix e
forall a. (forall s. ST s a) -> a
runST (Sz ix -> e -> ST s (MArray (PrimState (ST s)) U ix e)
forall r e ix (m :: * -> *).
(Manifest r e, Index ix, PrimMonad m) =>
Sz ix -> e -> m (MArray (PrimState m) r ix e)
newMArray Sz ix
sz e
e ST s (MArray s U ix e)
-> (MArray s U ix e -> ST s (Array U ix e)) -> ST s (Array U ix e)
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= Comp -> MArray (PrimState (ST s)) U ix e -> ST s (Array U ix e)
forall r e ix (m :: * -> *).
(Manifest r e, Index ix, PrimMonad m) =>
Comp -> MArray (PrimState m) r ix e -> m (Array r ix e)
unsafeFreeze Comp
comp)
{-# INLINE replicate #-}
iterArrayLinearST_ :: Scheduler s () -> Array U ix e -> (Int -> e -> ST s ()) -> ST s ()
iterArrayLinearST_ !Scheduler s ()
scheduler !Array U ix e
arr =
Scheduler s ()
-> Int -> (Int -> e) -> (Int -> e -> ST s ()) -> ST s ()
forall s (m :: * -> *) b.
MonadPrimBase s m =>
Scheduler s () -> Int -> (Int -> b) -> (Int -> b -> m ()) -> m ()
splitLinearlyWith_ Scheduler s ()
scheduler (Array U ix e -> Int
forall ix r e. (Index ix, Size r) => Array r ix e -> Int
elemsCount Array U ix e
arr) (Array U ix e -> Int -> e
forall r e ix. (Source r e, Index ix) => Array r ix e -> Int -> e
unsafeLinearIndex Array U ix e
arr)
{-# INLINE iterArrayLinearST_ #-}
instance (Unbox e, Index ix) => StrideLoad U ix e
instance Unbox e => Manifest U 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 #-}
sizeOfMArray :: MArray s U ix e -> Sz ix
sizeOfMArray (MUArray sz _) = Sz ix
sz
{-# INLINE sizeOfMArray #-}
unsafeResizeMArray :: Sz ix' -> MArray s U ix e -> MArray s U ix' e
unsafeResizeMArray Sz ix'
sz (MUArray _ mv) = Sz ix' -> MVector s e -> MArray s U ix' e
forall s ix e. Sz ix -> MVector s e -> MArray s U ix e
MUArray Sz ix'
sz MVector s e
mv
{-# INLINE unsafeResizeMArray #-}
unsafeLinearSliceMArray :: Int -> Sz1 -> MArray s U ix e -> MVector s U e
unsafeLinearSliceMArray Int
i Sz1
k (MUArray _ mv) = Sz1 -> MVector s e -> MVector s U e
forall s ix e. Sz ix -> MVector s e -> MArray s U ix e
MUArray Sz1
k (MVector s e -> MVector s U e) -> MVector s e -> MVector s U e
forall a b. (a -> b) -> a -> b
$ Int -> Int -> MVector s e -> MVector s e
forall a s. Unbox a => Int -> Int -> MVector s a -> MVector s a
MVU.unsafeSlice Int
i (Sz1 -> Int
forall ix. Sz ix -> ix
unSz Sz1
k) MVector s e
mv
{-# INLINE unsafeLinearSliceMArray #-}
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("(Manifest 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("(Manifest 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, 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, Index ix, Source r 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, Index ix, Source r e) =>
Array r ix e -> Steps m (ix, e)
S.isteps
{-# INLINE toStreamIx #-}
instance (Unbox e, IsList (Array 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.
(HasCallStack, Ragged L ix e, Manifest r 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.
(Ragged L ix e, Shape r ix, Source r e) =>
Array r ix e -> Array L ix e
toListArray
{-# INLINE toList #-}
instance (VU.Unbox e, Num e) => FoldNumeric U e where
unsafeDotProduct :: Array U ix e -> Array U ix e -> e
unsafeDotProduct = Array U ix e -> Array U ix e -> e
forall e ix r.
(Num e, Index ix, Source r e) =>
Array r ix e -> Array r ix e -> e
defaultUnsafeDotProduct
{-# INLINE unsafeDotProduct #-}
powerSumArray :: Array U ix e -> Int -> e
powerSumArray = Array U ix e -> Int -> e
forall ix r e.
(Index ix, Source r e, Num e) =>
Array r ix e -> Int -> e
defaultPowerSumArray
{-# INLINE powerSumArray #-}
foldArray :: (e -> e -> e) -> e -> Array U ix e -> e
foldArray = (e -> e -> e) -> e -> Array U ix e -> e
forall ix r e.
(Index ix, Source r e) =>
(e -> e -> e) -> e -> Array r ix e -> e
defaultFoldArray
{-# INLINE foldArray #-}
instance (VU.Unbox e, Num e) => Numeric U e where
unsafeLiftArray :: (e -> e) -> Array U ix e -> Array U ix e
unsafeLiftArray = (e -> e) -> Array U ix e -> Array U ix e
forall r ix e.
(Load r ix e, Source r e) =>
(e -> e) -> Array r ix e -> Array r ix e
defaultUnsafeLiftArray
{-# INLINE unsafeLiftArray #-}
unsafeLiftArray2 :: (e -> e -> e) -> Array U ix e -> Array U ix e -> Array U ix e
unsafeLiftArray2 = (e -> e -> e) -> Array U ix e -> Array U ix e -> Array U ix e
forall r ix e.
(Load r ix e, Source r e) =>
(e -> e -> e) -> Array r ix e -> Array r ix e -> Array r ix e
defaultUnsafeLiftArray2
{-# INLINE unsafeLiftArray2 #-}
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 => Comp -> VU.Vector e -> Vector U e
fromUnboxedVector :: Comp -> Vector e -> Vector U e
fromUnboxedVector Comp
comp Vector e
v = Comp -> Sz1 -> Vector e -> Vector U e
forall ix e. Comp -> Sz ix -> Vector e -> Array U ix e
UArray Comp
comp (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 :: Unbox e => VU.MVector s e -> MVector s U e
fromUnboxedMVector :: MVector s e -> MVector s U e
fromUnboxedMVector MVector s e
mv = Sz1 -> MVector s e -> MVector s U 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 #-}