{-# LANGUAGE CPP, MagicHash, UnboxedTuples, DeriveDataTypeable, BangPatterns #-}
{-# LANGUAGE RankNTypes #-}
{-# LANGUAGE TypeFamilies #-}
module Data.Primitive.Array (
Array(..), MutableArray(..),
newArray, readArray, writeArray, indexArray, indexArrayM, indexArray##,
freezeArray, thawArray, runArray, createArray,
unsafeFreezeArray, unsafeThawArray, sameMutableArray,
copyArray, copyMutableArray,
cloneArray, cloneMutableArray,
sizeofArray, sizeofMutableArray,
emptyArray,
fromListN, fromList,
arrayFromListN, arrayFromList,
mapArray',
traverseArrayP
) where
import Control.DeepSeq
import Control.Monad.Primitive
import GHC.Exts hiding (toList)
import qualified GHC.Exts as Exts
import Data.Typeable ( Typeable )
import Data.Data
(Data(..), DataType, mkDataType, mkNoRepType, Constr, mkConstr, Fixity(..), constrIndex)
import Control.Monad.ST (ST, runST)
import Control.Applicative
import Control.Monad (MonadPlus(..), when, liftM2)
import qualified Control.Monad.Fail as Fail
import Control.Monad.Fix
import qualified Data.Foldable as Foldable
import Control.Monad.Zip
import Data.Foldable (Foldable(..), toList)
#if MIN_VERSION_base(4,9,0)
import qualified GHC.ST as GHCST
import qualified Data.Foldable as F
import Data.Semigroup
#endif
import Data.Functor.Identity
#if MIN_VERSION_base(4,10,0)
import GHC.Exts (runRW#)
#elif MIN_VERSION_base(4,9,0)
import GHC.Base (runRW#)
#endif
import Text.Read (Read (..), parens, prec)
import Text.ParserCombinators.ReadPrec (ReadPrec)
import qualified Text.ParserCombinators.ReadPrec as RdPrc
import Text.ParserCombinators.ReadP
import Data.Functor.Classes (Eq1(..), Ord1(..), Show1(..), Read1(..))
data Array a = Array
{ array# :: Array# a }
deriving ( Typeable )
#if MIN_VERSION_deepseq(1,4,3)
instance NFData1 Array where
liftRnf r = Foldable.foldl' (\_ -> r) ()
#endif
instance NFData a => NFData (Array a) where
rnf = Foldable.foldl' (\_ -> rnf) ()
data MutableArray s a = MutableArray
{ marray# :: MutableArray# s a }
deriving ( Typeable )
sizeofArray :: Array a -> Int
sizeofArray a = I# (sizeofArray# (array# a))
{-# INLINE sizeofArray #-}
sizeofMutableArray :: MutableArray s a -> Int
sizeofMutableArray a = I# (sizeofMutableArray# (marray# a))
{-# INLINE sizeofMutableArray #-}
newArray :: PrimMonad m => Int -> a -> m (MutableArray (PrimState m) a)
{-# INLINE newArray #-}
newArray (I# n#) x = primitive
(\s# -> case newArray# n# x s# of
(# s'#, arr# #) ->
let ma = MutableArray arr#
in (# s'# , ma #))
readArray :: PrimMonad m => MutableArray (PrimState m) a -> Int -> m a
{-# INLINE readArray #-}
readArray arr (I# i#) = primitive (readArray# (marray# arr) i#)
writeArray :: PrimMonad m => MutableArray (PrimState m) a -> Int -> a -> m ()
{-# INLINE writeArray #-}
writeArray arr (I# i#) x = primitive_ (writeArray# (marray# arr) i# x)
indexArray :: Array a -> Int -> a
{-# INLINE indexArray #-}
indexArray arr (I# i#) = case indexArray# (array# arr) i# of (# x #) -> x
indexArray## :: Array a -> Int -> (# a #)
indexArray## arr (I# i) = indexArray# (array# arr) i
{-# INLINE indexArray## #-}
indexArrayM :: Monad m => Array a -> Int -> m a
{-# INLINE indexArrayM #-}
indexArrayM arr (I# i#)
= case indexArray# (array# arr) i# of (# x #) -> return x
freezeArray
:: PrimMonad m
=> MutableArray (PrimState m) a
-> Int
-> Int
-> m (Array a)
{-# INLINE freezeArray #-}
freezeArray (MutableArray ma#) (I# off#) (I# len#) =
primitive $ \s -> case freezeArray# ma# off# len# s of
(# s', a# #) -> (# s', Array a# #)
unsafeFreezeArray :: PrimMonad m => MutableArray (PrimState m) a -> m (Array a)
{-# INLINE unsafeFreezeArray #-}
unsafeFreezeArray arr
= primitive (\s# -> case unsafeFreezeArray# (marray# arr) s# of
(# s'#, arr'# #) ->
let a = Array arr'#
in (# s'#, a #))
thawArray
:: PrimMonad m
=> Array a
-> Int
-> Int
-> m (MutableArray (PrimState m) a)
{-# INLINE thawArray #-}
thawArray (Array a#) (I# off#) (I# len#) =
primitive $ \s -> case thawArray# a# off# len# s of
(# s', ma# #) -> (# s', MutableArray ma# #)
unsafeThawArray :: PrimMonad m => Array a -> m (MutableArray (PrimState m) a)
{-# INLINE unsafeThawArray #-}
unsafeThawArray a
= primitive (\s# -> case unsafeThawArray# (array# a) s# of
(# s'#, arr'# #) ->
let ma = MutableArray arr'#
in (# s'#, ma #))
sameMutableArray :: MutableArray s a -> MutableArray s a -> Bool
{-# INLINE sameMutableArray #-}
sameMutableArray arr brr
= isTrue# (sameMutableArray# (marray# arr) (marray# brr))
copyArray :: PrimMonad m
=> MutableArray (PrimState m) a
-> Int
-> Array a
-> Int
-> Int
-> m ()
{-# INLINE copyArray #-}
copyArray (MutableArray dst#) (I# doff#) (Array src#) (I# soff#) (I# len#)
= primitive_ (copyArray# src# soff# dst# doff# len#)
copyMutableArray :: PrimMonad m
=> MutableArray (PrimState m) a
-> Int
-> MutableArray (PrimState m) a
-> Int
-> Int
-> m ()
{-# INLINE copyMutableArray #-}
copyMutableArray (MutableArray dst#) (I# doff#)
(MutableArray src#) (I# soff#) (I# len#)
= primitive_ (copyMutableArray# src# soff# dst# doff# len#)
cloneArray :: Array a
-> Int
-> Int
-> Array a
{-# INLINE cloneArray #-}
cloneArray (Array arr#) (I# off#) (I# len#)
= case cloneArray# arr# off# len# of arr'# -> Array arr'#
cloneMutableArray :: PrimMonad m
=> MutableArray (PrimState m) a
-> Int
-> Int
-> m (MutableArray (PrimState m) a)
{-# INLINE cloneMutableArray #-}
cloneMutableArray (MutableArray arr#) (I# off#) (I# len#) = primitive
(\s# -> case cloneMutableArray# arr# off# len# s# of
(# s'#, arr'# #) -> (# s'#, MutableArray arr'# #))
emptyArray :: Array a
emptyArray =
runST $ newArray 0 (die "emptyArray" "impossible") >>= unsafeFreezeArray
{-# NOINLINE emptyArray #-}
runArray
:: (forall s. ST s (MutableArray s a))
-> Array a
#if !MIN_VERSION_base(4,9,0)
runArray m = runST $ m >>= unsafeFreezeArray
#else /* Below, runRW# is available. */
runArray m = Array (runArray# m)
runArray#
:: (forall s. ST s (MutableArray s a))
-> Array# a
runArray# m = case runRW# $ \s ->
case unST m s of { (# s', MutableArray mary# #) ->
unsafeFreezeArray# mary# s'} of (# _, ary# #) -> ary#
unST :: ST s a -> State# s -> (# State# s, a #)
unST (GHCST.ST f) = f
emptyArray# :: (# #) -> Array# a
emptyArray# _ = case emptyArray of Array ar -> ar
{-# NOINLINE emptyArray# #-}
#endif
createArray
:: Int
-> a
-> (forall s. MutableArray s a -> ST s ())
-> Array a
#if !MIN_VERSION_base(4,9,0)
createArray 0 _ _ = emptyArray
#else
createArray 0 _ _ = Array (emptyArray# (# #))
#endif
createArray n x f = runArray $ do
mary <- newArray n x
f mary
pure mary
die :: String -> String -> a
die fun problem = error $ "Data.Primitive.Array." ++ fun ++ ": " ++ problem
arrayLiftEq :: (a -> b -> Bool) -> Array a -> Array b -> Bool
arrayLiftEq p a1 a2 = sizeofArray a1 == sizeofArray a2 && loop (sizeofArray a1 - 1)
where loop i | i < 0 = True
| (# x1 #) <- indexArray## a1 i
, (# x2 #) <- indexArray## a2 i
, otherwise = p x1 x2 && loop (i - 1)
instance Eq a => Eq (Array a) where
a1 == a2 = arrayLiftEq (==) a1 a2
instance Eq1 Array where
#if MIN_VERSION_base(4,9,0) || MIN_VERSION_transformers(0,5,0)
liftEq = arrayLiftEq
#else
eq1 = arrayLiftEq (==)
#endif
instance Eq (MutableArray s a) where
ma1 == ma2 = isTrue# (sameMutableArray# (marray# ma1) (marray# ma2))
arrayLiftCompare :: (a -> b -> Ordering) -> Array a -> Array b -> Ordering
arrayLiftCompare elemCompare a1 a2 = loop 0
where
mn = sizeofArray a1 `min` sizeofArray a2
loop i
| i < mn
, (# x1 #) <- indexArray## a1 i
, (# x2 #) <- indexArray## a2 i
= elemCompare x1 x2 `mappend` loop (i + 1)
| otherwise = compare (sizeofArray a1) (sizeofArray a2)
instance Ord a => Ord (Array a) where
compare a1 a2 = arrayLiftCompare compare a1 a2
instance Ord1 Array where
#if MIN_VERSION_base(4,9,0) || MIN_VERSION_transformers(0,5,0)
liftCompare = arrayLiftCompare
#else
compare1 = arrayLiftCompare compare
#endif
instance Foldable Array where
foldr f = \z !ary ->
let
!sz = sizeofArray ary
go i
| i == sz = z
| (# x #) <- indexArray## ary i
= f x (go (i + 1))
in go 0
{-# INLINE foldr #-}
foldl f = \z !ary ->
let
go i
| i < 0 = z
| (# x #) <- indexArray## ary i
= f (go (i - 1)) x
in go (sizeofArray ary - 1)
{-# INLINE foldl #-}
foldr1 f = \ !ary ->
let
!sz = sizeofArray ary - 1
go i =
case indexArray## ary i of
(# x #) | i == sz -> x
| otherwise -> f x (go (i + 1))
in if sz < 0
then die "foldr1" "empty array"
else go 0
{-# INLINE foldr1 #-}
foldl1 f = \ !ary ->
let
!sz = sizeofArray ary - 1
go i =
case indexArray## ary i of
(# x #) | i == 0 -> x
| otherwise -> f (go (i - 1)) x
in if sz < 0
then die "foldl1" "empty array"
else go sz
{-# INLINE foldl1 #-}
foldr' f = \z !ary ->
let
go i !acc
| i == -1 = acc
| (# x #) <- indexArray## ary i
= go (i - 1) (f x acc)
in go (sizeofArray ary - 1) z
{-# INLINE foldr' #-}
foldl' f = \z !ary ->
let
!sz = sizeofArray ary
go i !acc
| i == sz = acc
| (# x #) <- indexArray## ary i
= go (i + 1) (f acc x)
in go 0 z
{-# INLINE foldl' #-}
null a = sizeofArray a == 0
{-# INLINE null #-}
length = sizeofArray
{-# INLINE length #-}
maximum ary | sz == 0 = die "maximum" "empty array"
| (# frst #) <- indexArray## ary 0
= go 1 frst
where
sz = sizeofArray ary
go i !e
| i == sz = e
| (# x #) <- indexArray## ary i
= go (i + 1) (max e x)
{-# INLINE maximum #-}
minimum ary | sz == 0 = die "minimum" "empty array"
| (# frst #) <- indexArray## ary 0
= go 1 frst
where sz = sizeofArray ary
go i !e
| i == sz = e
| (# x #) <- indexArray## ary i
= go (i + 1) (min e x)
{-# INLINE minimum #-}
sum = foldl' (+) 0
{-# INLINE sum #-}
product = foldl' (*) 1
{-# INLINE product #-}
newtype STA a = STA { _runSTA :: forall s. MutableArray# s a -> ST s (Array a) }
runSTA :: Int -> STA a -> Array a
runSTA !sz = \ (STA m) -> runST $ newArray_ sz >>= \ ar -> m (marray# ar)
{-# INLINE runSTA #-}
newArray_ :: Int -> ST s (MutableArray s a)
newArray_ !n = newArray n badTraverseValue
badTraverseValue :: a
badTraverseValue = die "traverse" "bad indexing"
{-# NOINLINE badTraverseValue #-}
instance Traversable Array where
traverse f = traverseArray f
{-# INLINE traverse #-}
traverseArray
:: Applicative f
=> (a -> f b)
-> Array a
-> f (Array b)
traverseArray f = \ !ary ->
let
!len = sizeofArray ary
go !i
| i == len = pure $ STA $ \mary -> unsafeFreezeArray (MutableArray mary)
| (# x #) <- indexArray## ary i
= liftA2 (\b (STA m) -> STA $ \mary ->
writeArray (MutableArray mary) i b >> m mary)
(f x) (go (i + 1))
in if len == 0
then pure emptyArray
else runSTA len <$> go 0
{-# INLINE [1] traverseArray #-}
{-# RULES
"traverse/ST" forall (f :: a -> ST s b). traverseArray f =
traverseArrayP f
"traverse/IO" forall (f :: a -> IO b). traverseArray f =
traverseArrayP f
"traverse/Id" forall (f :: a -> Identity b). traverseArray f =
(coerce :: (Array a -> Array (Identity b))
-> Array a -> Identity (Array b)) (fmap f)
#-}
traverseArrayP
:: PrimMonad m
=> (a -> m b)
-> Array a
-> m (Array b)
traverseArrayP f = \ !ary ->
let
!sz = sizeofArray ary
go !i !mary
| i == sz
= unsafeFreezeArray mary
| otherwise
= do
a <- indexArrayM ary i
b <- f a
writeArray mary i b
go (i + 1) mary
in do
mary <- newArray sz badTraverseValue
go 0 mary
{-# INLINE traverseArrayP #-}
mapArray' :: (a -> b) -> Array a -> Array b
mapArray' f a =
createArray (sizeofArray a) (die "mapArray'" "impossible") $ \mb ->
let go i | i == sizeofArray a
= return ()
| otherwise
= do x <- indexArrayM a i
let !y = f x
writeArray mb i y >> go (i + 1)
in go 0
{-# INLINE mapArray' #-}
arrayFromListN :: Int -> [a] -> Array a
arrayFromListN n l =
createArray n (die "fromListN" "uninitialized element") $ \sma ->
let go !ix [] = if ix == n
then return ()
else die "fromListN" "list length less than specified size"
go !ix (x : xs) = if ix < n
then do
writeArray sma ix x
go (ix+1) xs
else die "fromListN" "list length greater than specified size"
in go 0 l
arrayFromList :: [a] -> Array a
arrayFromList l = arrayFromListN (length l) l
instance Exts.IsList (Array a) where
type Item (Array a) = a
fromListN = arrayFromListN
fromList = arrayFromList
toList = toList
instance Functor Array where
fmap f a =
createArray (sizeofArray a) (die "fmap" "impossible") $ \mb ->
let go i | i == sizeofArray a
= return ()
| otherwise
= do x <- indexArrayM a i
writeArray mb i (f x) >> go (i + 1)
in go 0
e <$ a = createArray (sizeofArray a) e (\ !_ -> pure ())
instance Applicative Array where
pure x = runArray $ newArray 1 x
ab <*> a = createArray (szab * sza) (die "<*>" "impossible") $ \mb ->
let go1 i = when (i < szab) $
do
f <- indexArrayM ab i
go2 (i * sza) f 0
go1 (i + 1)
go2 off f j = when (j < sza) $
do
x <- indexArrayM a j
writeArray mb (off + j) (f x)
go2 off f (j + 1)
in go1 0
where szab = sizeofArray ab; sza = sizeofArray a
a *> b = createArray (sza * szb) (die "*>" "impossible") $ \mb ->
let go i | i < sza = copyArray mb (i * szb) b 0 szb *> go (i + 1)
| otherwise = return ()
in go 0
where sza = sizeofArray a; szb = sizeofArray b
a <* b = createArray (sza * szb) (die "<*" "impossible") $ \ma ->
let fill off i e | i < szb = writeArray ma (off + i) e >> fill off (i + 1) e
| otherwise = return ()
go i | i < sza
= do x <- indexArrayM a i
fill (i * szb) 0 x >> go (i + 1)
| otherwise = return ()
in go 0
where sza = sizeofArray a; szb = sizeofArray b
instance Alternative Array where
empty = emptyArray
a1 <|> a2 = createArray (sza1 + sza2) (die "<|>" "impossible") $ \ma ->
copyArray ma 0 a1 0 sza1 >> copyArray ma sza1 a2 0 sza2
where sza1 = sizeofArray a1; sza2 = sizeofArray a2
some a | sizeofArray a == 0 = emptyArray
| otherwise = die "some" "infinite arrays are not well defined"
many a | sizeofArray a == 0 = pure []
| otherwise = die "many" "infinite arrays are not well defined"
data ArrayStack a
= PushArray !(Array a) !(ArrayStack a)
| EmptyStack
instance Monad Array where
return = pure
(>>) = (*>)
ary >>= f = collect 0 EmptyStack (la - 1)
where
la = sizeofArray ary
collect sz stk i
| i < 0 = createArray sz (die ">>=" "impossible") $ fill 0 stk
| (# x #) <- indexArray## ary i
, let sb = f x
lsb = sizeofArray sb
= if lsb == 0
then collect sz stk (i - 1)
else collect (sz + lsb) (PushArray sb stk) (i - 1)
fill _ EmptyStack _ = return ()
fill off (PushArray sb sbs) smb
| let lsb = sizeofArray sb
= copyArray smb off sb 0 lsb
*> fill (off + lsb) sbs smb
#if !(MIN_VERSION_base(4,13,0))
fail = Fail.fail
#endif
instance Fail.MonadFail Array where
fail _ = empty
instance MonadPlus Array where
mzero = empty
mplus = (<|>)
zipW :: String -> (a -> b -> c) -> Array a -> Array b -> Array c
zipW s f aa ab = createArray mn (die s "impossible") $ \mc ->
let go i | i < mn
= do
x <- indexArrayM aa i
y <- indexArrayM ab i
writeArray mc i (f x y)
go (i + 1)
| otherwise = return ()
in go 0
where mn = sizeofArray aa `min` sizeofArray ab
{-# INLINE zipW #-}
instance MonadZip Array where
mzip aa ab = zipW "mzip" (,) aa ab
mzipWith f aa ab = zipW "mzipWith" f aa ab
munzip aab = runST $ do
let sz = sizeofArray aab
ma <- newArray sz (die "munzip" "impossible")
mb <- newArray sz (die "munzip" "impossible")
let go i | i < sz = do
(a, b) <- indexArrayM aab i
writeArray ma i a
writeArray mb i b
go (i + 1)
go _ = return ()
go 0
(,) <$> unsafeFreezeArray ma <*> unsafeFreezeArray mb
instance MonadFix Array where
mfix f = createArray (sizeofArray (f err))
(die "mfix" "impossible") $ flip fix 0 $
\r !i !mary -> when (i < sz) $ do
writeArray mary i (fix (\xi -> f xi `indexArray` i))
r (i + 1) mary
where
sz = sizeofArray (f err)
err = error "mfix for Data.Primitive.Array applied to strict function."
#if MIN_VERSION_base(4,9,0)
instance Semigroup (Array a) where
(<>) = (<|>)
sconcat = mconcat . F.toList
#endif
instance Monoid (Array a) where
mempty = empty
#if !(MIN_VERSION_base(4,11,0))
mappend = (<|>)
#endif
mconcat l = createArray sz (die "mconcat" "impossible") $ \ma ->
let go !_ [ ] = return ()
go off (a:as) =
copyArray ma off a 0 (sizeofArray a) >> go (off + sizeofArray a) as
in go 0 l
where sz = sum . fmap sizeofArray $ l
arrayLiftShowsPrec :: (Int -> a -> ShowS) -> ([a] -> ShowS) -> Int -> Array a -> ShowS
arrayLiftShowsPrec elemShowsPrec elemListShowsPrec p a = showParen (p > 10) $
showString "fromListN " . shows (sizeofArray a) . showString " "
. listLiftShowsPrec elemShowsPrec elemListShowsPrec 11 (toList a)
listLiftShowsPrec :: (Int -> a -> ShowS) -> ([a] -> ShowS) -> Int -> [a] -> ShowS
listLiftShowsPrec _ sl _ = sl
instance Show a => Show (Array a) where
showsPrec p a = arrayLiftShowsPrec showsPrec showList p a
instance Show1 Array where
#if MIN_VERSION_base(4,9,0) || MIN_VERSION_transformers(0,5,0)
liftShowsPrec = arrayLiftShowsPrec
#else
showsPrec1 = arrayLiftShowsPrec showsPrec showList
#endif
instance Read a => Read (Array a) where
readPrec = arrayLiftReadPrec readPrec readListPrec
instance Read1 Array where
#if MIN_VERSION_base(4,10,0)
liftReadPrec = arrayLiftReadPrec
#elif MIN_VERSION_base(4,9,0) || MIN_VERSION_transformers(0,5,0)
liftReadsPrec = arrayLiftReadsPrec
#else
readsPrec1 = arrayLiftReadsPrec readsPrec readList
#endif
arrayLiftReadPrec :: ReadPrec a -> ReadPrec [a] -> ReadPrec (Array a)
arrayLiftReadPrec _ read_list = parens $ prec app_prec $ RdPrc.lift skipSpaces >>
((fromList <$> read_list) RdPrc.+++
do
tag <- RdPrc.lift lexTag
case tag of
FromListTag -> fromList <$> read_list
FromListNTag -> liftM2 fromListN readPrec read_list)
where
app_prec = 10
data Tag = FromListTag | FromListNTag
lexTag :: ReadP Tag
lexTag = do
_ <- string "fromList"
s <- look
case s of
'N':c:_
| '0' <= c && c <= '9'
-> fail ""
| otherwise -> FromListNTag <$ get
_ -> return FromListTag
#if !MIN_VERSION_base(4,10,0)
arrayLiftReadsPrec :: (Int -> ReadS a) -> ReadS [a] -> Int -> ReadS (Array a)
arrayLiftReadsPrec reads_prec list_reads_prec = RdPrc.readPrec_to_S $
arrayLiftReadPrec (RdPrc.readS_to_Prec reads_prec) (RdPrc.readS_to_Prec (const list_reads_prec))
#endif
arrayDataType :: DataType
arrayDataType = mkDataType "Data.Primitive.Array.Array" [fromListConstr]
fromListConstr :: Constr
fromListConstr = mkConstr arrayDataType "fromList" [] Prefix
instance Data a => Data (Array a) where
toConstr _ = fromListConstr
dataTypeOf _ = arrayDataType
gunfold k z c = case constrIndex c of
1 -> k (z fromList)
_ -> error "gunfold"
gfoldl f z m = z fromList `f` toList m
instance (Typeable s, Typeable a) => Data (MutableArray s a) where
toConstr _ = error "toConstr"
gunfold _ _ = error "gunfold"
dataTypeOf _ = mkNoRepType "Data.Primitive.Array.MutableArray"