{-# LANGUAGE BangPatterns #-}
{-# LANGUAGE CPP #-}
{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE FlexibleInstances #-}
{-# LANGUAGE MagicHash #-}
{-# LANGUAGE MultiParamTypeClasses #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE TypeFamilies #-}
{-# LANGUAGE UndecidableInstances #-}
-- |
-- Module      : Data.Massiv.Array.Manifest.Storable
-- Copyright   : (c) Alexey Kuleshevich 2018-2020
-- License     : BSD3
-- Maintainer  : Alexey Kuleshevich <lehins@yandex.ru>
-- Stability   : experimental
-- Portability : non-portable
--
module Data.Massiv.Array.Manifest.Storable
  ( S (..)
  , Array(..)
  , Storable
  , toStorableVector
  , toStorableMVector
  , fromStorableVector
  , fromStorableMVector
  , withPtr
  , unsafeWithPtr
  , unsafeMallocMArray
  , 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 (compareArrays, eqArrays)
import Data.Massiv.Array.Manifest.Internal
import Data.Massiv.Array.Manifest.List as A
import Data.Massiv.Array.Manifest.Primitive (shrinkMutableByteArray)
import Data.Massiv.Array.Mutable
import Data.Massiv.Core.Common
import Data.Massiv.Core.List
import Data.Massiv.Core.Operations
import Data.Massiv.Vector.Stream as S (isteps, steps)
import Data.Primitive.ByteArray (MutableByteArray(..))
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.ForeignPtr (withForeignPtr, newForeignPtr)
import Foreign.Marshal.Array (advancePtr, copyArray)
import Foreign.Marshal.Alloc
import Foreign.Ptr
import Foreign.Storable
import GHC.Exts as GHC (IsList(..))
import GHC.ForeignPtr (ForeignPtr(..), ForeignPtrContents(..))
import Prelude hiding (mapM)
import System.IO.Unsafe (unsafePerformIO)
import Control.Exception

#include "massiv.h"

-- | Representation for `Storable` elements
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 #-}

-- | A pointer to the beginning of the storable array. It is unsafe since, if mutated, it can break
-- referential transparency.
--
-- @since 0.1.3
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 #-}


-- | A pointer to the beginning of the mutable array.
--
-- @since 0.1.3
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 #-}


-- | /O(1)/ - Unwrap storable array and pull out the underlying storable vector.
--
-- @since 0.2.1
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 #-}


-- | /O(1)/ - Unwrap storable mutable array and pull out the underlying storable mutable vector.
--
-- @since 0.2.1
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 #-}

-- | /O(1)/ - Cast a storable vector to a storable array.
--
-- @since 0.5.0
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 #-}

-- | /O(1)/ - Cast a mutable storable vector to a mutable storable array.
--
-- @since 0.5.0
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 #-}


-- | /O(1)/ - Yield the underlying `ForeignPtr` together with its length.
--
-- @since 0.3.0
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 #-}

-- | /O(1)/ - Yield the underlying `ForeignPtr` together with its length.
--
-- @since 0.3.0
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 #-}

-- | /O(1)/ - Wrap a `ForeignPtr` and it's size into a pure storable array.
--
-- @since 0.3.0
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 #-}

-- | /O(1)/ - Wrap a `ForeignPtr`, an offset and it's size into a pure storable array.
--
-- @since 0.3.0
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 #-}


-- | /O(1)/ - Wrap a `ForeignPtr` and it's size into a mutable storable array. It is still safe to
-- modify the pointer, unless the array gets frozen prior to modification.
--
-- @since 0.3.0
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 #-}


-- | /O(1)/ - Wrap a `ForeignPtr`, an offset and it's size into a mutable storable array. It is
-- still safe to modify the pointer, unless the array gets frozen prior to modification.
--
-- @since 0.3.0
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 #-}


-- | Allocate memory using @malloc@ on C heap, instead of on Haskell heap. Memory is left
-- uninitialized
--
-- @since 0.5.9
unsafeMallocMArray ::
     forall ix e. (Index ix, Storable e)
  => Sz ix
  -> IO (MArray RealWorld S ix e)
unsafeMallocMArray :: Sz ix -> IO (MArray RealWorld S ix e)
unsafeMallocMArray Sz ix
sz = do
  let n :: Int
n = Sz ix -> Int
forall ix. Index ix => Sz ix -> Int
totalElem Sz ix
sz
  ForeignPtr e
foreignPtr <- IO (ForeignPtr e) -> IO (ForeignPtr e)
forall a. IO a -> IO a
mask_ (IO (ForeignPtr e) -> IO (ForeignPtr e))
-> IO (ForeignPtr e) -> IO (ForeignPtr e)
forall a b. (a -> b) -> a -> b
$ do
    Ptr e
ptr <- Int -> IO (Ptr e)
forall a. Int -> IO (Ptr a)
mallocBytes (e -> Int
forall a. Storable a => a -> Int
sizeOf (e
forall a. HasCallStack => a
undefined :: e) Int -> Int -> Int
forall a. Num a => a -> a -> a
* Int
n)
    FinalizerPtr e -> Ptr e -> IO (ForeignPtr e)
forall a. FinalizerPtr a -> Ptr a -> IO (ForeignPtr a)
newForeignPtr FinalizerPtr e
forall a. FinalizerPtr a
finalizerFree Ptr e
ptr
  MArray RealWorld S ix e -> IO (MArray RealWorld S ix e)
forall (f :: * -> *) a. Applicative f => a -> f a
pure (MArray RealWorld S ix e -> IO (MArray RealWorld S ix e))
-> MArray RealWorld S ix e -> IO (MArray RealWorld S ix e)
forall a b. (a -> b) -> a -> b
$ Sz ix -> MVector RealWorld e -> MArray RealWorld S ix e
forall s ix e. Sz ix -> MVector s e -> MArray s S ix e
MSArray Sz ix
sz (ForeignPtr e -> Int -> MVector RealWorld e
forall a s. Storable a => ForeignPtr a -> Int -> MVector s a
MVS.unsafeFromForeignPtr0 ForeignPtr e
foreignPtr Int
n)
{-# INLINE unsafeMallocMArray #-}