Initial commit
This commit is contained in:
La Ancapo
commit
c101616e62
309 changed files with 53937 additions and 0 deletions
655
bundled/Basement/UArray/Base.hs
Normal file
655
bundled/Basement/UArray/Base.hs
Normal file
|
|
@ -0,0 +1,655 @@
|
|||
{-# LANGUAGE BangPatterns #-}
|
||||
{-# LANGUAGE DeriveDataTypeable #-}
|
||||
{-# LANGUAGE TypeFamilies #-}
|
||||
{-# LANGUAGE RebindableSyntax #-}
|
||||
{-# LANGUAGE NoImplicitPrelude #-}
|
||||
{-# LANGUAGE MagicHash #-}
|
||||
{-# LANGUAGE UnboxedTuples #-}
|
||||
{-# LANGUAGE ScopedTypeVariables #-}
|
||||
{-# LANGUAGE ViewPatterns #-}
|
||||
{-# LANGUAGE RankNTypes #-}
|
||||
{-# LANGUAGE FlexibleInstances, MultiParamTypeClasses #-}
|
||||
module Basement.UArray.Base
|
||||
( MUArray(..)
|
||||
, UArray(..)
|
||||
, MUArrayBackend(..)
|
||||
, UArrayBackend(..)
|
||||
-- * New mutable array creation
|
||||
, newUnpinned
|
||||
, newPinned
|
||||
, newNative
|
||||
, newNative_
|
||||
, new
|
||||
-- * Pinning status
|
||||
, isPinned
|
||||
, isMutablePinned
|
||||
-- * Mutable array accessor
|
||||
, unsafeRead
|
||||
, unsafeWrite
|
||||
-- * Freezing routines
|
||||
, unsafeFreezeShrink
|
||||
, unsafeFreeze
|
||||
, unsafeThaw
|
||||
, thaw
|
||||
, copy
|
||||
-- * Array accessor
|
||||
, unsafeIndex
|
||||
, unsafeIndexer
|
||||
, onBackend
|
||||
, onBackendPure
|
||||
, onBackendPure'
|
||||
, onBackendPrim
|
||||
, onMutableBackend
|
||||
, unsafeDewrap
|
||||
, unsafeDewrap2
|
||||
-- * Basic lowlevel functions
|
||||
, vFromListN
|
||||
, empty
|
||||
, length
|
||||
, offset
|
||||
, ValidRange(..)
|
||||
, offsetsValidRange
|
||||
, equal
|
||||
, equalMemcmp
|
||||
, compare
|
||||
, copyAt
|
||||
, unsafeCopyAtRO
|
||||
, toBlock
|
||||
-- * temporary
|
||||
, pureST
|
||||
) where
|
||||
|
||||
import GHC.Prim
|
||||
import GHC.Types
|
||||
import GHC.Ptr
|
||||
import GHC.ST
|
||||
import Basement.Compat.Primitive
|
||||
import Basement.Monad
|
||||
import Basement.PrimType
|
||||
import Basement.Compat.Base
|
||||
import Basement.Compat.C.Types
|
||||
import Basement.Compat.Semigroup
|
||||
import qualified Basement.Runtime as Runtime
|
||||
import Data.Proxy
|
||||
import qualified Basement.Compat.ExtList as List
|
||||
import qualified Basement.Alg.Class as Alg
|
||||
import Basement.Types.OffsetSize
|
||||
import Basement.FinalPtr
|
||||
import Basement.NormalForm
|
||||
import Basement.Block (MutableBlock(..), Block(..))
|
||||
import qualified Basement.Block as BLK
|
||||
import qualified Basement.Block.Mutable as MBLK
|
||||
import Basement.Numerical.Additive
|
||||
import Basement.Bindings.Memory
|
||||
import System.IO.Unsafe (unsafeDupablePerformIO)
|
||||
|
||||
-- | A Mutable array of types built on top of GHC primitive.
|
||||
--
|
||||
-- Element in this array can be modified in place.
|
||||
data MUArray ty st = MUArray {-# UNPACK #-} !(Offset ty)
|
||||
{-# UNPACK #-} !(CountOf ty)
|
||||
!(MUArrayBackend ty st)
|
||||
|
||||
data MUArrayBackend ty st = MUArrayMBA (MutableBlock ty st) | MUArrayAddr (FinalPtr ty)
|
||||
|
||||
|
||||
instance PrimType ty => Alg.Indexable (Ptr ty) ty where
|
||||
index (Ptr addr) = primAddrIndex addr
|
||||
|
||||
instance Alg.Indexable (Ptr Word8) Word64 where
|
||||
index (Ptr addr) = primAddrIndex addr
|
||||
|
||||
instance (PrimMonad prim, PrimType ty) => Alg.RandomAccess (Ptr ty) prim ty where
|
||||
read (Ptr addr) = primAddrRead addr
|
||||
write (Ptr addr) = primAddrWrite addr
|
||||
|
||||
-- | An array of type built on top of GHC primitive.
|
||||
--
|
||||
-- The elements need to have fixed sized and the representation is a
|
||||
-- packed contiguous array in memory that can easily be passed
|
||||
-- to foreign interface
|
||||
data UArray ty = UArray {-# UNPACK #-} !(Offset ty)
|
||||
{-# UNPACK #-} !(CountOf ty)
|
||||
!(UArrayBackend ty)
|
||||
deriving (Typeable)
|
||||
|
||||
data UArrayBackend ty = UArrayBA !(Block ty) | UArrayAddr !(FinalPtr ty)
|
||||
deriving (Typeable)
|
||||
|
||||
instance Data ty => Data (UArray ty) where
|
||||
dataTypeOf _ = arrayType
|
||||
toConstr _ = error "toConstr"
|
||||
gunfold _ _ = error "gunfold"
|
||||
|
||||
arrayType :: DataType
|
||||
arrayType = mkNoRepType "Basement.UArray"
|
||||
|
||||
instance NormalForm (UArray ty) where
|
||||
toNormalForm (UArray _ _ !_) = ()
|
||||
instance (PrimType ty, Show ty) => Show (UArray ty) where
|
||||
show v = show (toList v)
|
||||
instance (PrimType ty, Eq ty) => Eq (UArray ty) where
|
||||
(==) = equal
|
||||
instance (PrimType ty, Ord ty) => Ord (UArray ty) where
|
||||
{-# SPECIALIZE instance Ord (UArray Word8) #-}
|
||||
compare = vCompare
|
||||
|
||||
instance PrimType ty => Semigroup (UArray ty) where
|
||||
(<>) = append
|
||||
instance PrimType ty => Monoid (UArray ty) where
|
||||
mempty = empty
|
||||
mconcat = concat
|
||||
|
||||
instance PrimType ty => IsList (UArray ty) where
|
||||
type Item (UArray ty) = ty
|
||||
fromList = vFromList
|
||||
fromListN len = vFromListN (CountOf len)
|
||||
toList = vToList
|
||||
|
||||
length :: UArray ty -> CountOf ty
|
||||
length (UArray _ len _) = len
|
||||
{-# INLINE[1] length #-}
|
||||
|
||||
offset :: UArray ty -> Offset ty
|
||||
offset (UArray ofs _ _) = ofs
|
||||
{-# INLINE[1] offset #-}
|
||||
|
||||
data ValidRange ty = ValidRange {-# UNPACK #-} !(Offset ty) {-# UNPACK #-} !(Offset ty)
|
||||
|
||||
offsetsValidRange :: UArray ty -> ValidRange ty
|
||||
offsetsValidRange (UArray ofs len _) = ValidRange ofs (ofs `offsetPlusE` len)
|
||||
|
||||
-- | Return if the array is pinned in memory
|
||||
--
|
||||
-- note that Foreign array are considered pinned
|
||||
isPinned :: UArray ty -> PinnedStatus
|
||||
isPinned (UArray _ _ (UArrayAddr {})) = Pinned
|
||||
isPinned (UArray _ _ (UArrayBA blk)) = BLK.isPinned blk
|
||||
|
||||
-- | Return if a mutable array is pinned in memory
|
||||
isMutablePinned :: MUArray ty st -> PinnedStatus
|
||||
isMutablePinned (MUArray _ _ (MUArrayAddr {})) = Pinned
|
||||
isMutablePinned (MUArray _ _ (MUArrayMBA mb)) = BLK.isMutablePinned mb
|
||||
|
||||
-- | Create a new pinned mutable array of size @n.
|
||||
--
|
||||
-- all the cells are uninitialized and could contains invalid values.
|
||||
--
|
||||
-- All mutable arrays are allocated on a 64 bits aligned addresses
|
||||
newPinned :: forall prim ty . (PrimMonad prim, PrimType ty) => CountOf ty -> prim (MUArray ty (PrimState prim))
|
||||
newPinned n = MUArray 0 n . MUArrayMBA <$> MBLK.newPinned n
|
||||
|
||||
-- | Create a new unpinned mutable array of size @n elements.
|
||||
--
|
||||
-- If the size exceeds a GHC-defined threshold, then the memory will be
|
||||
-- pinned. To be certain about pinning status with small size, use 'newPinned'
|
||||
newUnpinned :: forall prim ty . (PrimMonad prim, PrimType ty) => CountOf ty -> prim (MUArray ty (PrimState prim))
|
||||
newUnpinned n = MUArray 0 n . MUArrayMBA <$> MBLK.new n
|
||||
|
||||
newNative :: (PrimMonad prim, PrimType ty)
|
||||
=> CountOf ty
|
||||
-> (MutableBlock ty (PrimState prim) -> prim a)
|
||||
-> prim (a, MUArray ty (PrimState prim))
|
||||
newNative n f = do
|
||||
mb <- MBLK.new n
|
||||
a <- f mb
|
||||
pure (a, MUArray 0 n (MUArrayMBA mb))
|
||||
|
||||
-- | Same as newNative but expect no extra return value from f
|
||||
newNative_ :: (PrimMonad prim, PrimType ty)
|
||||
=> CountOf ty
|
||||
-> (MutableBlock ty (PrimState prim) -> prim ())
|
||||
-> prim (MUArray ty (PrimState prim))
|
||||
newNative_ n f = do
|
||||
mb <- MBLK.new n
|
||||
f mb
|
||||
pure (MUArray 0 n (MUArrayMBA mb))
|
||||
|
||||
-- | Create a new mutable array of size @n.
|
||||
--
|
||||
-- When memory for a new array is allocated, we decide if that memory region
|
||||
-- should be pinned (will not be copied around by GC) or unpinned (can be
|
||||
-- moved around by GC) depending on its size.
|
||||
--
|
||||
-- You can change the threshold value used by setting the environment variable
|
||||
-- @HS_FOUNDATION_UARRAY_UNPINNED_MAX@.
|
||||
new :: (PrimMonad prim, PrimType ty) => CountOf ty -> prim (MUArray ty (PrimState prim))
|
||||
new sz
|
||||
| sizeRecast sz <= maxSizeUnpinned = newUnpinned sz
|
||||
| otherwise = newPinned sz
|
||||
where
|
||||
-- Safe to use here: If the value changes during runtime, this will only
|
||||
-- have an impact on newly created arrays.
|
||||
maxSizeUnpinned = Runtime.unsafeUArrayUnpinnedMaxSize
|
||||
{-# INLINE new #-}
|
||||
|
||||
-- | read from a cell in a mutable array without bounds checking.
|
||||
--
|
||||
-- Reading from invalid memory can return unpredictable and invalid values.
|
||||
-- use 'read' if unsure.
|
||||
unsafeRead :: (PrimMonad prim, PrimType ty) => MUArray ty (PrimState prim) -> Offset ty -> prim ty
|
||||
unsafeRead (MUArray start _ (MUArrayMBA (MutableBlock mba))) i = primMbaRead mba (start + i)
|
||||
unsafeRead (MUArray start _ (MUArrayAddr fptr)) i = withFinalPtr fptr $ \(Ptr addr) -> primAddrRead addr (start + i)
|
||||
{-# INLINE unsafeRead #-}
|
||||
|
||||
|
||||
-- | write to a cell in a mutable array without bounds checking.
|
||||
--
|
||||
-- Writing with invalid bounds will corrupt memory and your program will
|
||||
-- become unreliable. use 'write' if unsure.
|
||||
unsafeWrite :: (PrimMonad prim, PrimType ty) => MUArray ty (PrimState prim) -> Offset ty -> ty -> prim ()
|
||||
unsafeWrite (MUArray start _ (MUArrayMBA mb)) i v = MBLK.unsafeWrite mb (start+i) v
|
||||
unsafeWrite (MUArray start _ (MUArrayAddr fptr)) i v = withFinalPtr fptr $ \(Ptr addr) -> primAddrWrite addr (start+i) v
|
||||
{-# INLINE unsafeWrite #-}
|
||||
|
||||
-- | Return the element at a specific index from an array without bounds checking.
|
||||
--
|
||||
-- Reading from invalid memory can return unpredictable and invalid values.
|
||||
-- use 'index' if unsure.
|
||||
unsafeIndex :: forall ty . PrimType ty => UArray ty -> Offset ty -> ty
|
||||
unsafeIndex (UArray start _ (UArrayBA ba)) n = BLK.unsafeIndex ba (start + n)
|
||||
unsafeIndex (UArray start _ (UArrayAddr fptr)) n = withUnsafeFinalPtr fptr (\(Ptr addr) -> return (primAddrIndex addr (start+n)) :: IO ty)
|
||||
{-# INLINE unsafeIndex #-}
|
||||
|
||||
unsafeIndexer :: (PrimMonad prim, PrimType ty) => UArray ty -> ((Offset ty -> ty) -> prim a) -> prim a
|
||||
unsafeIndexer (UArray start _ (UArrayBA ba)) f = f (\n -> BLK.unsafeIndex ba (start + n))
|
||||
unsafeIndexer (UArray start _ (UArrayAddr fptr)) f = withFinalPtr fptr $ \(Ptr addr) -> f (\n -> primAddrIndex addr (start + n))
|
||||
{-# INLINE unsafeIndexer #-}
|
||||
|
||||
-- | Freeze a mutable array into an array.
|
||||
--
|
||||
-- the MUArray must not be changed after freezing.
|
||||
unsafeFreeze :: PrimMonad prim => MUArray ty (PrimState prim) -> prim (UArray ty)
|
||||
unsafeFreeze (MUArray start len (MUArrayMBA mba)) =
|
||||
UArray start len . UArrayBA <$> MBLK.unsafeFreeze mba
|
||||
unsafeFreeze (MUArray start len (MUArrayAddr fptr)) =
|
||||
pure $ UArray start len (UArrayAddr fptr)
|
||||
{-# INLINE unsafeFreeze #-}
|
||||
|
||||
unsafeFreezeShrink :: (PrimType ty, PrimMonad prim) => MUArray ty (PrimState prim) -> CountOf ty -> prim (UArray ty)
|
||||
unsafeFreezeShrink (MUArray start _ backend) n = unsafeFreeze (MUArray start n backend)
|
||||
{-# INLINE unsafeFreezeShrink #-}
|
||||
|
||||
-- | Thaw an immutable array.
|
||||
--
|
||||
-- The UArray must not be used after thawing.
|
||||
unsafeThaw :: (PrimType ty, PrimMonad prim) => UArray ty -> prim (MUArray ty (PrimState prim))
|
||||
unsafeThaw (UArray start len (UArrayBA blk)) = MUArray start len . MUArrayMBA <$> BLK.unsafeThaw blk
|
||||
unsafeThaw (UArray start len (UArrayAddr fptr)) = pure $ MUArray start len (MUArrayAddr fptr)
|
||||
{-# INLINE unsafeThaw #-}
|
||||
|
||||
-- | Thaw an array to a mutable array.
|
||||
--
|
||||
-- the array is not modified, instead a new mutable array is created
|
||||
-- and every values is copied, before returning the mutable array.
|
||||
thaw :: (PrimMonad prim, PrimType ty) => UArray ty -> prim (MUArray ty (PrimState prim))
|
||||
thaw array = do
|
||||
ma <- new (length array)
|
||||
unsafeCopyAtRO ma azero array (Offset 0) (length array)
|
||||
pure ma
|
||||
{-# INLINE thaw #-}
|
||||
|
||||
-- | Copy every cells of an existing array to a new array
|
||||
copy :: PrimType ty => UArray ty -> UArray ty
|
||||
copy array = runST (thaw array >>= unsafeFreeze)
|
||||
|
||||
|
||||
onBackend :: (Block ty -> a)
|
||||
-> (FinalPtr ty -> Ptr ty -> ST s a)
|
||||
-> UArray ty
|
||||
-> a
|
||||
onBackend onBa _ (UArray _ _ (UArrayBA ba)) = onBa ba
|
||||
onBackend _ onAddr (UArray _ _ (UArrayAddr fptr)) = withUnsafeFinalPtr fptr $ \ptr@(Ptr !_) ->
|
||||
onAddr fptr ptr
|
||||
{-# INLINE onBackend #-}
|
||||
|
||||
onBackendPure :: (Block ty -> a)
|
||||
-> (Ptr ty -> a)
|
||||
-> UArray ty
|
||||
-> a
|
||||
onBackendPure goBA goAddr arr = onBackend goBA (\_ -> pureST . goAddr) arr
|
||||
{-# INLINE onBackendPure #-}
|
||||
|
||||
onBackendPure' :: forall ty a . PrimType ty
|
||||
=> UArray ty
|
||||
-> (forall container. Alg.Indexable container ty
|
||||
=> container -> Offset ty -> Offset ty -> a)
|
||||
-> a
|
||||
onBackendPure' arr f = onBackendPure f' f' arr
|
||||
where f' :: Alg.Indexable container ty => container -> a
|
||||
f' c = f c start end
|
||||
where (ValidRange !start !end) = offsetsValidRange arr
|
||||
{-# INLINE onBackendPure' #-}
|
||||
|
||||
onBackendPrim :: PrimMonad prim
|
||||
=> (Block ty -> prim a)
|
||||
-> (FinalPtr ty -> prim a)
|
||||
-> UArray ty
|
||||
-> prim a
|
||||
onBackendPrim onBa _ (UArray _ _ (UArrayBA ba)) = onBa ba
|
||||
onBackendPrim _ onAddr (UArray _ _ (UArrayAddr fptr)) = onAddr fptr
|
||||
{-# INLINE onBackendPrim #-}
|
||||
|
||||
onMutableBackend :: PrimMonad prim
|
||||
=> (MutableBlock ty (PrimState prim) -> prim a)
|
||||
-> (FinalPtr ty -> prim a)
|
||||
-> MUArray ty (PrimState prim)
|
||||
-> prim a
|
||||
onMutableBackend onMba _ (MUArray _ _ (MUArrayMBA mba)) = onMba mba
|
||||
onMutableBackend _ onAddr (MUArray _ _ (MUArrayAddr fptr)) = onAddr fptr
|
||||
{-# INLINE onMutableBackend #-}
|
||||
|
||||
|
||||
unsafeDewrap :: (Block ty -> Offset ty -> a)
|
||||
-> (Ptr ty -> Offset ty -> ST s a)
|
||||
-> UArray ty
|
||||
-> a
|
||||
unsafeDewrap _ g (UArray start _ (UArrayAddr fptr)) = withUnsafeFinalPtr fptr $ \ptr -> g ptr start
|
||||
unsafeDewrap f _ (UArray start _ (UArrayBA ba)) = f ba start
|
||||
{-# INLINE unsafeDewrap #-}
|
||||
|
||||
unsafeDewrap2 :: (ByteArray# -> ByteArray# -> a)
|
||||
-> (Ptr ty -> Ptr ty -> ST s a)
|
||||
-> (ByteArray# -> Ptr ty -> ST s a)
|
||||
-> (Ptr ty -> ByteArray# -> ST s a)
|
||||
-> UArray ty
|
||||
-> UArray ty
|
||||
-> a
|
||||
unsafeDewrap2 f g h i (UArray _ _ back1) (UArray _ _ back2) =
|
||||
case (back1, back2) of
|
||||
(UArrayBA (Block ba1), UArrayBA (Block ba2)) -> f ba1 ba2
|
||||
(UArrayAddr fptr1, UArrayAddr fptr2) -> withUnsafeFinalPtr fptr1 $ \ptr1 -> withFinalPtr fptr2 $ \ptr2 -> g ptr1 ptr2
|
||||
(UArrayBA (Block ba1), UArrayAddr fptr2) -> withUnsafeFinalPtr fptr2 $ \ptr2 -> h ba1 ptr2
|
||||
(UArrayAddr fptr1, UArrayBA (Block ba2)) -> withUnsafeFinalPtr fptr1 $ \ptr1 -> i ptr1 ba2
|
||||
{-# INLINE [2] unsafeDewrap2 #-}
|
||||
|
||||
pureST :: a -> ST s a
|
||||
pureST = pure
|
||||
|
||||
-- | make an array from a list of elements.
|
||||
vFromList :: forall ty . PrimType ty => [ty] -> UArray ty
|
||||
vFromList l = runST $ do
|
||||
a <- newNative_ len copyList
|
||||
unsafeFreeze a
|
||||
where
|
||||
len = List.length l
|
||||
copyList :: MutableBlock ty s -> ST s ()
|
||||
copyList mb = loop 0 l
|
||||
where
|
||||
loop _ [] = pure ()
|
||||
loop !i (x:xs) = MBLK.unsafeWrite mb i x >> loop (i+1) xs
|
||||
|
||||
-- | Make an array from a list of elements with a size hint.
|
||||
--
|
||||
-- The list should be of the same size as the hint, as otherwise:
|
||||
--
|
||||
-- * The length of the list is smaller than the hint:
|
||||
-- the array allocated is of the size of the hint, but is sliced
|
||||
-- to only represent the valid bits
|
||||
-- * The length of the list is bigger than the hint:
|
||||
-- The allocated array is the size of the hint, and the list is truncated to
|
||||
-- fit.
|
||||
vFromListN :: forall ty . PrimType ty => CountOf ty -> [ty] -> UArray ty
|
||||
vFromListN len l = runST $ do
|
||||
(sz, ma) <- newNative len copyList
|
||||
unsafeFreezeShrink ma sz
|
||||
where
|
||||
copyList :: MutableBlock ty s -> ST s (CountOf ty)
|
||||
copyList mb = loop 0 l
|
||||
where
|
||||
loop !i [] = pure (offsetAsSize i)
|
||||
loop !i (x:xs)
|
||||
| i .==# len = pure (offsetAsSize i)
|
||||
| otherwise = MBLK.unsafeWrite mb i x >> loop (i+1) xs
|
||||
|
||||
-- | transform an array to a list.
|
||||
vToList :: forall ty . PrimType ty => UArray ty -> [ty]
|
||||
vToList a
|
||||
| len == 0 = []
|
||||
| otherwise = unsafeDewrap goBa goPtr a
|
||||
where
|
||||
!len = length a
|
||||
goBa (Block ba) start = loop start
|
||||
where
|
||||
!end = start `offsetPlusE` len
|
||||
loop !i | i == end = []
|
||||
| otherwise = primBaIndex ba i : loop (i+1)
|
||||
goPtr (Ptr addr) start = pureST (loop start)
|
||||
where
|
||||
!end = start `offsetPlusE` len
|
||||
loop !i | i == end = []
|
||||
| otherwise = primAddrIndex addr i : loop (i+1)
|
||||
|
||||
-- | Check if two vectors are identical
|
||||
equal :: (PrimType ty, Eq ty) => UArray ty -> UArray ty -> Bool
|
||||
equal a b
|
||||
| la /= lb = False
|
||||
| otherwise = unsafeDewrap2 goBaBa goPtrPtr goBaPtr goPtrBa a b
|
||||
where
|
||||
!start1 = offset a
|
||||
!start2 = offset b
|
||||
!end = start1 `offsetPlusE` la
|
||||
!la = length a
|
||||
!lb = length b
|
||||
goBaBa ba1 ba2 = loop start1 start2
|
||||
where
|
||||
loop !i !o | i == end = True
|
||||
| otherwise = primBaIndex ba1 i == primBaIndex ba2 o && loop (i+o1) (o+o1)
|
||||
goPtrPtr (Ptr addr1) (Ptr addr2) = pureST (loop start1 start2)
|
||||
where
|
||||
loop !i !o | i == end = True
|
||||
| otherwise = primAddrIndex addr1 i == primAddrIndex addr2 o && loop (i+o1) (o+o1)
|
||||
goBaPtr ba1 (Ptr addr2) = pureST (loop start1 start2)
|
||||
where
|
||||
loop !i !o | i == end = True
|
||||
| otherwise = primBaIndex ba1 i == primAddrIndex addr2 o && loop (i+o1) (o+o1)
|
||||
goPtrBa (Ptr addr1) ba2 = pureST (loop start1 start2)
|
||||
where
|
||||
loop !i !o | i == end = True
|
||||
| otherwise = primAddrIndex addr1 i == primBaIndex ba2 o && loop (i+o1) (o+o1)
|
||||
|
||||
o1 = Offset (I# 1#)
|
||||
{-# RULES "UArray/Eq/Word8" [3] equal = equalBytes #-}
|
||||
{-# INLINEABLE [2] equal #-}
|
||||
|
||||
equalBytes :: UArray Word8 -> UArray Word8 -> Bool
|
||||
equalBytes a b
|
||||
| la /= lb = False
|
||||
| otherwise = memcmp a b (sizeInBytes la) == 0
|
||||
where
|
||||
!la = length a
|
||||
!lb = length b
|
||||
|
||||
equalMemcmp :: PrimType ty => UArray ty -> UArray ty -> Bool
|
||||
equalMemcmp a b
|
||||
| la /= lb = False
|
||||
| otherwise = memcmp a b (sizeInBytes la) == 0
|
||||
where
|
||||
!la = length a
|
||||
!lb = length b
|
||||
|
||||
-- | Compare 2 vectors
|
||||
vCompare :: (Ord ty, PrimType ty) => UArray ty -> UArray ty -> Ordering
|
||||
vCompare a@(UArray start1 la _) b@(UArray start2 lb _) = unsafeDewrap2 goBaBa goPtrPtr goBaPtr goPtrBa a b
|
||||
where
|
||||
!end = start1 `offsetPlusE` min la lb
|
||||
o1 = Offset (I# 1#)
|
||||
goBaBa ba1 ba2 = loop start1 start2
|
||||
where
|
||||
loop !i !o | i == end = la `compare` lb
|
||||
| v1 == v2 = loop (i + o1) (o + o1)
|
||||
| otherwise = v1 `compare` v2
|
||||
where v1 = primBaIndex ba1 i
|
||||
v2 = primBaIndex ba2 o
|
||||
goPtrPtr (Ptr addr1) (Ptr addr2) = pureST (loop start1 start2)
|
||||
where
|
||||
loop !i !o | i == end = la `compare` lb
|
||||
| v1 == v2 = loop (i + o1) (o + o1)
|
||||
| otherwise = v1 `compare` v2
|
||||
where v1 = primAddrIndex addr1 i
|
||||
v2 = primAddrIndex addr2 o
|
||||
goBaPtr ba1 (Ptr addr2) = pureST (loop start1 start2)
|
||||
where
|
||||
loop !i !o | i == end = la `compare` lb
|
||||
| v1 == v2 = loop (i + o1) (o + o1)
|
||||
| otherwise = v1 `compare` v2
|
||||
where v1 = primBaIndex ba1 i
|
||||
v2 = primAddrIndex addr2 o
|
||||
goPtrBa (Ptr addr1) ba2 = pureST (loop start1 start2)
|
||||
where
|
||||
loop !i !o | i == end = la `compare` lb
|
||||
| v1 == v2 = loop (i + o1) (o + o1)
|
||||
| otherwise = v1 `compare` v2
|
||||
where v1 = primAddrIndex addr1 i
|
||||
v2 = primBaIndex ba2 o
|
||||
-- {-# SPECIALIZE [3] vCompare :: UArray Word8 -> UArray Word8 -> Ordering = vCompareBytes #-}
|
||||
{-# RULES "UArray/Ord/Word8" [3] vCompare = vCompareBytes #-}
|
||||
{-# INLINEABLE [2] vCompare #-}
|
||||
|
||||
vCompareBytes :: UArray Word8 -> UArray Word8 -> Ordering
|
||||
vCompareBytes = vCompareMemcmp
|
||||
|
||||
vCompareMemcmp :: (Ord ty, PrimType ty) => UArray ty -> UArray ty -> Ordering
|
||||
vCompareMemcmp a b = cintToOrdering $ memcmp a b sz
|
||||
where
|
||||
la = length a
|
||||
lb = length b
|
||||
sz = sizeInBytes $ min la lb
|
||||
cintToOrdering :: CInt -> Ordering
|
||||
cintToOrdering 0 = la `compare` lb
|
||||
cintToOrdering r | r < 0 = LT
|
||||
| otherwise = GT
|
||||
{-# SPECIALIZE [3] vCompareMemcmp :: UArray Word8 -> UArray Word8 -> Ordering #-}
|
||||
|
||||
memcmp :: PrimType ty => UArray ty -> UArray ty -> CountOf Word8 -> CInt
|
||||
memcmp a@(UArray (offsetInBytes -> o1) _ _) b@(UArray (offsetInBytes -> o2) _ _) sz = unsafeDewrap2
|
||||
(\s1 s2 -> unsafeDupablePerformIO $ sysHsMemcmpBaBa s1 o1 s2 o2 sz)
|
||||
(\s1 s2 -> unsafePrimToST $ sysHsMemcmpPtrPtr s1 o1 s2 o2 sz)
|
||||
(\s1 s2 -> unsafePrimToST $ sysHsMemcmpBaPtr s1 o1 s2 o2 sz)
|
||||
(\s1 s2 -> unsafePrimToST $ sysHsMemcmpPtrBa s1 o1 s2 o2 sz)
|
||||
a b
|
||||
{-# SPECIALIZE [3] memcmp :: UArray Word8 -> UArray Word8 -> CountOf Word8 -> CInt #-}
|
||||
|
||||
-- | Copy a number of elements from an array to another array with offsets
|
||||
copyAt :: forall prim ty . (PrimMonad prim, PrimType ty)
|
||||
=> MUArray ty (PrimState prim) -- ^ destination array
|
||||
-> Offset ty -- ^ offset at destination
|
||||
-> MUArray ty (PrimState prim) -- ^ source array
|
||||
-> Offset ty -- ^ offset at source
|
||||
-> CountOf ty -- ^ number of elements to copy
|
||||
-> prim ()
|
||||
copyAt (MUArray dstStart _ (MUArrayMBA (MutableBlock dstMba))) ed (MUArray srcStart _ (MUArrayMBA (MutableBlock srcBa))) es n =
|
||||
primitive $ \st -> (# copyMutableByteArray# srcBa os dstMba od nBytes st, () #)
|
||||
where
|
||||
!sz = primSizeInBytes (Proxy :: Proxy ty)
|
||||
!(Offset (I# os)) = offsetOfE sz (srcStart + es)
|
||||
!(Offset (I# od)) = offsetOfE sz (dstStart + ed)
|
||||
!(CountOf (I# nBytes)) = sizeOfE sz n
|
||||
copyAt (MUArray dstStart _ (MUArrayMBA (MutableBlock dstMba))) ed (MUArray srcStart _ (MUArrayAddr srcFptr)) es n =
|
||||
withFinalPtr srcFptr $ \srcPtr ->
|
||||
let !(Ptr srcAddr) = srcPtr `plusPtr` os
|
||||
in primitive $ \s -> (# copyAddrToByteArray# srcAddr dstMba od nBytes s, () #)
|
||||
where
|
||||
!sz = primSizeInBytes (Proxy :: Proxy ty)
|
||||
!(Offset os) = offsetOfE sz (srcStart + es)
|
||||
!(Offset (I# od)) = offsetOfE sz (dstStart + ed)
|
||||
!(CountOf (I# nBytes)) = sizeOfE sz n
|
||||
copyAt dst od src os n = loop od os
|
||||
where
|
||||
!endIndex = os `offsetPlusE` n
|
||||
loop !d !i
|
||||
| i == endIndex = return ()
|
||||
| otherwise = unsafeRead src i >>= unsafeWrite dst d >> loop (d+1) (i+1)
|
||||
|
||||
-- TODO Optimise with copyByteArray#
|
||||
-- | Copy @n@ sequential elements from the specified offset in a source array
|
||||
-- to the specified position in a destination array.
|
||||
--
|
||||
-- This function does not check bounds. Accessing invalid memory can return
|
||||
-- unpredictable and invalid values.
|
||||
unsafeCopyAtRO :: forall prim ty . (PrimMonad prim, PrimType ty)
|
||||
=> MUArray ty (PrimState prim) -- ^ destination array
|
||||
-> Offset ty -- ^ offset at destination
|
||||
-> UArray ty -- ^ source array
|
||||
-> Offset ty -- ^ offset at source
|
||||
-> CountOf ty -- ^ number of elements to copy
|
||||
-> prim ()
|
||||
unsafeCopyAtRO (MUArray dstStart _ (MUArrayMBA (MutableBlock dstMba))) ed (UArray srcStart _ (UArrayBA (Block srcBa))) es n =
|
||||
primitive $ \st -> (# copyByteArray# srcBa os dstMba od nBytes st, () #)
|
||||
where
|
||||
sz = primSizeInBytes (Proxy :: Proxy ty)
|
||||
!(Offset (I# os)) = offsetOfE sz (srcStart+es)
|
||||
!(Offset (I# od)) = offsetOfE sz (dstStart+ed)
|
||||
!(CountOf (I# nBytes)) = sizeOfE sz n
|
||||
unsafeCopyAtRO (MUArray dstStart _ (MUArrayMBA (MutableBlock dstMba))) ed (UArray srcStart _ (UArrayAddr srcFptr)) es n =
|
||||
withFinalPtr srcFptr $ \srcPtr ->
|
||||
let !(Ptr srcAddr) = srcPtr `plusPtr` os
|
||||
in primitive $ \s -> (# copyAddrToByteArray# srcAddr dstMba od nBytes s, () #)
|
||||
where
|
||||
sz = primSizeInBytes (Proxy :: Proxy ty)
|
||||
!(Offset os) = offsetOfE sz (srcStart+es)
|
||||
!(Offset (I# od)) = offsetOfE sz (dstStart+ed)
|
||||
!(CountOf (I# nBytes)) = sizeOfE sz n
|
||||
unsafeCopyAtRO dst od src os n = loop od os
|
||||
where
|
||||
!endIndex = os `offsetPlusE` n
|
||||
loop d i
|
||||
| i == endIndex = return ()
|
||||
| otherwise = unsafeWrite dst d (unsafeIndex src i) >> loop (d+1) (i+1)
|
||||
|
||||
empty_ :: Block ()
|
||||
empty_ = runST $ primitive $ \s1 ->
|
||||
case newByteArray# 0# s1 of { (# s2, mba #) ->
|
||||
case unsafeFreezeByteArray# mba s2 of { (# s3, ba #) ->
|
||||
(# s3, Block ba #) }}
|
||||
|
||||
empty :: UArray ty
|
||||
empty = UArray 0 0 (UArrayBA $ Block ba) where !(Block ba) = empty_
|
||||
|
||||
-- | Append 2 arrays together by creating a new bigger array
|
||||
append :: PrimType ty => UArray ty -> UArray ty -> UArray ty
|
||||
append a b
|
||||
| la == azero = b
|
||||
| lb == azero = a
|
||||
| otherwise = runST $ do
|
||||
r <- new (la+lb)
|
||||
ma <- unsafeThaw a
|
||||
mb <- unsafeThaw b
|
||||
copyAt r (Offset 0) ma (Offset 0) la
|
||||
copyAt r (sizeAsOffset la) mb (Offset 0) lb
|
||||
unsafeFreeze r
|
||||
where
|
||||
!la = length a
|
||||
!lb = length b
|
||||
|
||||
concat :: forall ty . PrimType ty => [UArray ty] -> UArray ty
|
||||
concat original = runST $ do
|
||||
r <- new total
|
||||
goCopy r 0 original
|
||||
unsafeFreeze r
|
||||
where
|
||||
!total = size 0 original
|
||||
-- size
|
||||
size !sz [] = sz
|
||||
size !sz (x:xs) = size (length x + sz) xs
|
||||
|
||||
zero = Offset 0
|
||||
|
||||
goCopy r = loop
|
||||
where
|
||||
loop _ [] = pure ()
|
||||
loop !i (x:xs) = do
|
||||
unsafeCopyAtRO r i x zero lx
|
||||
loop (i `offsetPlusE` lx) xs
|
||||
where !lx = length x
|
||||
|
||||
-- | Create a Block from a UArray.
|
||||
--
|
||||
-- Note that because of the slice, the destination block
|
||||
-- is re-allocated and copied, unless the slice point
|
||||
-- at the whole array
|
||||
toBlock :: PrimType ty => UArray ty -> Block ty
|
||||
toBlock arr@(UArray start len (UArrayBA blk))
|
||||
| start == 0 && BLK.length blk == len = blk
|
||||
| otherwise = toBlock $ copy arr
|
||||
toBlock arr = toBlock $ copy arr
|
||||
Loading…
Add table
Add a link
Reference in a new issue