{-# LANGUAGE BangPatterns #-}
{-# LANGUAGE CPP #-}
{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE FlexibleInstances #-}
{-# LANGUAGE MagicHash #-}
{-# LANGUAGE MultiParamTypeClasses #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE TypeFamilies #-}
{-# LANGUAGE UndecidableInstances #-}
module Data.Massiv.Array.Manifest.Storable
( S (..)
, Array(..)
, Storable
, toStorableVector
, toStorableMVector
, fromStorableVector
, fromStorableMVector
, withPtr
, unsafeWithPtr
, unsafeArrayToForeignPtr
, unsafeMArrayToForeignPtr
, unsafeArrayFromForeignPtr
, unsafeArrayFromForeignPtr0
, unsafeMArrayFromForeignPtr
, unsafeMArrayFromForeignPtr0
) where
import Control.DeepSeq (NFData(..), deepseq)
import Control.Monad.IO.Unlift
import Control.Monad.Primitive (unsafePrimToPrim)
import Data.Massiv.Array.Delayed.Pull (eqArrays, compareArrays)
import Data.Massiv.Array.Manifest.Internal
import Data.Massiv.Array.Manifest.Primitive (shrinkMutableByteArray)
import Data.Primitive.ByteArray (MutableByteArray(..))
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 qualified Data.Vector.Generic.Mutable as VGM
import qualified Data.Vector.Storable as VS
import qualified Data.Vector.Storable.Mutable as MVS
import Foreign.Ptr
import GHC.ForeignPtr (ForeignPtr(..), ForeignPtrContents(..))
import Foreign.ForeignPtr (withForeignPtr)
import Foreign.Storable
import Foreign.Marshal.Array (copyArray, advancePtr)
import GHC.Exts as GHC (IsList(..))
import Prelude hiding (mapM)
import System.IO.Unsafe (unsafePerformIO)
#include "massiv.h"
data S = S deriving Int -> S -> ShowS
[S] -> ShowS
S -> String
(Int -> S -> ShowS) -> (S -> String) -> ([S] -> ShowS) -> Show S
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
showList :: [S] -> ShowS
$cshowList :: [S] -> ShowS
show :: S -> String
$cshow :: S -> String
showsPrec :: Int -> S -> ShowS
$cshowsPrec :: Int -> S -> ShowS
Show
data instance Array S ix e = SArray { Array S ix e -> Comp
sComp :: !Comp
, Array S ix e -> Sz ix
sSize :: !(Sz ix)
, Array S ix e -> Vector e
sData :: !(VS.Vector e)
}
instance (Ragged L ix e, Show e, Storable e) => Show (Array S ix e) where
showsPrec :: Int -> Array S ix e -> ShowS
showsPrec = (Array S ix e -> Array S ix e) -> Int -> Array S 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 S ix e -> Array S ix e
forall a. a -> a
id
showList :: [Array S ix e] -> ShowS
showList = [Array S ix e] -> ShowS
forall arr. Show arr => [arr] -> ShowS
showArrayList
instance NFData ix => NFData (Array S ix e) where
rnf :: Array S ix e -> ()
rnf (SArray 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 (Storable e, Eq e, Index ix) => Eq (Array S ix e) where
== :: Array S ix e -> Array S ix e -> Bool
(==) = (e -> e -> Bool) -> Array S ix e -> Array S 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
eqArrays e -> e -> Bool
forall a. Eq a => a -> a -> Bool
(==)
{-# INLINE (==) #-}
instance (Storable e, Ord e, Index ix) => Ord (Array S ix e) where
compare :: Array S ix e -> Array S ix e -> Ordering
compare = (e -> e -> Ordering) -> Array S ix e -> Array S 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
compareArrays e -> e -> Ordering
forall a. Ord a => a -> a -> Ordering
compare
{-# INLINE compare #-}
instance (Storable e, Index ix) => Construct S ix e where
setComp :: Comp -> Array S ix e -> Array S ix e
setComp Comp
c Array S ix e
arr = Array S ix e
R:ArraySixe ix e
arr { sComp :: Comp
sComp = Comp
c }
{-# INLINE setComp #-}
makeArray :: Comp -> Sz ix -> (ix -> e) -> Array S ix e
makeArray !Comp
comp !Sz ix
sz ix -> e
f = IO (Array S ix e) -> Array S ix e
forall a. IO a -> a
unsafePerformIO (IO (Array S ix e) -> Array S ix e)
-> IO (Array S ix e) -> Array S ix e
forall a b. (a -> b) -> a -> b
$ Comp -> Sz ix -> (ix -> IO e) -> IO (Array S 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 (Storable e, Index ix) => Source S ix e where
unsafeLinearIndex :: Array S ix e -> Int -> e
unsafeLinearIndex (SArray _ _ v) =
INDEX_CHECK("(Source S ix e).unsafeLinearIndex", Sz . VS.length, VS.unsafeIndex) v
{-# INLINE unsafeLinearIndex #-}
unsafeLinearSlice :: Int -> Sz1 -> Array S ix e -> Array S Int e
unsafeLinearSlice Int
i Sz1
k (SArray c _ v) = Comp -> Sz1 -> Vector e -> Array S Int e
forall ix e. Comp -> Sz ix -> Vector e -> Array S ix e
SArray Comp
c Sz1
k (Vector e -> Array S Int e) -> Vector e -> Array S Int e
forall a b. (a -> b) -> a -> b
$ Int -> Int -> Vector e -> Vector e
forall a. Storable a => Int -> Int -> Vector a -> Vector a
VS.unsafeSlice Int
i (Sz1 -> Int
forall ix. Sz ix -> ix
unSz Sz1
k) Vector e
v
{-# INLINE unsafeLinearSlice #-}
instance Index ix => Resize S ix where
unsafeResize :: Sz ix' -> Array S ix e -> Array S ix' e
unsafeResize !Sz ix'
sz !Array S ix e
arr = Array S ix e
R:ArraySixe ix e
arr { sSize :: Sz ix'
sSize = Sz ix'
sz }
{-# INLINE unsafeResize #-}
instance (Storable e, Index ix) => Extract S ix e where
unsafeExtract :: ix -> Sz ix -> Array S ix e -> Array (R S) ix e
unsafeExtract !ix
sIx !Sz ix
newSz !Array S 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 S ix e -> Array M ix e
forall r ix e. Manifest r ix e => Array r ix e -> Array M ix e
toManifest Array S ix e
arr)
{-# INLINE unsafeExtract #-}
instance ( Storable e
, Index ix
, Index (Lower ix)
, Elt M ix e ~ Array M (Lower ix) e
, Elt S ix e ~ Array M (Lower ix) e
) =>
OuterSlice S ix e where
unsafeOuterSlice :: Array S ix e -> Int -> Elt S ix e
unsafeOuterSlice Array S 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 S ix e -> Array M ix e
forall r ix e. Manifest r ix e => Array r ix e -> Array M ix e
toManifest Array S ix e
arr)
{-# INLINE unsafeOuterSlice #-}
instance ( Storable e
, Index ix
, Index (Lower ix)
, Elt M ix e ~ Array M (Lower ix) e
, Elt S ix e ~ Array M (Lower ix) e
) =>
InnerSlice S ix e where
unsafeInnerSlice :: Array S ix e -> (Sz (Lower ix), Sz1) -> Int -> Elt S ix e
unsafeInnerSlice Array S 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 S ix e -> Array M ix e
forall r ix e. Manifest r ix e => Array r ix e -> Array M ix e
toManifest Array S ix e
arr)
{-# INLINE unsafeInnerSlice #-}
instance {-# OVERLAPPING #-} Storable e => Slice S Ix1 e where
unsafeSlice :: Array S Int e -> Int -> Sz1 -> Dim -> m (Elt S Int e)
unsafeSlice Array S Int e
arr Int
i Sz1
_ Dim
_ = e -> m e
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Array S Int e -> Int -> e
forall r ix e. Source r ix e => Array r ix e -> Int -> e
unsafeLinearIndex Array S Int e
arr Int
i)
{-# INLINE unsafeSlice #-}
instance (Index ix, Storable e) => Manifest S ix e where
unsafeLinearIndexM :: Array S ix e -> Int -> e
unsafeLinearIndexM (SArray _ _ v) =
INDEX_CHECK("(Manifest S ix e).unsafeLinearIndexM", Sz . VS.length, VS.unsafeIndex) v
{-# INLINE unsafeLinearIndexM #-}
instance (Index ix, Storable e) => Mutable S ix e where
data MArray s S ix e = MSArray !(Sz ix) !(VS.MVector s e)
msize :: MArray s S ix e -> Sz ix
msize (MSArray sz _) = Sz ix
sz
{-# INLINE msize #-}
unsafeThaw :: Array S ix e -> m (MArray (PrimState m) S ix e)
unsafeThaw (SArray _ sz v) = Sz ix -> MVector (PrimState m) e -> MArray (PrimState m) S ix e
forall s ix e. Sz ix -> MVector s e -> MArray s S ix e
MSArray Sz ix
sz (MVector (PrimState m) e -> MArray (PrimState m) S ix e)
-> m (MVector (PrimState m) e) -> m (MArray (PrimState m) S ix e)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Vector e -> m (MVector (PrimState m) e)
forall a (m :: * -> *).
(Storable a, PrimMonad m) =>
Vector a -> m (MVector (PrimState m) a)
VS.unsafeThaw Vector e
v
{-# INLINE unsafeThaw #-}
unsafeFreeze :: Comp -> MArray (PrimState m) S ix e -> m (Array S ix e)
unsafeFreeze Comp
comp (MSArray sz v) = Comp -> Sz ix -> Vector e -> Array S ix e
forall ix e. Comp -> Sz ix -> Vector e -> Array S ix e
SArray Comp
comp Sz ix
sz (Vector e -> Array S ix e) -> m (Vector e) -> m (Array S ix e)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> MVector (PrimState m) e -> m (Vector e)
forall a (m :: * -> *).
(Storable a, PrimMonad m) =>
MVector (PrimState m) a -> m (Vector a)
VS.unsafeFreeze MVector (PrimState m) e
v
{-# INLINE unsafeFreeze #-}
unsafeNew :: Sz ix -> m (MArray (PrimState m) S ix e)
unsafeNew Sz ix
sz = Sz ix -> MVector (PrimState m) e -> MArray (PrimState m) S ix e
forall s ix e. Sz ix -> MVector s e -> MArray s S ix e
MSArray Sz ix
sz (MVector (PrimState m) e -> MArray (PrimState m) S ix e)
-> m (MVector (PrimState m) e) -> m (MArray (PrimState m) S 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, Storable a) =>
Int -> m (MVector (PrimState m) a)
MVS.unsafeNew (Sz ix -> Int
forall ix. Index ix => Sz ix -> Int
totalElem Sz ix
sz)
{-# INLINE unsafeNew #-}
initialize :: MArray (PrimState m) S ix e -> m ()
initialize (MSArray _ 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 #-}
unsafeLinearRead :: MArray (PrimState m) S ix e -> Int -> m e
unsafeLinearRead (MSArray _ mv) =
INDEX_CHECK("(Mutable S ix e).unsafeLinearRead", Sz . MVS.length, MVS.unsafeRead) mv
{-# INLINE unsafeLinearRead #-}
unsafeLinearWrite :: MArray (PrimState m) S ix e -> Int -> e -> m ()
unsafeLinearWrite (MSArray _ mv) =
INDEX_CHECK("(Mutable S ix e).unsafeLinearWrite", Sz . MVS.length, MVS.unsafeWrite) mv
{-# INLINE unsafeLinearWrite #-}
unsafeLinearSet :: MArray (PrimState m) S ix e -> Int -> Sz1 -> e -> m ()
unsafeLinearSet (MSArray _ mv) Int
i Sz1
k = MVector (PrimState m) e -> e -> m ()
forall (v :: * -> * -> *) a (m :: * -> *).
(MVector v a, PrimMonad m) =>
v (PrimState m) a -> a -> m ()
VGM.basicSet (Int -> Int -> MVector (PrimState m) e -> MVector (PrimState m) e
forall a s. Storable a => Int -> Int -> MVector s a -> MVector s a
MVS.unsafeSlice Int
i (Sz1 -> Int
forall ix. Sz ix -> ix
unSz Sz1
k) MVector (PrimState m) e
mv)
{-# INLINE unsafeLinearSet #-}
unsafeLinearCopy :: MArray (PrimState m) S ix' e
-> Int -> MArray (PrimState m) S ix e -> Int -> Sz1 -> m ()
unsafeLinearCopy MArray (PrimState m) S ix' e
marrFrom Int
iFrom MArray (PrimState m) S ix e
marrTo Int
iTo (Sz Int
k) = do
let MSArray _ (MVS.MVector _ fpFrom) = MArray (PrimState m) S ix' e
marrFrom
MSArray _ (MVS.MVector _ fpTo) = MArray (PrimState m) S ix e
marrTo
IO () -> m ()
forall (m1 :: * -> *) (m2 :: * -> *) a.
(PrimBase m1, PrimMonad m2) =>
m1 a -> m2 a
unsafePrimToPrim (IO () -> m ()) -> IO () -> m ()
forall a b. (a -> b) -> a -> b
$
ForeignPtr e -> (Ptr e -> IO ()) -> IO ()
forall a b. ForeignPtr a -> (Ptr a -> IO b) -> IO b
withForeignPtr ForeignPtr e
fpFrom ((Ptr e -> IO ()) -> IO ()) -> (Ptr e -> IO ()) -> IO ()
forall a b. (a -> b) -> a -> b
$ \ Ptr e
ptrFrom ->
ForeignPtr e -> (Ptr e -> IO ()) -> IO ()
forall a b. ForeignPtr a -> (Ptr a -> IO b) -> IO b
withForeignPtr ForeignPtr e
fpTo ((Ptr e -> IO ()) -> IO ()) -> (Ptr e -> IO ()) -> IO ()
forall a b. (a -> b) -> a -> b
$ \ Ptr e
ptrTo -> do
let ptrFrom' :: Ptr e
ptrFrom' = Ptr e -> Int -> Ptr e
forall a. Storable a => Ptr a -> Int -> Ptr a
advancePtr Ptr e
ptrFrom Int
iFrom
ptrTo' :: Ptr e
ptrTo' = Ptr e -> Int -> Ptr e
forall a. Storable a => Ptr a -> Int -> Ptr a
advancePtr Ptr e
ptrTo Int
iTo
Ptr e -> Ptr e -> Int -> IO ()
forall a. Storable a => Ptr a -> Ptr a -> Int -> IO ()
copyArray Ptr e
ptrTo' Ptr e
ptrFrom' Int
k
{-# INLINE unsafeLinearCopy #-}
unsafeArrayLinearCopy :: Array S ix' e
-> Int -> MArray (PrimState m) S ix e -> Int -> Sz1 -> m ()
unsafeArrayLinearCopy Array S ix' e
arrFrom Int
iFrom MArray (PrimState m) S ix e
marrTo Int
iTo Sz1
sz = do
MArray (PrimState m) S ix' e
marrFrom <- Array S ix' e -> m (MArray (PrimState m) S ix' e)
forall r ix e (m :: * -> *).
(Mutable r ix e, PrimMonad m) =>
Array r ix e -> m (MArray (PrimState m) r ix e)
unsafeThaw Array S ix' e
arrFrom
MArray (PrimState m) S ix' e
-> Int -> MArray (PrimState m) S ix e -> Int -> Sz1 -> m ()
forall r ix e ix' (m :: * -> *).
(Mutable r ix e, Mutable r ix' e, PrimMonad m) =>
MArray (PrimState m) r ix' e
-> Int -> MArray (PrimState m) r ix e -> Int -> Sz1 -> m ()
unsafeLinearCopy MArray (PrimState m) S ix' e
marrFrom Int
iFrom MArray (PrimState m) S ix e
marrTo Int
iTo Sz1
sz
{-# INLINE unsafeArrayLinearCopy #-}
unsafeLinearShrink :: MArray (PrimState m) S ix e
-> Sz ix -> m (MArray (PrimState m) S ix e)
unsafeLinearShrink marr :: MArray (PrimState m) S ix e
marr@(MSArray _ mv@(MVS.MVector _ (ForeignPtr _ fpc))) Sz ix
sz = do
let shrinkMBA :: MutableByteArray RealWorld -> IO ()
shrinkMBA :: MutableByteArray RealWorld -> IO ()
shrinkMBA MutableByteArray RealWorld
mba = MutableByteArray (PrimState IO) -> Int -> IO ()
forall (m :: * -> *).
PrimMonad m =>
MutableByteArray (PrimState m) -> Int -> m ()
shrinkMutableByteArray MutableByteArray RealWorld
MutableByteArray (PrimState IO)
mba (Sz ix -> Int
forall ix. Index ix => Sz ix -> Int
totalElem Sz ix
sz Int -> Int -> Int
forall a. Num a => a -> a -> a
* e -> Int
forall a. Storable a => a -> Int
sizeOf (e
forall a. HasCallStack => a
undefined :: e))
{-# INLINE shrinkMBA #-}
case ForeignPtrContents
fpc of
MallocPtr MutableByteArray# RealWorld
mba# IORef Finalizers
_ -> do
IO () -> m ()
forall (m1 :: * -> *) (m2 :: * -> *) a.
(PrimBase m1, PrimMonad m2) =>
m1 a -> m2 a
unsafePrimToPrim (IO () -> m ()) -> IO () -> m ()
forall a b. (a -> b) -> a -> b
$ MutableByteArray RealWorld -> IO ()
shrinkMBA (MutableByteArray# RealWorld -> MutableByteArray RealWorld
forall s. MutableByteArray# s -> MutableByteArray s
MutableByteArray MutableByteArray# RealWorld
mba#)
MArray (PrimState m) S ix e -> m (MArray (PrimState m) S ix e)
forall (f :: * -> *) a. Applicative f => a -> f a
pure (MArray (PrimState m) S ix e -> m (MArray (PrimState m) S ix e))
-> MArray (PrimState m) S ix e -> m (MArray (PrimState m) S ix e)
forall a b. (a -> b) -> a -> b
$ Sz ix -> MVector (PrimState m) e -> MArray (PrimState m) S ix e
forall s ix e. Sz ix -> MVector s e -> MArray s S ix e
MSArray Sz ix
sz MVector (PrimState m) e
mv
PlainPtr MutableByteArray# RealWorld
mba# -> do
IO () -> m ()
forall (m1 :: * -> *) (m2 :: * -> *) a.
(PrimBase m1, PrimMonad m2) =>
m1 a -> m2 a
unsafePrimToPrim (IO () -> m ()) -> IO () -> m ()
forall a b. (a -> b) -> a -> b
$ MutableByteArray RealWorld -> IO ()
shrinkMBA (MutableByteArray# RealWorld -> MutableByteArray RealWorld
forall s. MutableByteArray# s -> MutableByteArray s
MutableByteArray MutableByteArray# RealWorld
mba#)
MArray (PrimState m) S ix e -> m (MArray (PrimState m) S ix e)
forall (f :: * -> *) a. Applicative f => a -> f a
pure (MArray (PrimState m) S ix e -> m (MArray (PrimState m) S ix e))
-> MArray (PrimState m) S ix e -> m (MArray (PrimState m) S ix e)
forall a b. (a -> b) -> a -> b
$ Sz ix -> MVector (PrimState m) e -> MArray (PrimState m) S ix e
forall s ix e. Sz ix -> MVector s e -> MArray s S ix e
MSArray Sz ix
sz MVector (PrimState m) e
mv
ForeignPtrContents
_ -> MArray (PrimState m) S ix e
-> Sz ix -> m (MArray (PrimState m) S ix e)
forall r ix e (m :: * -> *).
(Mutable r ix e, PrimMonad m) =>
MArray (PrimState m) r ix e
-> Sz ix -> m (MArray (PrimState m) r ix e)
unsafeDefaultLinearShrink MArray (PrimState m) S ix e
marr Sz ix
sz
{-# INLINE unsafeLinearShrink #-}
unsafeLinearGrow :: MArray (PrimState m) S ix e
-> Sz ix -> m (MArray (PrimState m) S ix e)
unsafeLinearGrow (MSArray oldSz mv) Sz ix
sz =
Sz ix -> MVector (PrimState m) e -> MArray (PrimState m) S ix e
forall s ix e. Sz ix -> MVector s e -> MArray s S ix e
MSArray Sz ix
sz (MVector (PrimState m) e -> MArray (PrimState m) S ix e)
-> m (MVector (PrimState m) e) -> m (MArray (PrimState m) S 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, Storable a) =>
MVector (PrimState m) a -> Int -> m (MVector (PrimState m) a)
MVS.unsafeGrow MVector (PrimState m) e
mv (Sz ix -> Int
forall ix. Index ix => Sz ix -> Int
totalElem Sz ix
sz Int -> Int -> Int
forall a. Num a => a -> a -> a
- Sz ix -> Int
forall ix. Index ix => Sz ix -> Int
totalElem Sz ix
oldSz)
{-# INLINE unsafeLinearGrow #-}
instance (Index ix, Storable e) => Load S ix e where
type R S = M
size :: Array S ix e -> Sz ix
size = Array S ix e -> Sz ix
forall ix e. Array S ix e -> Sz ix
sSize
{-# INLINE size #-}
getComp :: Array S ix e -> Comp
getComp = Array S ix e -> Comp
forall ix e. Array S ix e -> Comp
sComp
{-# INLINE getComp #-}
loadArrayM :: Scheduler m () -> Array S ix e -> (Int -> e -> m ()) -> m ()
loadArrayM !Scheduler m ()
scheduler !Array S 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 S ix e -> Int
forall r ix e. Load r ix e => Array r ix e -> Int
elemsCount Array S ix e
arr) (Array S ix e -> Int -> e
forall r ix e. Source r ix e => Array r ix e -> Int -> e
unsafeLinearIndex Array S ix e
arr)
{-# INLINE loadArrayM #-}
instance (Index ix, Storable e) => StrideLoad S ix e
instance (Index ix, Storable e) => Stream S ix e where
toStream :: Array S ix e -> Steps Id e
toStream = Array S 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 S ix e -> Steps Id (ix, e)
toStreamIx = Array S 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 (Storable e, Num e) => Numeric S e where
unsafeDotProduct :: Array S ix e -> Array S ix e -> e
unsafeDotProduct Array S ix e
a1 Array S ix e
a2 = e -> Int -> e
go e
0 Int
0
where
!len :: Int
len = Sz ix -> Int
forall ix. Index ix => Sz ix -> Int
totalElem (Array S ix e -> Sz ix
forall r ix e. Load r ix e => Array r ix e -> Sz ix
size Array S ix e
a1)
go :: e -> Int -> e
go !e
acc Int
i
| Int
i Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
len = e -> Int -> e
go (e
acc e -> e -> e
forall a. Num a => a -> a -> a
+ Array S ix e -> Int -> e
forall r ix e. Source r ix e => Array r ix e -> Int -> e
unsafeLinearIndex Array S ix e
a1 Int
i e -> e -> e
forall a. Num a => a -> a -> a
* Array S ix e -> Int -> e
forall r ix e. Source r ix e => Array r ix e -> Int -> e
unsafeLinearIndex Array S ix e
a2 Int
i) (Int
i Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1)
| Bool
otherwise = e
acc
{-# INLINE unsafeDotProduct #-}
powerSumArray :: Array S ix e -> Int -> e
powerSumArray Array S ix e
arr Int
p = e -> Int -> e
go e
0 Int
0
where
!len :: Int
len = Sz ix -> Int
forall ix. Index ix => Sz ix -> Int
totalElem (Array S ix e -> Sz ix
forall r ix e. Load r ix e => Array r ix e -> Sz ix
size Array S ix e
arr)
go :: e -> Int -> e
go !e
acc Int
i
| Int
i Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
len = e -> Int -> e
go (e
acc e -> e -> e
forall a. Num a => a -> a -> a
+ Array S ix e -> Int -> e
forall r ix e. Source r ix e => Array r ix e -> Int -> e
unsafeLinearIndex Array S ix e
arr Int
i e -> Int -> e
forall a b. (Num a, Integral b) => a -> b -> a
^ Int
p) (Int
i Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1)
| Bool
otherwise = e
acc
{-# INLINE powerSumArray #-}
foldArray :: (e -> e -> e) -> e -> Array S ix e -> e
foldArray e -> e -> e
f !e
initAcc Array S ix e
arr = e -> Int -> e
go e
initAcc Int
0
where
!len :: Int
len = Sz ix -> Int
forall ix. Index ix => Sz ix -> Int
totalElem (Array S ix e -> Sz ix
forall r ix e. Load r ix e => Array r ix e -> Sz ix
size Array S ix e
arr)
go :: e -> Int -> e
go !e
acc Int
i
| Int
i Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
len = e -> Int -> e
go (e -> e -> e
f e
acc (Array S ix e -> Int -> e
forall r ix e. Source r ix e => Array r ix e -> Int -> e
unsafeLinearIndex Array S ix e
arr Int
i)) (Int
i Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1)
| Bool
otherwise = e
acc
{-# INLINE foldArray #-}
unsafeLiftArray :: (e -> e) -> Array S ix e -> Array S ix e
unsafeLiftArray e -> e
f Array S ix e
arr = Comp -> Sz ix -> (Int -> e) -> Array S ix e
forall r ix e.
Construct r ix e =>
Comp -> Sz ix -> (Int -> e) -> Array r ix e
makeArrayLinear (Array S ix e -> Comp
forall r ix e. Load r ix e => Array r ix e -> Comp
getComp Array S ix e
arr) (Array S ix e -> Sz ix
forall r ix e. Load r ix e => Array r ix e -> Sz ix
size Array S ix e
arr) (e -> e
f (e -> e) -> (Int -> e) -> Int -> e
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Array S ix e -> Int -> e
forall r ix e. Source r ix e => Array r ix e -> Int -> e
unsafeLinearIndex Array S ix e
arr)
{-# INLINE unsafeLiftArray #-}
unsafeLiftArray2 :: (e -> e -> e) -> Array S ix e -> Array S ix e -> Array S ix e
unsafeLiftArray2 e -> e -> e
f Array S ix e
a1 Array S ix e
a2 =
Comp -> Sz ix -> (Int -> e) -> Array S ix e
forall r ix e.
Construct r ix e =>
Comp -> Sz ix -> (Int -> e) -> Array r ix e
makeArrayLinear
(Array S ix e -> Comp
forall r ix e. Load r ix e => Array r ix e -> Comp
getComp Array S ix e
a1 Comp -> Comp -> Comp
forall a. Semigroup a => a -> a -> a
<> Array S ix e -> Comp
forall r ix e. Load r ix e => Array r ix e -> Comp
getComp Array S ix e
a2)
(ix -> Sz ix
forall ix. ix -> Sz ix
SafeSz ((Int -> Int -> Int) -> ix -> ix -> ix
forall ix. Index ix => (Int -> Int -> Int) -> ix -> ix -> ix
liftIndex2 Int -> Int -> Int
forall a. Ord a => a -> a -> a
min (Sz ix -> ix
forall ix. Sz ix -> ix
unSz (Array S ix e -> Sz ix
forall r ix e. Load r ix e => Array r ix e -> Sz ix
size Array S ix e
a1)) (Sz ix -> ix
forall ix. Sz ix -> ix
unSz (Array S ix e -> Sz ix
forall r ix e. Load r ix e => Array r ix e -> Sz ix
size Array S ix e
a2)))) ((Int -> e) -> Array S ix e) -> (Int -> e) -> Array S ix e
forall a b. (a -> b) -> a -> b
$ \ !Int
i ->
e -> e -> e
f (Array S ix e -> Int -> e
forall r ix e. Source r ix e => Array r ix e -> Int -> e
unsafeLinearIndex Array S ix e
a1 Int
i) (Array S ix e -> Int -> e
forall r ix e. Source r ix e => Array r ix e -> Int -> e
unsafeLinearIndex Array S ix e
a2 Int
i)
{-# INLINE unsafeLiftArray2 #-}
instance (Storable e, Floating e) => NumericFloat S e
instance ( Storable e
, IsList (Array L ix e)
, Nested LN ix e
, Nested L ix e
, Ragged L ix e
) =>
IsList (Array S ix e) where
type Item (Array S ix e) = Item (Array L ix e)
fromList :: [Item (Array S ix e)] -> Array S ix e
fromList = Comp -> [ListItem ix e] -> Array S 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 S ix e -> [Item (Array S 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 S ix e -> Array L ix e)
-> Array S ix e
-> [ListItem ix e]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Array S 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 #-}
unsafeWithPtr :: (MonadUnliftIO m, Storable a) => Array S ix a -> (Ptr a -> m b) -> m b
unsafeWithPtr :: Array S ix a -> (Ptr a -> m b) -> m b
unsafeWithPtr Array S ix a
arr Ptr a -> m b
f = ((forall a. m a -> IO a) -> IO b) -> m b
forall (m :: * -> *) b.
MonadUnliftIO m =>
((forall a. m a -> IO a) -> IO b) -> m b
withRunInIO (((forall a. m a -> IO a) -> IO b) -> m b)
-> ((forall a. m a -> IO a) -> IO b) -> m b
forall a b. (a -> b) -> a -> b
$ \forall a. m a -> IO a
run -> Vector a -> (Ptr a -> IO b) -> IO b
forall a b. Storable a => Vector a -> (Ptr a -> IO b) -> IO b
VS.unsafeWith (Array S ix a -> Vector a
forall ix e. Array S ix e -> Vector e
sData Array S ix a
arr) (m b -> IO b
forall a. m a -> IO a
run (m b -> IO b) -> (Ptr a -> m b) -> Ptr a -> IO b
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Ptr a -> m b
f)
{-# INLINE unsafeWithPtr #-}
withPtr :: (MonadUnliftIO m, Storable a) => MArray RealWorld S ix a -> (Ptr a -> m b) -> m b
withPtr :: MArray RealWorld S ix a -> (Ptr a -> m b) -> m b
withPtr (MSArray _ mv) Ptr a -> m b
f = ((forall a. m a -> IO a) -> IO b) -> m b
forall (m :: * -> *) b.
MonadUnliftIO m =>
((forall a. m a -> IO a) -> IO b) -> m b
withRunInIO (((forall a. m a -> IO a) -> IO b) -> m b)
-> ((forall a. m a -> IO a) -> IO b) -> m b
forall a b. (a -> b) -> a -> b
$ \forall a. m a -> IO a
run -> IOVector a -> (Ptr a -> IO b) -> IO b
forall a b. Storable a => IOVector a -> (Ptr a -> IO b) -> IO b
MVS.unsafeWith IOVector a
mv (m b -> IO b
forall a. m a -> IO a
run (m b -> IO b) -> (Ptr a -> m b) -> Ptr a -> IO b
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Ptr a -> m b
f)
{-# INLINE withPtr #-}
toStorableVector :: Array S ix e -> VS.Vector e
toStorableVector :: Array S ix e -> Vector e
toStorableVector = Array S ix e -> Vector e
forall ix e. Array S ix e -> Vector e
sData
{-# INLINE toStorableVector #-}
toStorableMVector :: MArray s S ix e -> VS.MVector s e
toStorableMVector :: MArray s S ix e -> MVector s e
toStorableMVector (MSArray _ mv) = MVector s e
mv
{-# INLINE toStorableMVector #-}
fromStorableVector :: Storable e => Comp -> VS.Vector e -> Array S Ix1 e
fromStorableVector :: Comp -> Vector e -> Array S Int e
fromStorableVector Comp
comp Vector e
v = SArray :: forall ix e. Comp -> Sz ix -> Vector e -> Array S ix e
SArray {sComp :: Comp
sComp = Comp
comp, sSize :: Sz1
sSize = Int -> Sz1
forall ix. ix -> Sz ix
SafeSz (Vector e -> Int
forall a. Storable a => Vector a -> Int
VS.length Vector e
v), sData :: Vector e
sData = Vector e
v}
{-# INLINE fromStorableVector #-}
fromStorableMVector :: MVS.MVector s e -> MArray s S Ix1 e
fromStorableMVector :: MVector s e -> MArray s S Int e
fromStorableMVector mv :: MVector s e
mv@(MVS.MVector Int
len ForeignPtr e
_) = Sz1 -> MVector s e -> MArray s S Int e
forall s ix e. Sz ix -> MVector s e -> MArray s S ix e
MSArray (Int -> Sz1
forall ix. ix -> Sz ix
SafeSz Int
len) MVector s e
mv
{-# INLINE fromStorableMVector #-}
unsafeArrayToForeignPtr :: Storable e => Array S ix e -> (ForeignPtr e, Int)
unsafeArrayToForeignPtr :: Array S ix e -> (ForeignPtr e, Int)
unsafeArrayToForeignPtr = Vector e -> (ForeignPtr e, Int)
forall a. Storable a => Vector a -> (ForeignPtr a, Int)
VS.unsafeToForeignPtr0 (Vector e -> (ForeignPtr e, Int))
-> (Array S ix e -> Vector e)
-> Array S ix e
-> (ForeignPtr e, Int)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Array S ix e -> Vector e
forall ix e. Array S ix e -> Vector e
toStorableVector
{-# INLINE unsafeArrayToForeignPtr #-}
unsafeMArrayToForeignPtr :: Storable e => MArray s S ix e -> (ForeignPtr e, Int)
unsafeMArrayToForeignPtr :: MArray s S ix e -> (ForeignPtr e, Int)
unsafeMArrayToForeignPtr = MVector s e -> (ForeignPtr e, Int)
forall a s. Storable a => MVector s a -> (ForeignPtr a, Int)
MVS.unsafeToForeignPtr0 (MVector s e -> (ForeignPtr e, Int))
-> (MArray s S ix e -> MVector s e)
-> MArray s S ix e
-> (ForeignPtr e, Int)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. MArray s S ix e -> MVector s e
forall s ix e. MArray s S ix e -> MVector s e
toStorableMVector
{-# INLINE unsafeMArrayToForeignPtr #-}
unsafeArrayFromForeignPtr0 :: Storable e => Comp -> ForeignPtr e -> Sz1 -> Array S Ix1 e
unsafeArrayFromForeignPtr0 :: Comp -> ForeignPtr e -> Sz1 -> Array S Int e
unsafeArrayFromForeignPtr0 Comp
comp ForeignPtr e
ptr Sz1
sz =
SArray :: forall ix e. Comp -> Sz ix -> Vector e -> Array S ix e
SArray {sComp :: Comp
sComp = Comp
comp, sSize :: Sz1
sSize = Sz1
sz, sData :: Vector e
sData = ForeignPtr e -> Int -> Vector e
forall a. Storable a => ForeignPtr a -> Int -> Vector a
VS.unsafeFromForeignPtr0 ForeignPtr e
ptr (Sz1 -> Int
forall ix. Sz ix -> ix
unSz Sz1
sz)}
{-# INLINE unsafeArrayFromForeignPtr0 #-}
unsafeArrayFromForeignPtr :: Storable e => Comp -> ForeignPtr e -> Int -> Sz1 -> Array S Ix1 e
unsafeArrayFromForeignPtr :: Comp -> ForeignPtr e -> Int -> Sz1 -> Array S Int e
unsafeArrayFromForeignPtr Comp
comp ForeignPtr e
ptr Int
offset Sz1
sz =
SArray :: forall ix e. Comp -> Sz ix -> Vector e -> Array S ix e
SArray {sComp :: Comp
sComp = Comp
comp, sSize :: Sz1
sSize = Sz1
sz, sData :: Vector e
sData = ForeignPtr e -> Int -> Int -> Vector e
forall a. Storable a => ForeignPtr a -> Int -> Int -> Vector a
VS.unsafeFromForeignPtr ForeignPtr e
ptr Int
offset (Sz1 -> Int
forall ix. Sz ix -> ix
unSz Sz1
sz)}
{-# INLINE unsafeArrayFromForeignPtr #-}
unsafeMArrayFromForeignPtr0 :: Storable e => ForeignPtr e -> Sz1 -> MArray s S Ix1 e
unsafeMArrayFromForeignPtr0 :: ForeignPtr e -> Sz1 -> MArray s S Int e
unsafeMArrayFromForeignPtr0 ForeignPtr e
fp Sz1
sz =
Sz1 -> MVector s e -> MArray s S Int e
forall s ix e. Sz ix -> MVector s e -> MArray s S ix e
MSArray Sz1
sz (ForeignPtr e -> Int -> MVector s e
forall a s. Storable a => ForeignPtr a -> Int -> MVector s a
MVS.unsafeFromForeignPtr0 ForeignPtr e
fp (Sz1 -> Int
forall ix. Sz ix -> ix
unSz Sz1
sz))
{-# INLINE unsafeMArrayFromForeignPtr0 #-}
unsafeMArrayFromForeignPtr :: Storable e => ForeignPtr e -> Int -> Sz1 -> MArray s S Ix1 e
unsafeMArrayFromForeignPtr :: ForeignPtr e -> Int -> Sz1 -> MArray s S Int e
unsafeMArrayFromForeignPtr ForeignPtr e
fp Int
offset Sz1
sz =
Sz1 -> MVector s e -> MArray s S Int e
forall s ix e. Sz ix -> MVector s e -> MArray s S ix e
MSArray Sz1
sz (ForeignPtr e -> Int -> Int -> MVector s e
forall a s. Storable a => ForeignPtr a -> Int -> Int -> MVector s a
MVS.unsafeFromForeignPtr ForeignPtr e
fp Int
offset (Sz1 -> Int
forall ix. Sz ix -> ix
unSz Sz1
sz))
{-# INLINE unsafeMArrayFromForeignPtr #-}