{-# Language BangPatterns #-}
{-# Language CPP #-}
{-# Language DeriveDataTypeable #-}
{-# Language MagicHash #-}
{-# Language RankNTypes #-}
{-# Language ScopedTypeVariables #-}
{-# Language TypeFamilies #-}
{-# Language UnboxedTuples #-}
module Data.Primitive.UnliftedArray
(
UnliftedArray(..)
, MutableUnliftedArray(..)
, PrimUnlifted(..)
, unsafeNewUnliftedArray
, newUnliftedArray
, setUnliftedArray
, sizeofUnliftedArray
, sizeofMutableUnliftedArray
, readUnliftedArray
, writeUnliftedArray
, indexUnliftedArray
, indexUnliftedArrayM
, unsafeFreezeUnliftedArray
, freezeUnliftedArray
, thawUnliftedArray
, runUnliftedArray
, sameMutableUnliftedArray
, copyUnliftedArray
, copyMutableUnliftedArray
, cloneUnliftedArray
, cloneMutableUnliftedArray
, unliftedArrayToList
, unliftedArrayFromList
, unliftedArrayFromListN
, foldrUnliftedArray
, foldrUnliftedArray'
, foldlUnliftedArray
, foldlUnliftedArray'
, mapUnliftedArray
) where
import Data.Typeable
import Control.Applicative
import GHC.Prim
import GHC.Base (Int(..),build)
import Control.Monad.Primitive
import Control.Monad.ST (runST,ST)
import Data.Monoid (Monoid,mappend)
import Data.Primitive.Internal.Compat ( isTrue# )
import qualified Data.List as L
import Data.Primitive.Array (Array)
import qualified Data.Primitive.Array as A
import Data.Primitive.ByteArray (ByteArray)
import qualified Data.Primitive.ByteArray as BA
import qualified Data.Primitive.PrimArray as PA
import qualified Data.Primitive.SmallArray as SA
import qualified Data.Primitive.MutVar as MV
import qualified Data.Monoid
import qualified GHC.MVar as GM (MVar(..))
import qualified GHC.Conc as GC (TVar(..))
import qualified GHC.Stable as GSP (StablePtr(..))
import qualified GHC.Weak as GW (Weak(..))
import qualified GHC.Conc.Sync as GCS (ThreadId(..))
import qualified GHC.Exts as E
import qualified GHC.ST as GHCST
#if MIN_VERSION_base(4,9,0)
import Data.Semigroup (Semigroup)
import qualified Data.Semigroup
#endif
#if MIN_VERSION_base(4,10,0)
import GHC.Exts (runRW#)
#elif MIN_VERSION_base(4,9,0)
import GHC.Base (runRW#)
#endif
data UnliftedArray e = UnliftedArray ArrayArray#
deriving (Typeable)
data MutableUnliftedArray s e = MutableUnliftedArray (MutableArrayArray# s)
deriving (Typeable)
class PrimUnlifted a where
toArrayArray# :: a -> ArrayArray#
fromArrayArray# :: ArrayArray# -> a
instance PrimUnlifted (UnliftedArray e) where
toArrayArray# (UnliftedArray aa#) = aa#
fromArrayArray# aa# = UnliftedArray aa#
instance PrimUnlifted (MutableUnliftedArray s e) where
toArrayArray# (MutableUnliftedArray maa#) = unsafeCoerce# maa#
fromArrayArray# aa# = MutableUnliftedArray (unsafeCoerce# aa#)
instance PrimUnlifted (Array a) where
toArrayArray# (A.Array a#) = unsafeCoerce# a#
fromArrayArray# aa# = A.Array (unsafeCoerce# aa#)
instance PrimUnlifted (A.MutableArray s a) where
toArrayArray# (A.MutableArray ma#) = unsafeCoerce# ma#
fromArrayArray# aa# = A.MutableArray (unsafeCoerce# aa#)
instance PrimUnlifted ByteArray where
toArrayArray# (BA.ByteArray ba#) = unsafeCoerce# ba#
fromArrayArray# aa# = BA.ByteArray (unsafeCoerce# aa#)
instance PrimUnlifted (BA.MutableByteArray s) where
toArrayArray# (BA.MutableByteArray mba#) = unsafeCoerce# mba#
fromArrayArray# aa# = BA.MutableByteArray (unsafeCoerce# aa#)
instance PrimUnlifted (PA.PrimArray a) where
toArrayArray# (PA.PrimArray ba#) = unsafeCoerce# ba#
fromArrayArray# aa# = PA.PrimArray (unsafeCoerce# aa#)
instance PrimUnlifted (PA.MutablePrimArray s a) where
toArrayArray# (PA.MutablePrimArray mba#) = unsafeCoerce# mba#
fromArrayArray# aa# = PA.MutablePrimArray (unsafeCoerce# aa#)
instance PrimUnlifted (SA.SmallArray a) where
toArrayArray# (SA.SmallArray sa#) = unsafeCoerce# sa#
fromArrayArray# aa# = SA.SmallArray (unsafeCoerce# aa#)
instance PrimUnlifted (SA.SmallMutableArray s a) where
toArrayArray# (SA.SmallMutableArray sma#) = unsafeCoerce# sma#
fromArrayArray# aa# = SA.SmallMutableArray (unsafeCoerce# aa#)
instance PrimUnlifted (MV.MutVar s a) where
toArrayArray# (MV.MutVar mv#) = unsafeCoerce# mv#
fromArrayArray# aa# = MV.MutVar (unsafeCoerce# aa#)
instance PrimUnlifted (GM.MVar a) where
toArrayArray# (GM.MVar mv#) = unsafeCoerce# mv#
fromArrayArray# mv# = GM.MVar (unsafeCoerce# mv#)
instance PrimUnlifted (GC.TVar a) where
toArrayArray# (GC.TVar tv#) = unsafeCoerce# tv#
fromArrayArray# tv# = GC.TVar (unsafeCoerce# tv#)
instance PrimUnlifted (GSP.StablePtr a) where
toArrayArray# (GSP.StablePtr tv#) = unsafeCoerce# tv#
fromArrayArray# tv# = GSP.StablePtr (unsafeCoerce# tv#)
instance PrimUnlifted (GW.Weak a) where
toArrayArray# (GW.Weak tv#) = unsafeCoerce# tv#
fromArrayArray# tv# = GW.Weak (unsafeCoerce# tv#)
instance PrimUnlifted GCS.ThreadId where
toArrayArray# (GCS.ThreadId tv#) = unsafeCoerce# tv#
fromArrayArray# tv# = GCS.ThreadId (unsafeCoerce# tv#)
die :: String -> String -> a
die fun problem = error $ "Data.Primitive.UnliftedArray." ++ fun ++ ": " ++ problem
unsafeNewUnliftedArray
:: (PrimMonad m)
=> Int
-> m (MutableUnliftedArray (PrimState m) a)
unsafeNewUnliftedArray (I# i#) = primitive $ \s -> case newArrayArray# i# s of
(# s', maa# #) -> (# s', MutableUnliftedArray maa# #)
{-# inline unsafeNewUnliftedArray #-}
setUnliftedArray
:: (PrimMonad m, PrimUnlifted a)
=> MutableUnliftedArray (PrimState m) a
-> a
-> m ()
setUnliftedArray mua v = loop $ sizeofMutableUnliftedArray mua - 1
where
loop i | i < 0 = return ()
| otherwise = writeUnliftedArray mua i v >> loop (i-1)
{-# inline setUnliftedArray #-}
newUnliftedArray
:: (PrimMonad m, PrimUnlifted a)
=> Int
-> a
-> m (MutableUnliftedArray (PrimState m) a)
newUnliftedArray len v =
unsafeNewUnliftedArray len >>= \mua -> setUnliftedArray mua v >> return mua
{-# inline newUnliftedArray #-}
sizeofUnliftedArray :: UnliftedArray e -> Int
sizeofUnliftedArray (UnliftedArray aa#) = I# (sizeofArrayArray# aa#)
{-# inline sizeofUnliftedArray #-}
sizeofMutableUnliftedArray :: MutableUnliftedArray s e -> Int
sizeofMutableUnliftedArray (MutableUnliftedArray maa#)
= I# (sizeofMutableArrayArray# maa#)
{-# inline sizeofMutableUnliftedArray #-}
indexUnliftedArrayU
:: PrimUnlifted a
=> UnliftedArray a
-> Int
-> (# a #)
indexUnliftedArrayU (UnliftedArray src#) (I# i#)
= case indexArrayArrayArray# src# i# of
aa# -> (# fromArrayArray# aa# #)
{-# inline indexUnliftedArrayU #-}
indexUnliftedArray
:: PrimUnlifted a
=> UnliftedArray a
-> Int
-> a
indexUnliftedArray ua i
= case indexUnliftedArrayU ua i of (# v #) -> v
{-# inline indexUnliftedArray #-}
indexUnliftedArrayM
:: (PrimUnlifted a, Monad m)
=> UnliftedArray a
-> Int
-> m a
indexUnliftedArrayM ua i
= case indexUnliftedArrayU ua i of
(# v #) -> return v
{-# inline indexUnliftedArrayM #-}
readUnliftedArray
:: (PrimMonad m, PrimUnlifted a)
=> MutableUnliftedArray (PrimState m) a
-> Int
-> m a
readUnliftedArray (MutableUnliftedArray maa#) (I# i#)
= primitive $ \s -> case readArrayArrayArray# maa# i# s of
(# s', aa# #) -> (# s', fromArrayArray# aa# #)
{-# inline readUnliftedArray #-}
writeUnliftedArray
:: (PrimMonad m, PrimUnlifted a)
=> MutableUnliftedArray (PrimState m) a
-> Int
-> a
-> m ()
writeUnliftedArray (MutableUnliftedArray maa#) (I# i#) a
= primitive_ (writeArrayArrayArray# maa# i# (toArrayArray# a))
{-# inline writeUnliftedArray #-}
unsafeFreezeUnliftedArray
:: (PrimMonad m)
=> MutableUnliftedArray (PrimState m) a
-> m (UnliftedArray a)
unsafeFreezeUnliftedArray (MutableUnliftedArray maa#)
= primitive $ \s -> case unsafeFreezeArrayArray# maa# s of
(# s', aa# #) -> (# s', UnliftedArray aa# #)
{-# inline unsafeFreezeUnliftedArray #-}
sameMutableUnliftedArray
:: MutableUnliftedArray s a
-> MutableUnliftedArray s a
-> Bool
sameMutableUnliftedArray (MutableUnliftedArray maa1#) (MutableUnliftedArray maa2#)
= isTrue# (sameMutableArrayArray# maa1# maa2#)
{-# inline sameMutableUnliftedArray #-}
copyUnliftedArray
:: (PrimMonad m)
=> MutableUnliftedArray (PrimState m) a
-> Int
-> UnliftedArray a
-> Int
-> Int
-> m ()
copyUnliftedArray
(MutableUnliftedArray dst) (I# doff)
(UnliftedArray src) (I# soff) (I# ln) =
primitive_ $ copyArrayArray# src soff dst doff ln
{-# inline copyUnliftedArray #-}
copyMutableUnliftedArray
:: (PrimMonad m)
=> MutableUnliftedArray (PrimState m) a
-> Int
-> MutableUnliftedArray (PrimState m) a
-> Int
-> Int
-> m ()
copyMutableUnliftedArray
(MutableUnliftedArray dst) (I# doff)
(MutableUnliftedArray src) (I# soff) (I# ln) =
primitive_ $ copyMutableArrayArray# src soff dst doff ln
{-# inline copyMutableUnliftedArray #-}
freezeUnliftedArray
:: (PrimMonad m)
=> MutableUnliftedArray (PrimState m) a
-> Int
-> Int
-> m (UnliftedArray a)
freezeUnliftedArray src off len = do
dst <- unsafeNewUnliftedArray len
copyMutableUnliftedArray dst 0 src off len
unsafeFreezeUnliftedArray dst
{-# inline freezeUnliftedArray #-}
thawUnliftedArray
:: (PrimMonad m)
=> UnliftedArray a
-> Int
-> Int
-> m (MutableUnliftedArray (PrimState m) a)
thawUnliftedArray src off len = do
dst <- unsafeNewUnliftedArray len
copyUnliftedArray dst 0 src off len
return dst
{-# inline thawUnliftedArray #-}
#if !MIN_VERSION_base(4,9,0)
unsafeCreateUnliftedArray
:: Int
-> (forall s. MutableUnliftedArray s a -> ST s ())
-> UnliftedArray a
unsafeCreateUnliftedArray 0 _ = emptyUnliftedArray
unsafeCreateUnliftedArray n f = runUnliftedArray $ do
mary <- unsafeNewUnliftedArray n
f mary
pure mary
runUnliftedArray
:: (forall s. ST s (MutableUnliftedArray s a))
-> UnliftedArray a
runUnliftedArray m = runST $ m >>= unsafeFreezeUnliftedArray
#else /* Below, runRW# is available. */
unsafeCreateUnliftedArray
:: Int
-> (forall s. MutableUnliftedArray s a -> ST s ())
-> UnliftedArray a
unsafeCreateUnliftedArray 0 _ = UnliftedArray (emptyArrayArray# (# #))
unsafeCreateUnliftedArray n f = runUnliftedArray $ do
mary <- unsafeNewUnliftedArray n
f mary
pure mary
runUnliftedArray
:: (forall s. ST s (MutableUnliftedArray s a))
-> UnliftedArray a
runUnliftedArray m = UnliftedArray (runUnliftedArray# m)
runUnliftedArray#
:: (forall s. ST s (MutableUnliftedArray s a))
-> ArrayArray#
runUnliftedArray# m = case runRW# $ \s ->
case unST m s of { (# s', MutableUnliftedArray mary# #) ->
unsafeFreezeArrayArray# mary# s'} of (# _, ary# #) -> ary#
unST :: ST s a -> State# s -> (# State# s, a #)
unST (GHCST.ST f) = f
emptyArrayArray# :: (# #) -> ArrayArray#
emptyArrayArray# _ = case emptyUnliftedArray of UnliftedArray ar -> ar
{-# NOINLINE emptyArrayArray# #-}
#endif
cloneUnliftedArray
:: UnliftedArray a
-> Int
-> Int
-> UnliftedArray a
cloneUnliftedArray src off len =
runUnliftedArray (thawUnliftedArray src off len)
{-# inline cloneUnliftedArray #-}
cloneMutableUnliftedArray
:: (PrimMonad m)
=> MutableUnliftedArray (PrimState m) a
-> Int
-> Int
-> m (MutableUnliftedArray (PrimState m) a)
cloneMutableUnliftedArray src off len = do
dst <- unsafeNewUnliftedArray len
copyMutableUnliftedArray dst 0 src off len
return dst
{-# inline cloneMutableUnliftedArray #-}
instance Eq (MutableUnliftedArray s a) where
(==) = sameMutableUnliftedArray
instance (Eq a, PrimUnlifted a) => Eq (UnliftedArray a) where
aa1 == aa2 = sizeofUnliftedArray aa1 == sizeofUnliftedArray aa2
&& loop (sizeofUnliftedArray aa1 - 1)
where
loop i
| i < 0 = True
| otherwise = indexUnliftedArray aa1 i == indexUnliftedArray aa2 i && loop (i-1)
instance (Ord a, PrimUnlifted a) => Ord (UnliftedArray a) where
compare a1 a2 = loop 0
where
mn = sizeofUnliftedArray a1 `min` sizeofUnliftedArray a2
loop i
| i < mn
, x1 <- indexUnliftedArray a1 i
, x2 <- indexUnliftedArray a2 i
= compare x1 x2 `mappend` loop (i+1)
| otherwise = compare (sizeofUnliftedArray a1) (sizeofUnliftedArray a2)
instance (Show a, PrimUnlifted a) => Show (UnliftedArray a) where
showsPrec p a = showParen (p > 10) $
showString "fromListN " . shows (sizeofUnliftedArray a) . showString " "
. shows (unliftedArrayToList a)
#if MIN_VERSION_base(4,9,0)
instance PrimUnlifted a => Semigroup (UnliftedArray a) where
(<>) = concatUnliftedArray
#endif
instance PrimUnlifted a => Monoid (UnliftedArray a) where
mempty = emptyUnliftedArray
#if !(MIN_VERSION_base(4,11,0))
mappend = concatUnliftedArray
#endif
emptyUnliftedArray :: UnliftedArray a
emptyUnliftedArray = runUnliftedArray (unsafeNewUnliftedArray 0)
{-# NOINLINE emptyUnliftedArray #-}
concatUnliftedArray :: UnliftedArray a -> UnliftedArray a -> UnliftedArray a
concatUnliftedArray x y = unsafeCreateUnliftedArray (sizeofUnliftedArray x + sizeofUnliftedArray y) $ \m -> do
copyUnliftedArray m 0 x 0 (sizeofUnliftedArray x)
copyUnliftedArray m (sizeofUnliftedArray x) y 0 (sizeofUnliftedArray y)
{-# INLINE foldrUnliftedArray #-}
foldrUnliftedArray :: forall a b. PrimUnlifted a => (a -> b -> b) -> b -> UnliftedArray a -> b
foldrUnliftedArray f z arr = go 0
where
!sz = sizeofUnliftedArray arr
go !i
| sz > i = f (indexUnliftedArray arr i) (go (i+1))
| otherwise = z
{-# INLINE foldrUnliftedArray' #-}
foldrUnliftedArray' :: forall a b. PrimUnlifted a => (a -> b -> b) -> b -> UnliftedArray a -> b
foldrUnliftedArray' f z0 arr = go (sizeofUnliftedArray arr - 1) z0
where
go !i !acc
| i < 0 = acc
| otherwise = go (i - 1) (f (indexUnliftedArray arr i) acc)
{-# INLINE foldlUnliftedArray #-}
foldlUnliftedArray :: forall a b. PrimUnlifted a => (b -> a -> b) -> b -> UnliftedArray a -> b
foldlUnliftedArray f z arr = go (sizeofUnliftedArray arr - 1)
where
go !i
| i < 0 = z
| otherwise = f (go (i - 1)) (indexUnliftedArray arr i)
{-# INLINE foldlUnliftedArray' #-}
foldlUnliftedArray' :: forall a b. PrimUnlifted a => (b -> a -> b) -> b -> UnliftedArray a -> b
foldlUnliftedArray' f z0 arr = go 0 z0
where
!sz = sizeofUnliftedArray arr
go !i !acc
| i < sz = go (i + 1) (f acc (indexUnliftedArray arr i))
| otherwise = acc
{-# INLINE mapUnliftedArray #-}
mapUnliftedArray :: (PrimUnlifted a, PrimUnlifted b)
=> (a -> b)
-> UnliftedArray a
-> UnliftedArray b
mapUnliftedArray f arr = unsafeCreateUnliftedArray sz $ \marr -> do
let go !ix = if ix < sz
then do
let b = f (indexUnliftedArray arr ix)
writeUnliftedArray marr ix b
go (ix + 1)
else return ()
go 0
where
!sz = sizeofUnliftedArray arr
{-# INLINE unliftedArrayToList #-}
unliftedArrayToList :: PrimUnlifted a => UnliftedArray a -> [a]
unliftedArrayToList xs = build (\c n -> foldrUnliftedArray c n xs)
unliftedArrayFromList :: PrimUnlifted a => [a] -> UnliftedArray a
unliftedArrayFromList xs = unliftedArrayFromListN (L.length xs) xs
unliftedArrayFromListN :: forall a. PrimUnlifted a => Int -> [a] -> UnliftedArray a
unliftedArrayFromListN len vs = unsafeCreateUnliftedArray len run where
run :: forall s. MutableUnliftedArray s a -> ST s ()
run arr = do
let go :: [a] -> Int -> ST s ()
go [] !ix = if ix == len
then return ()
else die "unliftedArrayFromListN" "list length less than specified size"
go (a : as) !ix = if ix < len
then do
writeUnliftedArray arr ix a
go as (ix + 1)
else die "unliftedArrayFromListN" "list length greater than specified size"
go vs 0
#if MIN_VERSION_base(4,7,0)
instance PrimUnlifted a => E.IsList (UnliftedArray a) where
type Item (UnliftedArray a) = a
fromList = unliftedArrayFromList
fromListN = unliftedArrayFromListN
toList = unliftedArrayToList
#endif