{-# 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-2022
-- License     : BSD3
-- Maintainer  : Alexey Kuleshevich <lehins@yandex.ru>
-- Stability   : experimental
-- Portability : non-portable
module Data.Massiv.Array.Manifest.Storable (
  S (..),
  Array (..),
  MArray (..),
  Storable,
  toStorableVector,
  toStorableMVector,
  fromStorableVector,
  fromStorableMVector,
  withPtr,
  unsafeWithPtr,
  unsafeMallocMArray,
  unsafeArrayToForeignPtr,
  unsafeMArrayToForeignPtr,
  unsafeArrayFromForeignPtr,
  unsafeArrayFromForeignPtr0,
  unsafeMArrayFromForeignPtr,
  unsafeMArrayFromForeignPtr0,
) where

import Control.DeepSeq (NFData (..), deepseq)
import Control.Exception
import Control.Monad
import Control.Monad.IO.Unlift
import Control.Monad.Primitive
import Data.Massiv.Array.Delayed.Pull (D, compareArrays, eqArrays)
import Data.Massiv.Array.Manifest.Internal
import Data.Massiv.Array.Manifest.List as A
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
import Data.Primitive.Ptr (setPtr)
import Data.Word
import Foreign.ForeignPtr
import Foreign.Marshal.Alloc
import Foreign.Marshal.Array (advancePtr, copyArray)
import Foreign.Ptr
import Foreign.Storable
import GHC.Exts as GHC
import GHC.ForeignPtr
import System.IO.Unsafe (unsafePerformIO)
import Unsafe.Coerce
import Prelude hiding (mapM)

import qualified Data.Vector.Generic.Mutable as MVG
import qualified Data.Vector.Storable as VS
import qualified Data.Vector.Storable.Mutable as MVS

-- | Representation for `Storable` elements
data S = S deriving (Int -> S -> ShowS
[S] -> ShowS
S -> String
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
  { forall ix e. Array S ix e -> Comp
sComp :: !Comp
  , forall ix e. Array S ix e -> Sz ix
sSize :: !(Sz ix)
  , forall ix e. Array S ix e -> ForeignPtr e
sData :: {-# UNPACK #-} !(ForeignPtr e)
  }

data instance MArray s S ix e = MSArray !(Sz ix) {-# UNPACK #-} !(ForeignPtr e)

instance (Ragged L ix e, Show e, Storable e) => Show (Array S ix e) where
  showsPrec :: Int -> Array S ix e -> ShowS
showsPrec = 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 forall a. a -> a
id
  showList :: [Array S ix e] -> ShowS
showList = forall arr. Show arr => [arr] -> ShowS
showArrayList

instance NFData ix => NFData (Array S ix e) where
  rnf :: Array S ix e -> ()
rnf (SArray Comp
c Sz ix
sz ForeignPtr e
_v) = Comp
c forall a b. NFData a => a -> b -> b
`deepseq` Sz ix
sz forall a b. NFData a => a -> b -> b
`deepseq` ()
  {-# INLINE rnf #-}

instance NFData ix => NFData (MArray s S ix e) where
  rnf :: MArray s S ix e -> ()
rnf (MSArray Sz ix
sz ForeignPtr e
_mv) = Sz ix
sz 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
(==) = 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 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 = 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 forall a. Ord a => a -> a -> Ordering
compare
  {-# INLINE compare #-}

instance Strategy S where
  getComp :: forall ix e. Array S ix e -> Comp
getComp = forall ix e. Array S ix e -> Comp
sComp
  {-# INLINE getComp #-}
  setComp :: forall ix e. Comp -> Array S ix e -> Array S ix e
setComp Comp
c Array S ix e
arr = Array S ix e
arr{sComp :: Comp
sComp = Comp
c}
  {-# INLINE setComp #-}
  repr :: S
repr = S
S

plusFp :: ForeignPtr a -> Int -> ForeignPtr b
plusFp :: forall a b. ForeignPtr a -> Int -> ForeignPtr b
plusFp (ForeignPtr Addr#
addr ForeignPtrContents
c) (I# Int#
d) = forall a. Addr# -> ForeignPtrContents -> ForeignPtr a
ForeignPtr (Addr# -> Int# -> Addr#
plusAddr# Addr#
addr Int#
d) ForeignPtrContents
c

advanceForeignPtr :: forall e. Storable e => ForeignPtr e -> Int -> ForeignPtr e
advanceForeignPtr :: forall e. Storable e => ForeignPtr e -> Int -> ForeignPtr e
advanceForeignPtr ForeignPtr e
fp Int
i = forall a b. ForeignPtr a -> Int -> ForeignPtr b
plusFp ForeignPtr e
fp (Int
i forall a. Num a => a -> a -> a
* forall a. Storable a => a -> Int
sizeOf (forall a. HasCallStack => a
undefined :: e))
{-# INLINE advanceForeignPtr #-}

indexForeignPtr :: Storable e => ForeignPtr e -> Int -> e
indexForeignPtr :: forall e. Storable e => ForeignPtr e -> Int -> e
indexForeignPtr ForeignPtr e
fp Int
i = forall a. IO a -> a
unsafeInlineIO forall a b. (a -> b) -> a -> b
$ forall a b. ForeignPtr a -> (Ptr a -> IO b) -> IO b
unsafeWithForeignPtr ForeignPtr e
fp forall a b. (a -> b) -> a -> b
$ \Ptr e
p -> forall a. Storable a => Ptr a -> Int -> IO a
peekElemOff Ptr e
p Int
i
{-# INLINE indexForeignPtr #-}

instance Storable e => Source S e where
  unsafeLinearIndex :: forall ix. Index ix => Array S ix e -> Int -> e
unsafeLinearIndex (SArray Comp
_ Sz ix
_sz ForeignPtr e
fp) =
    forall a ix e.
String -> (a -> Sz ix) -> (a -> ix -> e) -> a -> ix -> e
indexAssert String
"S.unsafeLinearIndex" (forall a b. a -> b -> a
const (forall ix. Index ix => Sz ix -> Sz1
toLinearSz Sz ix
_sz)) forall e. Storable e => ForeignPtr e -> Int -> e
indexForeignPtr ForeignPtr e
fp
  {-# INLINE unsafeLinearIndex #-}

  unsafeOuterSlice :: forall ix.
(Index ix, Index (Lower ix)) =>
Array S ix e -> Sz (Lower ix) -> Int -> Array S (Lower ix) e
unsafeOuterSlice (SArray Comp
c Sz ix
_ ForeignPtr e
fp) Sz (Lower ix)
szL Int
i =
    let k :: Int
k = forall ix. Index ix => Sz ix -> Int
totalElem Sz (Lower ix)
szL
     in forall ix e. Comp -> Sz ix -> ForeignPtr e -> Array S ix e
SArray Comp
c Sz (Lower ix)
szL forall a b. (a -> b) -> a -> b
$ forall e. Storable e => ForeignPtr e -> Int -> ForeignPtr e
advanceForeignPtr ForeignPtr e
fp (Int
i forall a. Num a => a -> a -> a
* Int
k)
  {-# INLINE unsafeOuterSlice #-}

  unsafeLinearSlice :: forall ix. Index ix => Int -> Sz1 -> Array S ix e -> Array S Int e
unsafeLinearSlice Int
i Sz1
k (SArray Comp
c Sz ix
_ ForeignPtr e
fp) =
    forall ix e. Comp -> Sz ix -> ForeignPtr e -> Array S ix e
SArray Comp
c Sz1
k forall a b. (a -> b) -> a -> b
$ forall e. Storable e => ForeignPtr e -> Int -> ForeignPtr e
advanceForeignPtr ForeignPtr e
fp Int
i
  {-# INLINE unsafeLinearSlice #-}

instance Index ix => Shape S ix where
  maxLinearSize :: forall e. Array S ix e -> Maybe Sz1
maxLinearSize = forall a. a -> Maybe a
Just forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall ix. ix -> Sz ix
SafeSz forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall ix r e. (Index ix, Size r) => Array r ix e -> Int
elemsCount
  {-# INLINE maxLinearSize #-}

instance Size S where
  size :: forall ix e. Array S ix e -> Sz ix
size = forall ix e. Array S ix e -> Sz ix
sSize
  {-# INLINE size #-}
  unsafeResize :: forall ix ix' e.
(Index ix, Index ix') =>
Sz ix' -> Array S ix e -> Array S ix' e
unsafeResize !Sz ix'
sz !Array S ix e
arr = Array S ix e
arr{sSize :: Sz ix'
sSize = Sz ix'
sz}
  {-# INLINE unsafeResize #-}

instance Storable e => Manifest S e where
  unsafeLinearIndexM :: forall ix. Index ix => Array S ix e -> Int -> e
unsafeLinearIndexM (SArray Comp
_ Sz ix
_sz ForeignPtr e
fp) =
    forall a ix e.
String -> (a -> Sz ix) -> (a -> ix -> e) -> a -> ix -> e
indexAssert String
"S.unsafeLinearIndex" (forall a b. a -> b -> a
const (forall ix. Index ix => Sz ix -> Sz1
toLinearSz Sz ix
_sz)) forall e. Storable e => ForeignPtr e -> Int -> e
indexForeignPtr ForeignPtr e
fp
  {-# INLINE unsafeLinearIndexM #-}

  sizeOfMArray :: forall ix s. Index ix => MArray s S ix e -> Sz ix
sizeOfMArray (MSArray Sz ix
sz ForeignPtr e
_) = Sz ix
sz
  {-# INLINE sizeOfMArray #-}

  unsafeResizeMArray :: forall ix' ix s.
(Index ix', Index ix) =>
Sz ix' -> MArray s S ix e -> MArray s S ix' e
unsafeResizeMArray Sz ix'
sz (MSArray Sz ix
_ ForeignPtr e
fp) = forall s ix e. Sz ix -> ForeignPtr e -> MArray s S ix e
MSArray Sz ix'
sz ForeignPtr e
fp
  {-# INLINE unsafeResizeMArray #-}

  unsafeLinearSliceMArray :: forall ix s.
Index ix =>
Int -> Sz1 -> MArray s S ix e -> MVector s S e
unsafeLinearSliceMArray Int
i Sz1
k (MSArray Sz ix
_ ForeignPtr e
fp) = forall s ix e. Sz ix -> ForeignPtr e -> MArray s S ix e
MSArray Sz1
k forall a b. (a -> b) -> a -> b
$ forall e. Storable e => ForeignPtr e -> Int -> ForeignPtr e
advanceForeignPtr ForeignPtr e
fp Int
i
  {-# INLINE unsafeLinearSliceMArray #-}

  unsafeThaw :: forall ix (m :: * -> *).
(Index ix, PrimMonad m) =>
Array S ix e -> m (MArray (PrimState m) S ix e)
unsafeThaw (SArray Comp
_ Sz ix
sz ForeignPtr e
fp) = forall (f :: * -> *) a. Applicative f => a -> f a
pure forall a b. (a -> b) -> a -> b
$ forall s ix e. Sz ix -> ForeignPtr e -> MArray s S ix e
MSArray Sz ix
sz ForeignPtr e
fp
  {-# INLINE unsafeThaw #-}

  unsafeFreeze :: forall ix (m :: * -> *).
(Index ix, PrimMonad m) =>
Comp -> MArray (PrimState m) S ix e -> m (Array S ix e)
unsafeFreeze Comp
comp (MSArray Sz ix
sz ForeignPtr e
v) = forall (f :: * -> *) a. Applicative f => a -> f a
pure forall a b. (a -> b) -> a -> b
$ forall ix e. Comp -> Sz ix -> ForeignPtr e -> Array S ix e
SArray Comp
comp Sz ix
sz ForeignPtr e
v
  {-# INLINE unsafeFreeze #-}

  unsafeNew :: forall ix (m :: * -> *).
(Index ix, PrimMonad m) =>
Sz ix -> m (MArray (PrimState m) S ix e)
unsafeNew Sz ix
sz = do
    let !n :: Int
n = forall ix. Index ix => Sz ix -> Int
totalElem Sz ix
sz
        dummy :: e
dummy = forall a. HasCallStack => a
undefined :: e
        !eSize :: Int
eSize = forall a. Storable a => a -> Int
sizeOf e
dummy
    forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
when (Int
n forall a. Ord a => a -> a -> Bool
> (forall a. Bounded a => a
maxBound :: Int) forall a. Integral a => a -> a -> a
`div` Int
eSize) forall a b. (a -> b) -> a -> b
$ forall a. HasCallStack => String -> a
error forall a b. (a -> b) -> a -> b
$ String
"Array size is too big: " forall a. [a] -> [a] -> [a]
++ forall a. Show a => a -> String
show Sz ix
sz
    forall (m :: * -> *) a. PrimMonad m => IO a -> m a
unsafeIOToPrim forall a b. (a -> b) -> a -> b
$ do
      ForeignPtr e
fp <- forall a. Int -> Int -> IO (ForeignPtr a)
mallocPlainForeignPtrAlignedBytes (Int
n forall a. Num a => a -> a -> a
* forall a. Storable a => a -> Int
sizeOf e
dummy) (forall a. Storable a => a -> Int
alignment e
dummy)
      forall (f :: * -> *) a. Applicative f => a -> f a
pure forall a b. (a -> b) -> a -> b
$ forall s ix e. Sz ix -> ForeignPtr e -> MArray s S ix e
MSArray Sz ix
sz ForeignPtr e
fp
  {-# INLINE unsafeNew #-}

  initialize :: forall ix (m :: * -> *).
(Index ix, PrimMonad m) =>
MArray (PrimState m) S ix e -> m ()
initialize (MSArray Sz ix
sz ForeignPtr e
fp) =
    forall (m :: * -> *) a. PrimMonad m => IO a -> m a
unsafeIOToPrim forall a b. (a -> b) -> a -> b
$
      forall a b. ForeignPtr a -> (Ptr a -> IO b) -> IO b
unsafeWithForeignPtr ForeignPtr e
fp forall a b. (a -> b) -> a -> b
$ \Ptr e
p ->
        forall a (m :: * -> *).
(Prim a, PrimMonad m) =>
Ptr a -> Int -> a -> m ()
setPtr (forall a b. Ptr a -> Ptr b
castPtr Ptr e
p) (forall ix. Index ix => Sz ix -> Int
totalElem Sz ix
sz forall a. Num a => a -> a -> a
* forall a. Storable a => a -> Int
sizeOf (forall a. HasCallStack => a
undefined :: e)) (Word8
0 :: Word8)
  {-# INLINE initialize #-}

  unsafeLinearRead :: forall ix (m :: * -> *).
(Index ix, PrimMonad m) =>
MArray (PrimState m) S ix e -> Int -> m e
unsafeLinearRead (MSArray Sz ix
_sz ForeignPtr e
fp) Int
o =
    forall (m :: * -> *) a. PrimMonad m => IO a -> m a
unsafeIOToPrim forall a b. (a -> b) -> a -> b
$
      forall a ix e.
String -> (a -> Sz ix) -> (a -> ix -> e) -> a -> ix -> e
indexAssert
        String
"S.unsafeLinearRead"
        (forall a b. a -> b -> a
const (forall ix. Index ix => Sz ix -> Sz1
toLinearSz Sz ix
_sz))
        (\ForeignPtr e
_ Int
_ -> forall a b. ForeignPtr a -> (Ptr a -> IO b) -> IO b
unsafeWithForeignPtr ForeignPtr e
fp (forall a. Storable a => Ptr a -> Int -> IO a
`peekElemOff` Int
o))
        ForeignPtr e
fp
        Int
o
  {-# INLINE unsafeLinearRead #-}

  unsafeLinearWrite :: forall ix (m :: * -> *).
(Index ix, PrimMonad m) =>
MArray (PrimState m) S ix e -> Int -> e -> m ()
unsafeLinearWrite (MSArray Sz ix
_sz ForeignPtr e
fp) Int
o e
e =
    forall (m :: * -> *) a. PrimMonad m => IO a -> m a
unsafeIOToPrim forall a b. (a -> b) -> a -> b
$
      forall a ix e.
String -> (a -> Sz ix) -> (a -> ix -> e) -> a -> ix -> e
indexAssert
        String
"S.unsafeLinearWrite"
        (forall a b. a -> b -> a
const (forall ix. Index ix => Sz ix -> Sz1
toLinearSz Sz ix
_sz))
        (\ForeignPtr e
_ Int
_ -> forall a b. ForeignPtr a -> (Ptr a -> IO b) -> IO b
unsafeWithForeignPtr ForeignPtr e
fp (\Ptr e
p -> forall a. Storable a => Ptr a -> Int -> a -> IO ()
pokeElemOff Ptr e
p Int
o e
e))
        ForeignPtr e
fp
        Int
o
  {-# INLINE unsafeLinearWrite #-}

  unsafeLinearSet :: forall ix (m :: * -> *).
(Index ix, PrimMonad m) =>
MArray (PrimState m) S ix e -> Int -> Sz1 -> e -> m ()
unsafeLinearSet (MSArray Sz ix
_ ForeignPtr e
fp) Int
i Sz1
k e
e =
    forall (m :: * -> *) a. PrimMonad m => ST (PrimState m) a -> m a
stToPrim (forall (v :: * -> * -> *) a s. MVector v a => v s a -> a -> ST s ()
MVG.basicSet (forall a s. ForeignPtr a -> Int -> MVector s a
MVS.unsafeFromForeignPtr0 (forall e. Storable e => ForeignPtr e -> Int -> ForeignPtr e
advanceForeignPtr ForeignPtr e
fp Int
i) (forall ix. Sz ix -> ix
unSz Sz1
k)) e
e)
  {-# INLINE unsafeLinearSet #-}

  unsafeLinearCopy :: forall ix' ix (m :: * -> *).
(Index ix', Index ix, PrimMonad m) =>
MArray (PrimState m) S ix' e
-> Int -> MArray (PrimState m) S ix e -> Int -> Sz1 -> m ()
unsafeLinearCopy (MSArray Sz ix'
_ ForeignPtr e
fpFrom) Int
iFrom (MSArray Sz ix
_ ForeignPtr e
fpTo) Int
iTo (Sz Int
k) = do
    forall (m1 :: * -> *) (m2 :: * -> *) a.
(PrimBase m1, PrimMonad m2) =>
m1 a -> m2 a
unsafePrimToPrim forall a b. (a -> b) -> a -> b
$
      forall a b. ForeignPtr a -> (Ptr a -> IO b) -> IO b
withForeignPtr ForeignPtr e
fpFrom forall a b. (a -> b) -> a -> b
$ \Ptr e
ptrFrom ->
        forall a b. ForeignPtr a -> (Ptr a -> IO b) -> IO b
withForeignPtr ForeignPtr e
fpTo forall a b. (a -> b) -> a -> b
$ \Ptr e
ptrTo -> do
          let ptrFrom' :: Ptr e
ptrFrom' = forall a. Storable a => Ptr a -> Int -> Ptr a
advancePtr Ptr e
ptrFrom Int
iFrom
              ptrTo' :: Ptr e
ptrTo' = forall a. Storable a => Ptr a -> Int -> Ptr a
advancePtr Ptr e
ptrTo Int
iTo
          forall a. Storable a => Ptr a -> Ptr a -> Int -> IO ()
copyArray Ptr e
ptrTo' Ptr e
ptrFrom' Int
k
  {-# INLINE unsafeLinearCopy #-}

  unsafeArrayLinearCopy :: forall ix' ix (m :: * -> *).
(Index ix', Index ix, PrimMonad m) =>
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 <- forall r e ix (m :: * -> *).
(Manifest r e, Index ix, PrimMonad m) =>
Array r ix e -> m (MArray (PrimState m) r ix e)
unsafeThaw Array S ix' e
arrFrom
    forall r e ix' ix (m :: * -> *).
(Manifest r e, Index ix', Index ix, 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 :: forall ix (m :: * -> *).
(Index ix, PrimMonad m) =>
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 Sz ix
_ fp :: ForeignPtr e
fp@(ForeignPtr Addr#
_ ForeignPtrContents
fpc)) Sz ix
sz = do
    let shrinkMBA :: MutableByteArray RealWorld -> IO ()
        shrinkMBA :: MutableByteArray RealWorld -> IO ()
shrinkMBA MutableByteArray RealWorld
mba = forall (m :: * -> *).
PrimMonad m =>
MutableByteArray (PrimState m) -> Int -> m ()
shrinkMutableByteArray MutableByteArray RealWorld
mba (forall ix. Index ix => Sz ix -> Int
totalElem Sz ix
sz forall a. Num a => a -> a -> a
* forall a. Storable a => a -> Int
sizeOf (forall a. HasCallStack => a
undefined :: e))
        {-# INLINE shrinkMBA #-}
    case ForeignPtrContents
fpc of
      MallocPtr MutableByteArray# RealWorld
mba# IORef Finalizers
_ -> do
        forall (m1 :: * -> *) (m2 :: * -> *) a.
(PrimBase m1, PrimMonad m2) =>
m1 a -> m2 a
unsafePrimToPrim forall a b. (a -> b) -> a -> b
$ MutableByteArray RealWorld -> IO ()
shrinkMBA (forall s. MutableByteArray# s -> MutableByteArray s
MutableByteArray MutableByteArray# RealWorld
mba#)
        forall (f :: * -> *) a. Applicative f => a -> f a
pure forall a b. (a -> b) -> a -> b
$ forall s ix e. Sz ix -> ForeignPtr e -> MArray s S ix e
MSArray Sz ix
sz ForeignPtr e
fp
      PlainPtr MutableByteArray# RealWorld
mba# -> do
        forall (m1 :: * -> *) (m2 :: * -> *) a.
(PrimBase m1, PrimMonad m2) =>
m1 a -> m2 a
unsafePrimToPrim forall a b. (a -> b) -> a -> b
$ MutableByteArray RealWorld -> IO ()
shrinkMBA (forall s. MutableByteArray# s -> MutableByteArray s
MutableByteArray MutableByteArray# RealWorld
mba#)
        forall (f :: * -> *) a. Applicative f => a -> f a
pure forall a b. (a -> b) -> a -> b
$ forall s ix e. Sz ix -> ForeignPtr e -> MArray s S ix e
MSArray Sz ix
sz ForeignPtr e
fp
      ForeignPtrContents
_ -> forall r e ix (m :: * -> *).
(Manifest r e, Index ix, 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 #-}

instance (Index ix, Storable e) => Load S ix e where
  makeArray :: Comp -> Sz ix -> (ix -> e) -> Array S ix e
makeArray Comp
comp Sz ix
sz ix -> e
f = forall r ix e r'.
(Manifest r e, Load r' ix e) =>
Array r' ix e -> Array r ix e
compute (forall r ix e.
Load r ix e =>
Comp -> Sz ix -> (ix -> e) -> Array r ix e
makeArray Comp
comp Sz ix
sz ix -> e
f :: Array D ix e)
  {-# INLINE makeArray #-}

  makeArrayLinear :: Comp -> Sz ix -> (Int -> e) -> Array S ix e
makeArrayLinear !Comp
comp !Sz ix
sz Int -> e
f = forall a. IO a -> a
unsafePerformIO forall a b. (a -> b) -> a -> b
$ 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 (forall (f :: * -> *) a. Applicative f => a -> f a
pure forall b c a. (b -> c) -> (a -> b) -> a -> c
. Int -> e
f)
  {-# INLINE makeArrayLinear #-}

  replicate :: Comp -> Sz ix -> e -> Array S ix e
replicate Comp
comp !Sz ix
sz !e
e = forall a. (forall s. ST s a) -> a
runST (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 forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= 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_ :: forall s.
Scheduler s () -> Array S ix e -> (Int -> e -> ST s ()) -> ST s ()
iterArrayLinearST_ !Scheduler s ()
scheduler !Array S ix e
arr =
    forall s (m :: * -> *) b.
MonadPrimBase s m =>
Scheduler s () -> Int -> (Int -> b) -> (Int -> b -> m ()) -> m ()
splitLinearlyWith_ Scheduler s ()
scheduler (forall ix r e. (Index ix, Size r) => Array r ix e -> Int
elemsCount Array S ix e
arr) (forall r e ix. (Source r e, Index ix) => Array r ix e -> Int -> e
unsafeLinearIndex Array S ix e
arr)
  {-# INLINE iterArrayLinearST_ #-}

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 = 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 S ix e -> Steps Id (ix, e)
toStreamIx = 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 (Storable e, Num e) => FoldNumeric S e where
  unsafeDotProduct :: forall ix. Index ix => Array S ix e -> Array S ix e -> e
unsafeDotProduct = forall e ix r.
(Num e, Index ix, Source r e) =>
Array r ix e -> Array r ix e -> e
defaultUnsafeDotProduct
  {-# INLINE unsafeDotProduct #-}
  powerSumArray :: forall ix. Index ix => Array S ix e -> Int -> e
powerSumArray = forall ix r e.
(Index ix, Source r e, Num e) =>
Array r ix e -> Int -> e
defaultPowerSumArray
  {-# INLINE powerSumArray #-}
  foldArray :: forall ix. Index ix => (e -> e -> e) -> e -> Array S ix e -> e
foldArray = forall ix r e.
(Index ix, Source r e) =>
(e -> e -> e) -> e -> Array r ix e -> e
defaultFoldArray
  {-# INLINE foldArray #-}

instance (Storable e, Num e) => Numeric S e where
  unsafeLiftArray :: forall ix. Index ix => (e -> e) -> Array S ix e -> Array S ix e
unsafeLiftArray = 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 :: forall ix.
Index ix =>
(e -> e -> e) -> Array S ix e -> Array S ix e -> Array S ix e
unsafeLiftArray2 = 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 #-}

instance (Storable e, Floating e) => NumericFloat S e

instance (Storable e, IsList (Array 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 = 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 S ix e -> [Item (Array S ix e)]
toList = forall l. IsList l => l -> [Item l]
GHC.toList forall b c a. (b -> c) -> (a -> b) -> a -> c
. 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 #-}

-- | 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 => Array S ix e -> (Ptr e -> m b) -> m b
unsafeWithPtr :: forall (m :: * -> *) ix e b.
MonadUnliftIO m =>
Array S ix e -> (Ptr e -> m b) -> m b
unsafeWithPtr Array S ix e
arr Ptr e -> m b
f = forall (m :: * -> *) b.
MonadUnliftIO m =>
((forall a. m a -> IO a) -> IO b) -> m b
withRunInIO forall a b. (a -> b) -> a -> b
$ \forall a. m a -> IO a
run -> forall a b. ForeignPtr a -> (Ptr a -> IO b) -> IO b
unsafeWithForeignPtr (forall ix e. Array S ix e -> ForeignPtr e
sData Array S ix e
arr) (forall a. m a -> IO a
run forall b c a. (b -> c) -> (a -> b) -> a -> c
. Ptr e -> m b
f)
{-# INLINE unsafeWithPtr #-}

-- | A pointer to the beginning of the mutable array.
--
-- @since 0.1.3
withPtr :: MonadUnliftIO m => MArray RealWorld S ix e -> (Ptr e -> m b) -> m b
withPtr :: forall (m :: * -> *) ix e b.
MonadUnliftIO m =>
MArray RealWorld S ix e -> (Ptr e -> m b) -> m b
withPtr (MSArray Sz ix
_ ForeignPtr e
fp) Ptr e -> m b
f = forall (m :: * -> *) b.
MonadUnliftIO m =>
((forall a. m a -> IO a) -> IO b) -> m b
withRunInIO forall a b. (a -> b) -> a -> b
$ \forall a. m a -> IO a
run -> forall a b. ForeignPtr a -> (Ptr a -> IO b) -> IO b
unsafeWithForeignPtr ForeignPtr e
fp (forall a. m a -> IO a
run forall b c a. (b -> c) -> (a -> b) -> a -> c
. Ptr e -> m b
f)
{-# INLINE withPtr #-}

-- | /O(1)/ - Unwrap storable array and pull out the underlying storable vector.
--
-- @since 0.2.1
toStorableVector :: Index ix => Array S ix e -> VS.Vector e
toStorableVector :: forall ix e. Index ix => Array S ix e -> Vector e
toStorableVector Array S ix e
arr =
  -- this hack is needed to workaround the redundant Storable constraint
  -- see haskell/vector#394
  forall a b. a -> b
unsafeCoerce forall a b. (a -> b) -> a -> b
$
    forall a. ForeignPtr a -> Int -> Vector a
VS.unsafeFromForeignPtr0 (forall a b. ForeignPtr a -> ForeignPtr b
castForeignPtr (forall ix e. Array S ix e -> ForeignPtr e
sData Array S ix e
arr) :: ForeignPtr Word) (forall ix. Index ix => Sz ix -> Int
totalElem (forall ix e. Array S ix e -> Sz ix
sSize Array S ix e
arr))
{-# INLINE toStorableVector #-}

-- | /O(1)/ - Unwrap storable mutable array and pull out the underlying storable mutable vector.
--
-- @since 0.2.1
toStorableMVector :: Index ix => MArray s S ix e -> VS.MVector s e
toStorableMVector :: forall ix s e. Index ix => MArray s S ix e -> MVector s e
toStorableMVector (MSArray Sz ix
sz ForeignPtr e
fp) = forall s a. Int -> ForeignPtr a -> MVector s a
MVS.MVector (forall ix. Index ix => Sz ix -> Int
totalElem Sz ix
sz) ForeignPtr e
fp
{-# INLINE toStorableMVector #-}

-- | /O(1)/ - Cast a storable vector to a storable array.
--
-- @since 0.5.0
fromStorableVector :: Comp -> VS.Vector e -> Vector S e
fromStorableVector :: forall e. Comp -> Vector e -> Vector S e
fromStorableVector Comp
comp Vector e
v =
  -- unasfeCoerce hack below is needed to workaround the redundant Storable
  -- constraint see haskell/vector#394
  case forall a. Vector a -> (ForeignPtr a, Int)
VS.unsafeToForeignPtr0 (forall a b. a -> b
unsafeCoerce Vector e
v :: VS.Vector Word) of
    (ForeignPtr Word
fp, Int
k) -> SArray{sComp :: Comp
sComp = Comp
comp, sSize :: Sz1
sSize = forall ix. ix -> Sz ix
SafeSz Int
k, sData :: ForeignPtr e
sData = forall a b. ForeignPtr a -> ForeignPtr b
castForeignPtr ForeignPtr Word
fp}
{-# INLINE fromStorableVector #-}

-- | /O(1)/ - Cast a mutable storable vector to a mutable storable array.
--
-- @since 0.5.0
fromStorableMVector :: MVS.MVector s e -> MVector s S e
fromStorableMVector :: forall s e. MVector s e -> MVector s S e
fromStorableMVector (MVS.MVector Int
n ForeignPtr e
fp) = forall s ix e. Sz ix -> ForeignPtr e -> MArray s S ix e
MSArray (forall ix. ix -> Sz ix
SafeSz Int
n) ForeignPtr e
fp
{-# INLINE fromStorableMVector #-}

-- | /O(1)/ - Yield the underlying `ForeignPtr` together with its length.
--
-- @since 0.3.0
unsafeArrayToForeignPtr :: Index ix => Array S ix e -> (ForeignPtr e, Int)
unsafeArrayToForeignPtr :: forall ix e. Index ix => Array S ix e -> (ForeignPtr e, Int)
unsafeArrayToForeignPtr (SArray Comp
_ Sz ix
sz ForeignPtr e
fp) = (ForeignPtr e
fp, forall ix. Index ix => Sz ix -> Int
totalElem Sz ix
sz)
{-# INLINE unsafeArrayToForeignPtr #-}

-- | /O(1)/ - Yield the underlying `ForeignPtr` together with its length.
--
-- @since 0.3.0
unsafeMArrayToForeignPtr :: Index ix => MArray s S ix e -> (ForeignPtr e, Int)
unsafeMArrayToForeignPtr :: forall ix s e. Index ix => MArray s S ix e -> (ForeignPtr e, Int)
unsafeMArrayToForeignPtr (MSArray Sz ix
sz ForeignPtr e
fp) = (ForeignPtr e
fp, forall ix. Index ix => Sz ix -> Int
totalElem Sz ix
sz)
{-# INLINE unsafeMArrayToForeignPtr #-}

-- | /O(1)/ - Wrap a `ForeignPtr` and it's size into a pure storable array.
--
-- @since 0.3.0
unsafeArrayFromForeignPtr0 :: Comp -> ForeignPtr e -> Sz1 -> Vector S e
unsafeArrayFromForeignPtr0 :: forall e. Comp -> ForeignPtr e -> Sz1 -> Vector S e
unsafeArrayFromForeignPtr0 Comp
comp ForeignPtr e
fp Sz1
sz = SArray{sComp :: Comp
sComp = Comp
comp, sSize :: Sz1
sSize = Sz1
sz, sData :: ForeignPtr e
sData = ForeignPtr e
fp}
{-# 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 :: forall e.
Storable e =>
Comp -> ForeignPtr e -> Int -> Sz1 -> Array S Int e
unsafeArrayFromForeignPtr Comp
comp ForeignPtr e
ptr Int
offset Sz1
sz =
  SArray{sComp :: Comp
sComp = Comp
comp, sSize :: Sz1
sSize = Sz1
sz, sData :: ForeignPtr e
sData = forall e. Storable e => ForeignPtr e -> Int -> ForeignPtr e
advanceForeignPtr ForeignPtr e
ptr Int
offset}
{-# 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 :: ForeignPtr e -> Sz1 -> MArray s S Ix1 e
unsafeMArrayFromForeignPtr0 :: forall e s. ForeignPtr e -> Sz1 -> MArray s S Int e
unsafeMArrayFromForeignPtr0 ForeignPtr e
fp Sz1
sz = forall s ix e. Sz ix -> ForeignPtr e -> MArray s S ix e
MSArray Sz1
sz ForeignPtr e
fp
{-# 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 :: forall e s.
Storable e =>
ForeignPtr e -> Int -> Sz1 -> MArray s S Int e
unsafeMArrayFromForeignPtr ForeignPtr e
fp Int
offset Sz1
sz = forall s ix e. Sz ix -> ForeignPtr e -> MArray s S ix e
MSArray Sz1
sz (forall e. Storable e => ForeignPtr e -> Int -> ForeignPtr e
advanceForeignPtr ForeignPtr e
fp Int
offset)
{-# 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 m
   . (Index ix, Storable e, PrimMonad m)
  => Sz ix
  -> m (MArray (PrimState m) S ix e)
unsafeMallocMArray :: forall ix e (m :: * -> *).
(Index ix, Storable e, PrimMonad m) =>
Sz ix -> m (MArray (PrimState m) S ix e)
unsafeMallocMArray Sz ix
sz = forall (m1 :: * -> *) (m2 :: * -> *) a.
(PrimBase m1, PrimMonad m2) =>
m1 a -> m2 a
unsafePrimToPrim forall a b. (a -> b) -> a -> b
$ do
  let n :: Int
n = forall ix. Index ix => Sz ix -> Int
totalElem Sz ix
sz
  ForeignPtr e
foreignPtr <- forall a. IO a -> IO a
mask_ forall a b. (a -> b) -> a -> b
$ do
    Ptr e
ptr <- forall a. Int -> IO (Ptr a)
mallocBytes (forall a. Storable a => a -> Int
sizeOf (forall a. HasCallStack => a
undefined :: e) forall a. Num a => a -> a -> a
* Int
n)
    forall a. FinalizerPtr a -> Ptr a -> IO (ForeignPtr a)
newForeignPtr forall a. FinalizerPtr a
finalizerFree Ptr e
ptr
  forall (f :: * -> *) a. Applicative f => a -> f a
pure forall a b. (a -> b) -> a -> b
$ forall s ix e. Sz ix -> ForeignPtr e -> MArray s S ix e
MSArray Sz ix
sz ForeignPtr e
foreignPtr
{-# INLINE unsafeMallocMArray #-}

#if !MIN_VERSION_base(4,15,0)
-- | A compatibility wrapper for 'GHC.ForeignPtr.unsafeWithForeignPtr' provided
-- by GHC 9.0.1 and later.
--
-- Only to be used when the continuation is known not to
-- unconditionally diverge lest unsoundness can result.
unsafeWithForeignPtr :: ForeignPtr a -> (Ptr a -> IO b) -> IO b
unsafeWithForeignPtr = withForeignPtr
#endif