{-# LANGUAGE BangPatterns #-}
{-# LANGUAGE TypeFamilyDependencies #-}
{-# LANGUAGE CPP #-}
{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE FlexibleInstances #-}
{-# LANGUAGE FunctionalDependencies #-}
{-# LANGUAGE KindSignatures #-}
{-# LANGUAGE MagicHash #-}
{-# LANGUAGE MultiParamTypeClasses #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE TypeFamilies #-}
{-# LANGUAGE UnboxedTuples #-}
{-# LANGUAGE UnliftedFFITypes #-}
module Std.Data.Array (
Arr(..)
, RealWorld
, Array(..)
, MutableArray(..)
, SmallArray(..)
, SmallMutableArray(..)
, uninitialized
, PrimArray(..)
, MutablePrimArray(..)
, Prim(..)
, newPinnedPrimArray, newAlignedPinnedPrimArray
, copyPrimArrayToPtr, copyMutablePrimArrayToPtr, copyPtrToMutablePrimArray
, primArrayContents, mutablePrimArrayContents, withPrimArrayContents, withMutablePrimArrayContents
, isPrimArrayPinned, isMutablePrimArrayPinned
, UnliftedArray(..)
, MutableUnliftedArray(..)
, PrimUnlifted(..)
, ArrayException(..)
, castArray
, castMutableArray
) where
import Control.Exception (ArrayException (..), throw)
import Control.Monad.Primitive
import Data.Primitive.Array
import Data.Primitive.ByteArray
import Data.Primitive.PrimArray
import Data.Primitive.Ptr (copyPtrToMutablePrimArray)
import Data.Primitive.SmallArray
import Data.Primitive.Types
import Data.Primitive.UnliftedArray
import GHC.Prim
import GHC.Ptr (Ptr (..))
import GHC.ST
import GHC.Types
import Std.Data.PrimArray.Cast
uninitialized :: a
uninitialized = throw (UndefinedElement "Data.Array.uninitialized")
class Arr (marr :: * -> * -> *) (arr :: * -> * ) a | arr -> marr, marr -> arr where
type IArr marr = (ar :: * -> *) | ar -> marr
type MArr arr = (mar :: * -> * -> *) | mar -> arr
newArr :: (PrimMonad m, PrimState m ~ s) => Int -> m (marr s a)
newArrWith :: (PrimMonad m, PrimState m ~ s) => Int -> a -> m (marr s a)
readArr :: (PrimMonad m, PrimState m ~ s) => marr s a -> Int -> m a
writeArr :: (PrimMonad m, PrimState m ~ s) => marr s a -> Int -> a -> m ()
setArr :: (PrimMonad m, PrimState m ~ s) => marr s a -> Int -> Int -> a -> m ()
indexArr :: arr a -> Int -> a
indexArr' :: arr a -> Int -> (# a #)
indexArrM :: (Monad m) => arr a -> Int -> m a
freezeArr :: (PrimMonad m, PrimState m ~ s) => marr s a -> Int -> Int -> m (arr a)
thawArr :: (PrimMonad m, PrimState m ~ s) => arr a -> Int -> Int -> m (marr s a)
unsafeFreezeArr :: (PrimMonad m, PrimState m ~ s) => marr s a -> m (arr a)
unsafeThawArr :: (PrimMonad m, PrimState m ~ s) => arr a -> m (marr s a)
copyArr :: (PrimMonad m, PrimState m ~ s) => marr s a -> Int -> arr a -> Int -> Int -> m ()
copyMutableArr :: (PrimMonad m, PrimState m ~ s) => marr s a -> Int -> marr s a -> Int -> Int -> m ()
moveArr :: (PrimMonad m, PrimState m ~ s) => marr s a -> Int -> marr s a -> Int -> Int -> m ()
cloneArr :: arr a -> Int -> Int -> arr a
cloneMutableArr :: (PrimMonad m, PrimState m ~ s) => marr s a -> Int -> Int -> m (marr s a)
resizeMutableArr :: (PrimMonad m, PrimState m ~ s) => marr s a -> Int -> m (marr s a)
shrinkMutableArr :: (PrimMonad m, PrimState m ~ s) => marr s a -> Int -> m ()
sameMutableArr :: marr s a -> marr s a -> Bool
sizeofArr :: arr a -> Int
sizeofMutableArr :: (PrimMonad m, PrimState m ~ s) => marr s a -> m Int
sameArr :: arr a -> arr a -> Bool
instance Arr MutableArray Array a where
type MArr Array = MutableArray
type IArr MutableArray = Array
newArr n = newArray n uninitialized
{-# INLINE newArr #-}
newArrWith = newArray
{-# INLINE newArrWith #-}
readArr = readArray
{-# INLINE readArr #-}
writeArr = writeArray
{-# INLINE writeArr #-}
setArr marr s l x = go s
where
!sl = s + l
go !i | i >= sl = return ()
| otherwise = writeArray marr i x >> go (i+1)
{-# INLINE setArr #-}
indexArr = indexArray
{-# INLINE indexArr #-}
indexArr' (Array arr#) (I# i#) = indexArray# arr# i#
{-# INLINE indexArr' #-}
indexArrM = indexArrayM
{-# INLINE indexArrM #-}
freezeArr = freezeArray
{-# INLINE freezeArr #-}
thawArr = thawArray
{-# INLINE thawArr #-}
unsafeFreezeArr = unsafeFreezeArray
{-# INLINE unsafeFreezeArr #-}
unsafeThawArr = unsafeThawArray
{-# INLINE unsafeThawArr #-}
copyArr = copyArray
{-# INLINE copyArr #-}
copyMutableArr = copyMutableArray
{-# INLINE copyMutableArr #-}
moveArr marr1 s1 marr2 s2 l
| l <= 0 = return ()
| sameMutableArray marr1 marr2 =
case compare s1 s2 of
LT ->
let !d = s2 - s1
!s2l = s2 + l
go !i | i >= s2l = return ()
| otherwise = do x <- readArray marr2 i
writeArray marr1 (i-d) x
go (i+1)
in go s2
EQ -> return ()
GT ->
let !d = s1 - s2
go !i | i < s2 = return ()
| otherwise = do x <- readArray marr2 i
writeArray marr1 (i+d) x
go (i-1)
in go (s2+l-1)
| otherwise = copyMutableArray marr1 s1 marr2 s2 l
{-# INLINE moveArr #-}
cloneArr = cloneArray
{-# INLINE cloneArr #-}
cloneMutableArr = cloneMutableArray
{-# INLINE cloneMutableArr #-}
resizeMutableArr marr n = do
marr' <- newArray n uninitialized
copyMutableArray marr' 0 marr 0 (sizeofMutableArray marr)
return marr'
{-# INLINE resizeMutableArr #-}
shrinkMutableArr _ _ = return ()
{-# INLINE shrinkMutableArr #-}
sameMutableArr = sameMutableArray
{-# INLINE sameMutableArr #-}
sizeofArr = sizeofArray
{-# INLINE sizeofArr #-}
sizeofMutableArr = return . sizeofMutableArray
{-# INLINE sizeofMutableArr #-}
sameArr (Array arr1#) (Array arr2#) = isTrue# (
sameMutableArray# (unsafeCoerce# arr1#) (unsafeCoerce# arr2#))
{-# INLINE sameArr #-}
instance Arr SmallMutableArray SmallArray a where
type MArr SmallArray = SmallMutableArray
type IArr SmallMutableArray = SmallArray
newArr n = newSmallArray n uninitialized
{-# INLINE newArr #-}
newArrWith = newSmallArray
{-# INLINE newArrWith #-}
readArr = readSmallArray
{-# INLINE readArr #-}
writeArr = writeSmallArray
{-# INLINE writeArr #-}
setArr marr s l x = go s
where
!sl = s + l
go !i | i >= sl = return ()
| otherwise = writeSmallArray marr i x >> go (i+1)
{-# INLINE setArr #-}
indexArr = indexSmallArray
{-# INLINE indexArr #-}
indexArr' (SmallArray arr#) (I# i#) = indexSmallArray# arr# i#
{-# INLINE indexArr' #-}
indexArrM = indexSmallArrayM
{-# INLINE indexArrM #-}
freezeArr = freezeSmallArray
{-# INLINE freezeArr #-}
thawArr = thawSmallArray
{-# INLINE thawArr #-}
unsafeFreezeArr = unsafeFreezeSmallArray
{-# INLINE unsafeFreezeArr #-}
unsafeThawArr = unsafeThawSmallArray
{-# INLINE unsafeThawArr #-}
copyArr = copySmallArray
{-# INLINE copyArr #-}
copyMutableArr = copySmallMutableArray
{-# INLINE copyMutableArr #-}
moveArr marr1 s1 marr2 s2 l
| l <= 0 = return ()
| sameMutableArr marr1 marr2 =
case compare s1 s2 of
LT ->
let !d = s2 - s1
!s2l = s2 + l
go !i | i >= s2l = return ()
| otherwise = do x <- readSmallArray marr2 i
writeSmallArray marr1 (i-d) x
go (i+1)
in go s2
EQ -> return ()
GT ->
let !d = s1 - s2
go !i | i < s2 = return ()
| otherwise = do x <- readSmallArray marr2 i
writeSmallArray marr1 (i+d) x
go (i-1)
in go (s2+l-1)
| otherwise = copySmallMutableArray marr1 s1 marr2 s2 l
{-# INLINE moveArr #-}
cloneArr = cloneSmallArray
{-# INLINE cloneArr #-}
cloneMutableArr = cloneSmallMutableArray
{-# INLINE cloneMutableArr #-}
resizeMutableArr marr n = do
marr' <- newSmallArray n uninitialized
copySmallMutableArray marr' 0 marr 0 (sizeofSmallMutableArray marr)
return marr'
{-# INLINE resizeMutableArr #-}
shrinkMutableArr _ _ = return ()
{-# INLINE shrinkMutableArr #-}
sameMutableArr (SmallMutableArray smarr1#) (SmallMutableArray smarr2#) =
isTrue# (sameSmallMutableArray# smarr1# smarr2#)
{-# INLINE sameMutableArr #-}
sizeofArr = sizeofSmallArray
{-# INLINE sizeofArr #-}
sizeofMutableArr = return . sizeofSmallMutableArray
{-# INLINE sizeofMutableArr #-}
sameArr (SmallArray arr1#) (SmallArray arr2#) = isTrue# (
sameSmallMutableArray# (unsafeCoerce# arr1#) (unsafeCoerce# arr2#))
{-# INLINE sameArr #-}
instance Prim a => Arr MutablePrimArray PrimArray a where
type MArr PrimArray = MutablePrimArray
type IArr MutablePrimArray = PrimArray
newArr = newPrimArray
{-# INLINE newArr #-}
newArrWith n x = do
marr <- newPrimArray n
setPrimArray marr 0 n x
return marr
{-# INLINE newArrWith #-}
readArr = readPrimArray
{-# INLINE readArr #-}
writeArr = writePrimArray
{-# INLINE writeArr #-}
setArr = setPrimArray
{-# INLINE setArr #-}
indexArr = indexPrimArray
{-# INLINE indexArr #-}
indexArr' arr i = (# indexPrimArray arr i #)
{-# INLINE indexArr' #-}
indexArrM arr i = return (indexPrimArray arr i)
{-# INLINE indexArrM #-}
freezeArr marr s l = do
marr' <- newPrimArray l
copyMutablePrimArray marr' 0 marr s l
unsafeFreezePrimArray marr'
{-# INLINE freezeArr #-}
thawArr arr s l = do
marr' <- newPrimArray l
copyPrimArray marr' 0 arr s l
return marr'
{-# INLINE thawArr #-}
unsafeFreezeArr = unsafeFreezePrimArray
{-# INLINE unsafeFreezeArr #-}
unsafeThawArr = unsafeThawPrimArray
{-# INLINE unsafeThawArr #-}
copyArr = copyPrimArray
{-# INLINE copyArr #-}
copyMutableArr = copyMutablePrimArray
{-# INLINE copyMutableArr #-}
moveArr (MutablePrimArray dst) doff (MutablePrimArray src) soff n =
moveByteArray (MutableByteArray dst) (doff*siz) (MutableByteArray src) (soff*siz) (n*siz)
where siz = sizeOf (undefined :: a)
{-# INLINE moveArr #-}
cloneArr arr s l = runST (do
marr <- newPrimArray l
copyPrimArray marr 0 arr s l
unsafeFreezePrimArray marr
)
{-# INLINE cloneArr #-}
cloneMutableArr marr s l = do
marr' <- newPrimArray l
copyMutablePrimArray marr' 0 marr s l
return marr'
{-# INLINE cloneMutableArr #-}
resizeMutableArr = resizeMutablePrimArray
{-# INLINE resizeMutableArr #-}
shrinkMutableArr = shrinkMutablePrimArray
{-# INLINE shrinkMutableArr #-}
sameMutableArr = sameMutablePrimArray
{-# INLINE sameMutableArr #-}
sizeofArr = sizeofPrimArray
{-# INLINE sizeofArr #-}
sizeofMutableArr = getSizeofMutablePrimArray
{-# INLINE sizeofMutableArr #-}
sameArr (PrimArray ba1#) (PrimArray ba2#) =
isTrue# (sameMutableByteArray# (unsafeCoerce# ba1#) (unsafeCoerce# ba2#))
{-# INLINE sameArr #-}
instance PrimUnlifted a => Arr MutableUnliftedArray UnliftedArray a where
type MArr UnliftedArray = MutableUnliftedArray
type IArr MutableUnliftedArray = UnliftedArray
newArr = unsafeNewUnliftedArray
{-# INLINE newArr #-}
newArrWith = newUnliftedArray
{-# INLINE newArrWith #-}
readArr = readUnliftedArray
{-# INLINE readArr #-}
writeArr = writeUnliftedArray
{-# INLINE writeArr #-}
setArr marr s l x = go s
where
!sl = s + l
go !i | i >= sl = return ()
| otherwise = writeUnliftedArray marr i x >> go (i+1)
{-# INLINE setArr #-}
indexArr = indexUnliftedArray
{-# INLINE indexArr #-}
indexArr' arr i = (# indexUnliftedArray arr i #)
{-# INLINE indexArr' #-}
indexArrM = indexUnliftedArrayM
{-# INLINE indexArrM #-}
freezeArr = freezeUnliftedArray
{-# INLINE freezeArr #-}
thawArr = thawUnliftedArray
{-# INLINE thawArr #-}
unsafeFreezeArr = unsafeFreezeUnliftedArray
{-# INLINE unsafeFreezeArr #-}
unsafeThawArr (UnliftedArray arr#) = primitive ( \ s0# ->
let (# s1#, marr# #) = unsafeThawArray# (unsafeCoerce# arr#) s0#
in (# s1#, MutableUnliftedArray (unsafeCoerce# marr#) #)
)
{-# INLINE unsafeThawArr #-}
copyArr = copyUnliftedArray
{-# INLINE copyArr #-}
copyMutableArr = copyMutableUnliftedArray
{-# INLINE copyMutableArr #-}
moveArr marr1 s1 marr2 s2 l
| l <= 0 = return ()
| sameMutableUnliftedArray marr1 marr2 =
case compare s1 s2 of
LT ->
let !d = s2 - s1
!s2l = s2 + l
go !i | i >= s2l = return ()
| otherwise = do x <- readUnliftedArray marr2 i
writeUnliftedArray marr1 (i-d) x
go (i+1)
in go s2
EQ -> return ()
GT ->
let !d = s1 - s2
go !i | i < s2 = return ()
| otherwise = do x <- readUnliftedArray marr2 i
writeUnliftedArray marr1 (i+d) x
go (i-1)
in go (s2+l-1)
| otherwise = copyMutableUnliftedArray marr1 s1 marr2 s2 l
{-# INLINE moveArr #-}
cloneArr = cloneUnliftedArray
{-# INLINE cloneArr #-}
cloneMutableArr = cloneMutableUnliftedArray
{-# INLINE cloneMutableArr #-}
resizeMutableArr marr n = do
marr' <- newUnliftedArray n uninitialized
copyMutableUnliftedArray marr' 0 marr 0 (sizeofMutableUnliftedArray marr)
return marr'
{-# INLINE resizeMutableArr #-}
shrinkMutableArr _ _ = return ()
{-# INLINE shrinkMutableArr #-}
sameMutableArr = sameMutableUnliftedArray
{-# INLINE sameMutableArr #-}
sizeofArr = sizeofUnliftedArray
{-# INLINE sizeofArr #-}
sizeofMutableArr = return . sizeofMutableUnliftedArray
{-# INLINE sizeofMutableArr #-}
sameArr (UnliftedArray arr1#) (UnliftedArray arr2#) = isTrue# (
sameMutableArrayArray# (unsafeCoerce# arr1#) (unsafeCoerce# arr2#))
{-# INLINE sameArr #-}
newPinnedPrimArray :: forall m a. (PrimMonad m, Prim a) => Int -> m (MutablePrimArray (PrimState m) a)
{-# INLINE newPinnedPrimArray #-}
newPinnedPrimArray n = do
(MutableByteArray mba#) <- newPinnedByteArray (n*siz)
return (MutablePrimArray mba#)
where siz = sizeOf (undefined :: a)
newAlignedPinnedPrimArray
:: forall m a. (PrimMonad m, Prim a) => Int -> m (MutablePrimArray (PrimState m) a)
{-# INLINE newAlignedPinnedPrimArray #-}
newAlignedPinnedPrimArray n = do
(MutableByteArray mba#) <- newAlignedPinnedByteArray (n*siz) align
return (MutablePrimArray mba#)
where siz = sizeOf (undefined :: a)
align = alignment (undefined :: a)
primArrayContents :: PrimArray a -> Ptr a
{-# INLINE primArrayContents #-}
primArrayContents (PrimArray ba) =
let addr# = byteArrayContents# ba in Ptr addr#
mutablePrimArrayContents :: MutablePrimArray s a -> Ptr a
{-# INLINE mutablePrimArrayContents #-}
mutablePrimArrayContents (MutablePrimArray mba#) =
let addr# = byteArrayContents# (unsafeCoerce# mba#) in Ptr addr#
withPrimArrayContents :: PrimArray a -> (Ptr a -> IO b) -> IO b
{-# INLINE withPrimArrayContents #-}
withPrimArrayContents (PrimArray ba#) f = do
let addr# = byteArrayContents# ba#
ptr = Ptr addr#
b <- f ptr
primitive_ (touch# ba#)
return b
withMutablePrimArrayContents :: MutablePrimArray RealWorld a -> (Ptr a -> IO b) -> IO b
{-# INLINE withMutablePrimArrayContents #-}
withMutablePrimArrayContents (MutablePrimArray mba#) f = do
let addr# = byteArrayContents# (unsafeCoerce# mba#)
ptr = Ptr addr#
b <- f ptr
primitive_ (touch# mba#)
return b
isPrimArrayPinned :: PrimArray a -> Bool
{-# INLINE isPrimArrayPinned #-}
isPrimArrayPinned (PrimArray ba#) = tagToEnum# (isByteArrayPinned# ba#)
isMutablePrimArrayPinned :: MutablePrimArray s a -> Bool
{-# INLINE isMutablePrimArrayPinned #-}
isMutablePrimArrayPinned (MutablePrimArray mba#) =
tagToEnum# (isByteArrayPinned# (unsafeCoerce# mba#))
castArray :: (Arr marr arr a, Cast a b) => arr a -> arr b
castArray = unsafeCoerce#
castMutableArray :: (Arr marr arr a, Cast a b) => marr s a -> marr s b
castMutableArray = unsafeCoerce#