module Foundation.Array.Boxed
( Array
, MArray
, empty
, length
, copy
, copyAt
, unsafeCopyAtRO
, thaw
, new
, unsafeFreeze
, unsafeThaw
, freeze
, unsafeWrite
, unsafeRead
, unsafeIndex
, write
, read
, index
, singleton
, null
, take
, drop
, splitAt
, revTake
, revDrop
, revSplitAt
, splitOn
, sub
, intersperse
, span
, break
, cons
, snoc
, uncons
, unsnoc
, sortBy
, filter
, reverse
, find
, foldl'
, foldr
, foldl
, builderAppend
, builderBuild
) where
import GHC.Prim
import GHC.Types
import GHC.ST
import Foundation.Numerical
import Foundation.Internal.Base
import Foundation.Internal.MonadTrans
import Foundation.Internal.Types
import Foundation.Primitive.Types
import Foundation.Primitive.Monad
import Foundation.Array.Common
import Foundation.Boot.Builder
import qualified Foundation.Boot.List as List
import qualified Prelude
data Array a = Array !(Offset a)
!(Size a)
(Array# a)
deriving (Typeable)
instance Data ty => Data (Array ty) where
dataTypeOf _ = arrayType
toConstr _ = error "toConstr"
gunfold _ _ = error "gunfold"
arrayType :: DataType
arrayType = mkNoRepType "Foundation.Array"
data MArray a st = MArray !(Offset a)
!(Size a)
(MutableArray# st a)
deriving (Typeable)
instance Functor Array where
fmap = map
instance Monoid (Array a) where
mempty = empty
mappend = append
mconcat = concat
instance Show a => Show (Array a) where
show v = show (toList v)
instance Eq a => Eq (Array a) where
(==) = equal
instance Ord a => Ord (Array a) where
compare = vCompare
instance IsList (Array ty) where
type Item (Array ty) = ty
fromList = vFromList
toList = vToList
mutableLength :: MArray ty st -> Int
mutableLength (MArray _ (Size len) _) = len
index :: Array ty -> Int -> ty
index array n
| n < 0 || n >= len = throw (OutOfBound OOB_Index n len)
| otherwise = unsafeIndex array n
where len = length array
unsafeIndex :: Array ty -> Int -> ty
unsafeIndex (Array start _ a) ofs = primArrayIndex a (start+Offset ofs)
read :: PrimMonad prim => MArray ty (PrimState prim) -> Int -> prim ty
read array n
| n < 0 || n >= len = primThrow (OutOfBound OOB_Read n len)
| otherwise = unsafeRead array n
where len = mutableLength array
unsafeRead :: PrimMonad prim => MArray ty (PrimState prim) -> Int -> prim ty
unsafeRead (MArray start _ ma) i = primMutableArrayRead ma (start + Offset i)
write :: PrimMonad prim => MArray ty (PrimState prim) -> Int -> ty -> prim ()
write array n val
| n < 0 || n >= len = primThrow (OutOfBound OOB_Write n len)
| otherwise = unsafeWrite array n val
where len = mutableLength array
unsafeWrite :: PrimMonad prim => MArray ty (PrimState prim) -> Int -> ty -> prim ()
unsafeWrite (MArray start _ ma) ofs v =
primMutableArrayWrite ma (start + Offset ofs) v
unsafeFreeze :: PrimMonad prim => MArray ty (PrimState prim) -> prim (Array ty)
unsafeFreeze (MArray ofs sz ma) = primitive $ \s1 ->
case unsafeFreezeArray# ma s1 of
(# s2, a #) -> (# s2, Array ofs sz a #)
unsafeThaw :: PrimMonad prim => Array ty -> prim (MArray ty (PrimState prim))
unsafeThaw (Array ofs sz a) = primitive $ \st -> (# st, MArray ofs sz (unsafeCoerce# a) #)
thaw :: PrimMonad prim => Array ty -> prim (MArray ty (PrimState prim))
thaw array = do
m <- new (lengthSize array)
unsafeCopyAtRO m (Offset 0) array (Offset 0) (lengthSize array)
return m
freeze :: PrimMonad prim => MArray ty (PrimState prim) -> prim (Array ty)
freeze marray = do
m <- new sz
copyAt m (Offset 0) marray (Offset 0) sz
unsafeFreeze m
where
sz = Size $ mutableLength marray
copy :: Array ty -> Array ty
copy a = runST (unsafeThaw a >>= freeze)
copyAt :: PrimMonad prim
=> MArray ty (PrimState prim)
-> Offset ty
-> MArray ty (PrimState prim)
-> Offset ty
-> Size ty
-> prim ()
copyAt dst od src os n = loop od os
where !endIndex = os `offsetPlusE` n
loop (Offset d) s@(Offset i)
| s == endIndex = return ()
| otherwise = unsafeRead src i >>= unsafeWrite dst d >> loop (Offset $ d+1) (Offset $ i+1)
unsafeCopyAtRO :: PrimMonad prim
=> MArray ty (PrimState prim)
-> Offset ty
-> Array ty
-> Offset ty
-> Size ty
-> prim ()
unsafeCopyAtRO (MArray (Offset (I# dstart)) _ da) (Offset (I# dofs))
(Array (Offset (I# sstart)) _ sa) (Offset (I# sofs))
(Size (I# n)) =
primitive $ \st ->
(# copyArray# sa (sstart +# sofs) da (dstart +# dofs) n st, () #)
unsafeCopyFrom :: Array ty
-> Size ty
-> (Array ty -> Offset ty -> MArray ty s -> ST s ())
-> ST s (Array ty)
unsafeCopyFrom v' newLen f = new newLen >>= fill (Offset 0) f >>= unsafeFreeze
where len = lengthSize v'
endIdx = Offset 0 `offsetPlusE` len
fill i f' r'
| i == endIdx = return r'
| otherwise = do f' v' i r'
fill (i + Offset 1) f' r'
new :: PrimMonad prim => Size ty -> prim (MArray ty (PrimState prim))
new sz@(Size (I# n)) = primitive $ \s1 ->
case newArray# n (error "vector: internal error uninitialized vector") s1 of
(# s2, ma #) -> (# s2, MArray (Offset 0) sz ma #)
create :: Int
-> (Int -> ty)
-> Array ty
create n initializer = runST (new (Size n) >>= iter initializer)
where
iter :: PrimMonad prim => (Int -> ty) -> MArray ty (PrimState prim) -> prim (Array ty)
iter f ma = loop (Offset 0)
where
!end = Offset 0 `offsetPlusE` Size n
loop s@(Offset i)
| s == end = unsafeFreeze ma
| otherwise = unsafeWrite ma i (f i) >> loop (Offset $ i+1)
equal :: Eq a => Array a -> Array a -> Bool
equal a b = (len == length b) && eachEqual 0
where
len = length a
eachEqual !i
| i == len = True
| unsafeIndex a i /= unsafeIndex b i = False
| otherwise = eachEqual (i+1)
vCompare :: Ord a => Array a -> Array a -> Ordering
vCompare a b = loop 0
where
!la = length a
!lb = length b
loop n
| n == la = if la == lb then EQ else LT
| n == lb = GT
| otherwise =
case unsafeIndex a n `compare` unsafeIndex b n of
EQ -> loop (n+1)
r -> r
empty :: Array a
empty = runST $ onNewArray 0 (\_ s -> s)
length :: Array a -> Int
length (Array _ (Size len) _) = len
lengthSize :: Array a -> Size a
lengthSize (Array _ sz _) = sz
vFromList :: [a] -> Array a
vFromList l = runST (new len >>= loop 0 l)
where
len = Size $ List.length l
loop _ [] ma = unsafeFreeze ma
loop i (x:xs) ma = unsafeWrite ma i x >> loop (i+1) xs ma
vToList :: Array a -> [a]
vToList v = fmap (unsafeIndex v) [0..(length v 1)]
append :: Array ty -> Array ty -> Array ty
append a b = 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 = lengthSize a
lb = lengthSize b
concat :: [Array ty] -> Array ty
concat l = runST $ do
r <- new (Size $ Prelude.sum $ fmap length l)
loop r (Offset 0) l
unsafeFreeze r
where loop _ _ [] = return ()
loop r i (x:xs) = do
mx <- unsafeThaw x
copyAt r i mx (Offset 0) lx
loop r (i `offsetPlusE` lx) xs
where lx = lengthSize x
onNewArray :: PrimMonad m
=> Int
-> (MutableArray# (PrimState m) a -> State# (PrimState m) -> State# (PrimState m))
-> m (Array a)
onNewArray len@(I# len#) f = primitive $ \st -> do
case newArray# len# (error "onArray") st of { (# st2, mv #) ->
case f mv st2 of { st3 ->
case unsafeFreezeArray# mv st3 of { (# st4, a #) ->
(# st4, Array (Offset 0) (Size len) a #) }}}
null :: Array ty -> Bool
null = (==) 0 . length
take :: Int -> Array ty -> Array ty
take nbElems a@(Array start len arr)
| nbElems <= 0 = empty
| n == len = a
| otherwise = Array start n arr
where
n = min (Size nbElems) len
drop :: Int -> Array ty -> Array ty
drop nbElems a@(Array start len arr)
| nbElems <= 0 = a
| n == len = empty
| otherwise = Array (start `offsetPlusE` n) (len n) arr
where
n = min (Size nbElems) len
splitAt :: Int -> Array ty -> (Array ty, Array ty)
splitAt nbElems a@(Array start len arr)
| nbElems <= 0 = (empty, a)
| n == len = (a, empty)
| otherwise =
(Array start n arr, Array (start `offsetPlusE` n) (len n) arr)
where
n = min (Size nbElems) len
revTake :: Int -> Array ty -> Array ty
revTake nbElems v = drop (length v nbElems) v
revDrop :: Int -> Array ty -> Array ty
revDrop nbElems v = take (length v nbElems) v
revSplitAt :: Int -> Array ty -> (Array ty, Array ty)
revSplitAt n v = (drop idx v, take idx v)
where idx = length v n
splitOn :: (ty -> Bool) -> Array ty -> [Array ty]
splitOn predicate vec
| len == Size 0 = [mempty]
| otherwise = loop (Offset 0) (Offset 0)
where
!len = lengthSize vec
!endIdx = Offset 0 `offsetPlusE` len
loop prevIdx idx@(Offset i)
| idx == endIdx = [sub vec prevIdx idx]
| otherwise =
let e = unsafeIndex vec i
idx' = idx + Offset 1
in if predicate e
then sub vec prevIdx idx : loop idx' idx'
else loop prevIdx idx'
sub :: Array ty -> Offset ty -> Offset ty -> Array ty
sub (Array start len a) startIdx expectedEndIdx
| startIdx == endIdx = empty
| otherwise = Array (start + startIdx) newLen a
where
newLen = endIdx startIdx
endIdx = min expectedEndIdx (sizeAsOffset len)
break :: (ty -> Bool) -> Array ty -> (Array ty, Array ty)
break predicate v = findBreak 0
where
findBreak i
| i == length v = (v, empty)
| otherwise =
if predicate (unsafeIndex v i)
then splitAt i v
else findBreak (i+1)
intersperse :: ty -> Array ty -> Array ty
intersperse sep v
| len <= Size 1 = v
| otherwise = runST $ unsafeCopyFrom v ((len + len) Size 1) (go (Offset 0 `offsetPlusE` (len Size 1)) sep)
where len = lengthSize v
go :: Offset ty -> ty -> Array ty -> Offset ty -> MArray ty s -> ST s ()
go endI sep' oldV oldI@(Offset oi) newV
| oldI == endI = unsafeWrite newV dst e
| otherwise = do
unsafeWrite newV dst e
unsafeWrite newV (dst + 1) sep'
where
e = unsafeIndex oldV oi
(Offset dst) = oldI + oldI
span :: (ty -> Bool) -> Array ty -> (Array ty, Array ty)
span p = break (not . p)
map :: (a -> b) -> Array a -> Array b
map f a = create (length a) (\i -> f $ unsafeIndex a i)
singleton :: ty -> Array ty
singleton e = runST $ do
a <- new 1
unsafeWrite a 0 e
unsafeFreeze a
cons :: ty -> Array ty -> Array ty
cons e vec
| len == Size 0 = singleton e
| otherwise = runST $ do
mv <- new (len + Size 1)
unsafeWrite mv 0 e
unsafeCopyAtRO mv (Offset 1) vec (Offset 0) len
unsafeFreeze mv
where
!len = lengthSize vec
snoc :: Array ty -> ty -> Array ty
snoc vec e
| len == Size 0 = singleton e
| otherwise = runST $ do
mv <- new (len + Size 1)
unsafeCopyAtRO mv (Offset 0) vec (Offset 0) len
unsafeWrite mv lastI e
unsafeFreeze mv
where
!len@(Size lastI) = lengthSize vec
uncons :: Array ty -> Maybe (ty, Array ty)
uncons vec
| len == 0 = Nothing
| otherwise = Just (unsafeIndex vec 0, drop 1 vec)
where
!len = length vec
unsnoc :: Array ty -> Maybe (Array ty, ty)
unsnoc vec
| len == 0 = Nothing
| otherwise = Just (take (len 1) vec, unsafeIndex vec (len1))
where
!len = length vec
find :: (ty -> Bool) -> Array ty -> Maybe ty
find predicate vec = loop 0
where
!len = length vec
loop i
| i == len = Nothing
| otherwise =
let e = unsafeIndex vec i
in if predicate e then Just e else loop (i+1)
sortBy :: (ty -> ty -> Ordering) -> Array ty -> Array ty
sortBy xford vec = runST (thaw vec >>= doSort xford)
where
len = length vec
doSort :: PrimMonad prim => (ty -> ty -> Ordering) -> MArray ty (PrimState prim) -> prim (Array ty)
doSort ford ma = qsort 0 (len 1) >> unsafeFreeze ma
where
qsort lo hi
| lo >= hi = return ()
| otherwise = do
p <- partition lo hi
qsort lo (p1)
qsort (p+1) hi
partition lo hi = do
pivot <- unsafeRead ma hi
let loop i j
| j == hi = return i
| otherwise = do
aj <- unsafeRead ma j
i' <- if ford aj pivot == GT
then return i
else do
ai <- unsafeRead ma i
unsafeWrite ma j ai
unsafeWrite ma i aj
return $ i + 1
loop i' (j+1)
i <- loop lo lo
ai <- unsafeRead ma i
ahi <- unsafeRead ma hi
unsafeWrite ma hi ai
unsafeWrite ma i ahi
return i
filter :: (ty -> Bool) -> Array ty -> Array ty
filter predicate vec = runST (new len >>= copyFilterFreeze predicate (unsafeIndex vec))
where
!len = lengthSize vec
!end = Offset 0 `offsetPlusE` len
copyFilterFreeze :: PrimMonad prim => (ty -> Bool) -> (Int -> ty) -> MArray ty (PrimState prim) -> prim (Array ty)
copyFilterFreeze predi getVec mvec = loop (Offset 0) (Offset 0) >>= freezeUntilIndex mvec
where
loop d@(Offset di) s@(Offset si)
| s == end = return d
| predi v = unsafeWrite mvec di v >> loop (d+Offset 1) (s+Offset 1)
| otherwise = loop d (s+Offset 1)
where
v = getVec si
freezeUntilIndex :: PrimMonad prim => MArray ty (PrimState prim) -> Offset ty -> prim (Array ty)
freezeUntilIndex mvec d = do
m <- new (offsetAsSize d)
copyAt m (Offset 0) mvec (Offset 0) (offsetAsSize d)
unsafeFreeze m
unsafeFreezeShrink :: PrimMonad prim => MArray ty (PrimState prim) -> Size ty -> prim (Array ty)
unsafeFreezeShrink (MArray start _ ma) n = unsafeFreeze (MArray start n ma)
reverse :: Array ty -> Array ty
reverse a = create len toEnd
where
len = length a
toEnd i = unsafeIndex a (len i 1)
foldl :: (a -> ty -> a) -> a -> Array ty -> a
foldl f initialAcc vec = loop 0 initialAcc
where
len = length vec
loop !i acc
| i == len = acc
| otherwise = loop (i+1) (f acc (unsafeIndex vec i))
foldr :: (ty -> a -> a) -> a -> Array ty -> a
foldr f initialAcc vec = loop 0
where
len = length vec
loop !i
| i == len = initialAcc
| otherwise = unsafeIndex vec i `f` loop (i+1)
foldl' :: (a -> ty -> a) -> a -> Array ty -> a
foldl' f initialAcc vec = loop 0 initialAcc
where
len = length vec
loop !i !acc
| i == len = acc
| otherwise = loop (i+1) (f acc (unsafeIndex vec i))
builderAppend :: PrimMonad state => ty -> Builder (Array ty) (MArray ty) ty state ()
builderAppend v = Builder $ State $ \(i, st) ->
if offsetAsSize i == chunkSize st
then do
cur <- unsafeFreeze (curChunk st)
newChunk <- new (chunkSize st)
unsafeWrite newChunk 0 v
return ((), (Offset 1, st { prevChunks = cur : prevChunks st
, prevChunksSize = chunkSize st + prevChunksSize st
, curChunk = newChunk
}))
else do
let (Offset i') = i
unsafeWrite (curChunk st) i' v
return ((), (i + Offset 1, st))
builderBuild :: PrimMonad m => Int -> Builder (Array ty) (MArray ty) ty m () -> m (Array ty)
builderBuild sizeChunksI ab
| sizeChunksI <= 0 = builderBuild 64 ab
| otherwise = do
first <- new sizeChunks
((), (i, st)) <- runState (runBuilder ab) (Offset 0, BuildingState [] (Size 0) first sizeChunks)
cur <- unsafeFreezeShrink (curChunk st) (offsetAsSize i)
let totalSize = prevChunksSize st + offsetAsSize i
new totalSize >>= fillFromEnd totalSize (cur : prevChunks st) >>= unsafeFreeze
where
sizeChunks = Size sizeChunksI
fillFromEnd _ [] mua = return mua
fillFromEnd !end (x:xs) mua = do
let sz = lengthSize x
unsafeCopyAtRO mua (sizeAsOffset (end sz)) x (Offset 0) sz
fillFromEnd (end sz) xs mua