1-- |
2-- Module      : Basement.BoxedArray
3-- License     : BSD-style
4-- Maintainer  : Vincent Hanquez <vincent@snarc.org>
5-- Stability   : experimental
6-- Portability : portable
7--
8-- Simple boxed array abstraction
9--
10{-# LANGUAGE MagicHash #-}
11{-# LANGUAGE BangPatterns #-}
12{-# LANGUAGE UnboxedTuples #-}
13{-# LANGUAGE ScopedTypeVariables #-}
14{-# LANGUAGE MultiParamTypeClasses #-}
15{-# LANGUAGE FlexibleInstances #-}
16module Basement.BoxedArray
17    ( Array
18    , MArray
19    , empty
20    , length
21    , mutableLength
22    , copy
23    , unsafeCopyAtRO
24    , thaw
25    , new
26    , create
27    , unsafeFreeze
28    , unsafeThaw
29    , freeze
30    , unsafeWrite
31    , unsafeRead
32    , unsafeIndex
33    , write
34    , read
35    , index
36    , singleton
37    , replicate
38    , null
39    , take
40    , drop
41    , splitAt
42    , revTake
43    , revDrop
44    , revSplitAt
45    , splitOn
46    , sub
47    , intersperse
48    , span
49    , spanEnd
50    , break
51    , breakEnd
52    , mapFromUnboxed
53    , mapToUnboxed
54    , cons
55    , snoc
56    , uncons
57    , unsnoc
58    -- , findIndex
59    , sortBy
60    , filter
61    , reverse
62    , elem
63    , find
64    , foldl'
65    , foldr
66    , foldl1'
67    , foldr1
68    , all
69    , any
70    , isPrefixOf
71    , isSuffixOf
72    , builderAppend
73    , builderBuild
74    , builderBuild_
75    ) where
76
77import           GHC.Prim
78import           GHC.Types
79import           GHC.ST
80import           Data.Proxy
81import           Basement.Numerical.Additive
82import           Basement.Numerical.Subtractive
83import           Basement.NonEmpty
84import           Basement.Compat.Base
85import qualified Basement.Alg.Class as Alg
86import qualified Basement.Alg.Mutable as Alg
87import           Basement.Compat.MonadTrans
88import           Basement.Compat.Semigroup
89import           Basement.Compat.Primitive
90import           Basement.Types.OffsetSize
91import           Basement.PrimType
92import           Basement.NormalForm
93import           Basement.Monad
94import           Basement.UArray.Base (UArray)
95import qualified Basement.UArray.Base as UArray
96import           Basement.Exception
97import           Basement.MutableBuilder
98import qualified Basement.Compat.ExtList as List
99
100-- | Array of a
101data Array a = Array {-# UNPACK #-} !(Offset a)
102                     {-# UNPACK #-} !(CountOf a)
103                                    (Array# a)
104    deriving (Typeable)
105
106instance Data ty => Data (Array ty) where
107    dataTypeOf _ = arrayType
108    toConstr _   = error "toConstr"
109    gunfold _ _  = error "gunfold"
110
111arrayType :: DataType
112arrayType = mkNoRepType "Foundation.Array"
113
114instance NormalForm a => NormalForm (Array a) where
115    toNormalForm arr = loop 0
116      where
117        !sz = length arr
118        loop !i
119            | i .==# sz = ()
120            | otherwise = unsafeIndex arr i `seq` loop (i+1)
121
122-- | Mutable Array of a
123data MArray a st = MArray {-# UNPACK #-} !(Offset a)
124                          {-# UNPACK #-} !(CountOf a)
125                                         (MutableArray# st a)
126    deriving (Typeable)
127
128instance Functor Array where
129    fmap = map
130
131instance Semigroup (Array a) where
132    (<>) = append
133instance Monoid (Array a) where
134    mempty  = empty
135    mappend = append
136    mconcat = concat
137
138instance Show a => Show (Array a) where
139    show v = show (toList v)
140
141instance Eq a => Eq (Array a) where
142    (==) = equal
143instance Ord a => Ord (Array a) where
144    compare = vCompare
145
146instance IsList (Array ty) where
147    type Item (Array ty) = ty
148    fromList = vFromList
149    fromListN len = vFromListN (CountOf len)
150    toList = vToList
151
152-- | return the numbers of elements in a mutable array
153mutableLength :: MArray ty st -> Int
154mutableLength (MArray _ (CountOf len) _) = len
155{-# INLINE mutableLength #-}
156
157-- | return the numbers of elements in a mutable array
158mutableLengthSize :: MArray ty st -> CountOf ty
159mutableLengthSize (MArray _ size _) = size
160{-# INLINE mutableLengthSize #-}
161
162-- | Return the element at a specific index from an array.
163--
164-- If the index @n is out of bounds, an error is raised.
165index :: Array ty -> Offset ty -> ty
166index array n
167    | isOutOfBound n len = outOfBound OOB_Index n len
168    | otherwise          = unsafeIndex array n
169  where len = length array
170{-# INLINE index #-}
171
172-- | Return the element at a specific index from an array without bounds checking.
173--
174-- Reading from invalid memory can return unpredictable and invalid values.
175-- use 'index' if unsure.
176unsafeIndex :: Array ty -> Offset ty -> ty
177unsafeIndex (Array start _ a) ofs = primArrayIndex a (start+ofs)
178{-# INLINE unsafeIndex #-}
179
180-- | read a cell in a mutable array.
181--
182-- If the index is out of bounds, an error is raised.
183read :: PrimMonad prim => MArray ty (PrimState prim) -> Offset ty -> prim ty
184read array n
185    | isOutOfBound n len = primOutOfBound OOB_Read n len
186    | otherwise          = unsafeRead array n
187  where len = mutableLengthSize array
188{-# INLINE read #-}
189
190-- | read from a cell in a mutable array without bounds checking.
191--
192-- Reading from invalid memory can return unpredictable and invalid values.
193-- use 'read' if unsure.
194unsafeRead :: PrimMonad prim => MArray ty (PrimState prim) -> Offset ty -> prim ty
195unsafeRead (MArray start _ ma) i = primMutableArrayRead ma (start + i)
196{-# INLINE unsafeRead #-}
197
198-- | Write to a cell in a mutable array.
199--
200-- If the index is out of bounds, an error is raised.
201write :: PrimMonad prim => MArray ty (PrimState prim) -> Offset ty -> ty -> prim ()
202write array n val
203    | isOutOfBound n len = primOutOfBound OOB_Write n len
204    | otherwise          = unsafeWrite array n val
205  where len = mutableLengthSize array
206{-# INLINE write #-}
207
208-- | write to a cell in a mutable array without bounds checking.
209--
210-- Writing with invalid bounds will corrupt memory and your program will
211-- become unreliable. use 'write' if unsure.
212unsafeWrite :: PrimMonad prim => MArray ty (PrimState prim) -> Offset ty -> ty -> prim ()
213unsafeWrite (MArray start _ ma) ofs v =
214    primMutableArrayWrite ma (start + ofs) v
215{-# INLINE unsafeWrite #-}
216
217-- | Freeze a mutable array into an array.
218--
219-- the MArray must not be changed after freezing.
220unsafeFreeze :: PrimMonad prim => MArray ty (PrimState prim) -> prim (Array ty)
221unsafeFreeze (MArray ofs sz ma) = primitive $ \s1 ->
222    case unsafeFreezeArray# ma s1 of
223        (# s2, a #) -> (# s2, Array ofs sz a #)
224{-# INLINE unsafeFreeze #-}
225
226-- | Thaw an immutable array.
227--
228-- The Array must not be used after thawing.
229unsafeThaw :: PrimMonad prim => Array ty -> prim (MArray ty (PrimState prim))
230unsafeThaw (Array ofs sz a) = primitive $ \st -> (# st, MArray ofs sz (unsafeCoerce# a) #)
231{-# INLINE unsafeThaw #-}
232
233-- | Thaw an array to a mutable array.
234--
235-- the array is not modified, instead a new mutable array is created
236-- and every values is copied, before returning the mutable array.
237thaw :: PrimMonad prim => Array ty -> prim (MArray ty (PrimState prim))
238thaw array = do
239    m <- new (length array)
240    unsafeCopyAtRO m (Offset 0) array (Offset 0) (length array)
241    pure m
242{-# INLINE thaw #-}
243
244freeze :: PrimMonad prim => MArray ty (PrimState prim) -> prim (Array ty)
245freeze marray = do
246    m <- new sz
247    copyAt m (Offset 0) marray (Offset 0) sz
248    unsafeFreeze m
249  where
250    sz = mutableLengthSize marray
251
252-- | Copy the element to a new element array
253copy :: Array ty -> Array ty
254copy a = runST (unsafeThaw a >>= freeze)
255
256-- | Copy a number of elements from an array to another array with offsets
257copyAt :: PrimMonad prim
258       => MArray ty (PrimState prim) -- ^ destination array
259       -> Offset ty                  -- ^ offset at destination
260       -> MArray ty (PrimState prim) -- ^ source array
261       -> Offset ty                  -- ^ offset at source
262       -> CountOf ty                 -- ^ number of elements to copy
263       -> prim ()
264copyAt dst od src os n = loop od os
265  where -- !endIndex = os `offsetPlusE` n
266        loop d s
267            | s .==# n  = pure ()
268            | otherwise = unsafeRead src s >>= unsafeWrite dst d >> loop (d+1) (s+1)
269
270-- | Copy @n@ sequential elements from the specified offset in a source array
271--   to the specified position in a destination array.
272--
273--   This function does not check bounds. Accessing invalid memory can return
274--   unpredictable and invalid values.
275unsafeCopyAtRO :: PrimMonad prim
276               => MArray ty (PrimState prim) -- ^ destination array
277               -> Offset ty                  -- ^ offset at destination
278               -> Array ty                   -- ^ source array
279               -> Offset ty                  -- ^ offset at source
280               -> CountOf ty                    -- ^ number of elements to copy
281               -> prim ()
282unsafeCopyAtRO (MArray (Offset (I# dstart)) _ da) (Offset (I# dofs))
283               (Array  (Offset (I# sstart)) _ sa) (Offset (I# sofs))
284               (CountOf (I# n)) =
285    primitive $ \st ->
286        (# copyArray# sa (sstart +# sofs) da (dstart +# dofs) n st, () #)
287
288-- | Allocate a new array with a fill function that has access to the elements of
289--   the source array.
290unsafeCopyFrom :: Array ty -- ^ Source array
291               -> CountOf ty  -- ^ Length of the destination array
292               -> (Array ty -> Offset ty -> MArray ty s -> ST s ())
293               -- ^ Function called for each element in the source array
294               -> ST s (Array ty) -- ^ Returns the filled new array
295unsafeCopyFrom v' newLen f = new newLen >>= fill (Offset 0) f >>= unsafeFreeze
296  where len = length v'
297        endIdx = Offset 0 `offsetPlusE` len
298        fill i f' r'
299            | i == endIdx = pure r'
300            | otherwise   = do f' v' i r'
301                               fill (i + Offset 1) f' r'
302
303-- | Create a new mutable array of size @n.
304--
305-- all the cells are uninitialized and could contains invalid values.
306--
307-- All mutable arrays are allocated on a 64 bits aligned addresses
308-- and always contains a number of bytes multiples of 64 bits.
309new :: PrimMonad prim => CountOf ty -> prim (MArray ty (PrimState prim))
310new sz@(CountOf (I# n)) = primitive $ \s1 ->
311    case newArray# n (error "vector: internal error uninitialized vector") s1 of
312        (# s2, ma #) -> (# s2, MArray (Offset 0) sz ma #)
313
314-- | Create a new array of size @n by settings each cells through the
315-- function @f.
316create :: forall ty . CountOf ty -- ^ the size of the array
317       -> (Offset ty -> ty)   -- ^ the function that set the value at the index
318       -> Array ty            -- ^ the array created
319create n initializer = runST (new n >>= iter initializer)
320  where
321    iter :: PrimMonad prim => (Offset ty -> ty) -> MArray ty (PrimState prim) -> prim (Array ty)
322    iter f ma = loop 0
323      where
324        loop s
325            | s .==# n  = unsafeFreeze ma
326            | otherwise = unsafeWrite ma s (f s) >> loop (s+1)
327        {-# INLINE loop #-}
328    {-# INLINE iter #-}
329
330-----------------------------------------------------------------------
331-- higher level collection implementation
332-----------------------------------------------------------------------
333equal :: Eq a => Array a -> Array a -> Bool
334equal a b = (len == length b) && eachEqual 0
335  where
336    len = length a
337    eachEqual !i
338        | i .==# len                         = True
339        | unsafeIndex a i /= unsafeIndex b i = False
340        | otherwise                          = eachEqual (i+1)
341
342vCompare :: Ord a => Array a -> Array a -> Ordering
343vCompare a b = loop 0
344  where
345    !la = length a
346    !lb = length b
347    loop n
348        | n .==# la = if la == lb then EQ else LT
349        | n .==# lb = GT
350        | otherwise =
351            case unsafeIndex a n `compare` unsafeIndex b n of
352                EQ -> loop (n+1)
353                r  -> r
354
355empty :: Array a
356empty = runST $ onNewArray 0 (\_ s -> s)
357
358length :: Array a -> CountOf a
359length (Array _ sz _) = sz
360
361vFromList :: [a] -> Array a
362vFromList l = runST (new len >>= loop 0 l)
363  where
364    len = List.length l
365    loop _ []     ma = unsafeFreeze ma
366    loop i (x:xs) ma = unsafeWrite ma i x >> loop (i+1) xs ma
367
368-- | just like vFromList but with a length hint.
369--
370-- The resulting array is guarantee to have been allocated to the length
371-- specified, but the slice might point to the initialized cells only in
372-- case the length is bigger than the list.
373--
374-- If the length is too small, then the list is truncated.
375--
376vFromListN :: forall a . CountOf a -> [a] -> Array a
377vFromListN len l = runST $ do
378    ma <- new len
379    sz <- loop 0 l ma
380    unsafeFreezeShrink ma sz
381  where
382    -- TODO rewrite without ma as parameter
383    loop :: Offset a -> [a] -> MArray a s -> ST s (CountOf a)
384    loop i []     _  = return (offsetAsSize i)
385    loop i (x:xs) ma
386        | i .==# len = return (offsetAsSize i)
387        | otherwise  = unsafeWrite ma i x >> loop (i+1) xs ma
388
389vToList :: Array a -> [a]
390vToList v
391    | len == 0  = []
392    | otherwise = fmap (unsafeIndex v) [0..sizeLastOffset len]
393  where !len = length v
394
395-- | Append 2 arrays together by creating a new bigger array
396append :: Array ty -> Array ty -> Array ty
397append a b = runST $ do
398    r  <- new (la+lb)
399    unsafeCopyAtRO r (Offset 0) a (Offset 0) la
400    unsafeCopyAtRO r (sizeAsOffset la) b (Offset 0) lb
401    unsafeFreeze r
402  where la = length a
403        lb = length b
404
405concat :: [Array ty] -> Array ty
406concat l = runST $ do
407    r <- new (mconcat $ fmap length l)
408    loop r (Offset 0) l
409    unsafeFreeze r
410  where loop _ _ []     = pure ()
411        loop r i (x:xs) = do
412            unsafeCopyAtRO r i x (Offset 0) lx
413            loop r (i `offsetPlusE` lx) xs
414          where lx = length x
415
416{-
417modify :: PrimMonad m
418       => Array a
419       -> (MArray (PrimState m) a -> m ())
420       -> m (Array a)
421modify (Array a) f = primitive $ \st -> do
422    case thawArray# a 0# (sizeofArray# a) st of
423        (# st2, mv #) ->
424            case internal_ (f $ MArray mv) st2 of
425                st3 ->
426                    case unsafeFreezeArray# mv st3 of
427                        (# st4, a' #) -> (# st4, Array a' #)
428-}
429
430-----------------------------------------------------------------------
431-- helpers
432
433onNewArray :: PrimMonad m
434           => Int
435           -> (MutableArray# (PrimState m) a -> State# (PrimState m) -> State# (PrimState m))
436           -> m (Array a)
437onNewArray len@(I# len#) f = primitive $ \st -> do
438    case newArray# len# (error "onArray") st of { (# st2, mv #) ->
439    case f mv st2                            of { st3           ->
440    case unsafeFreezeArray# mv st3           of { (# st4, a #)  ->
441        (# st4, Array (Offset 0) (CountOf len) a #) }}}
442
443-----------------------------------------------------------------------
444
445
446null :: Array ty -> Bool
447null = (==) 0 . length
448
449take :: CountOf ty -> Array ty -> Array ty
450take nbElems a@(Array start len arr)
451    | nbElems <= 0 = empty
452    | n == len     = a
453    | otherwise    = Array start n arr
454  where
455    n = min nbElems len
456
457drop :: CountOf ty -> Array ty -> Array ty
458drop nbElems a@(Array start len arr)
459    | nbElems <= 0                               = a
460    | Just nbTails <- len - nbElems, nbTails > 0 = Array (start `offsetPlusE` nbElems) nbTails arr
461    | otherwise                                  = empty
462
463splitAt :: CountOf ty -> Array ty -> (Array ty, Array ty)
464splitAt nbElems a@(Array start len arr)
465    | nbElems <= 0 = (empty, a)
466    | Just nbTails <- len - nbElems, nbTails > 0 = ( Array start                         nbElems arr
467                                                   , Array (start `offsetPlusE` nbElems) nbTails arr)
468    | otherwise = (a, empty)
469
470-- inverse a CountOf that is specified from the end (e.g. take n elements from the end)
471countFromStart :: Array ty -> CountOf ty -> CountOf ty
472countFromStart v sz@(CountOf sz')
473    | sz >= len = CountOf 0
474    | otherwise = CountOf (len' - sz')
475  where len@(CountOf len') = length v
476
477revTake :: CountOf ty -> Array ty -> Array ty
478revTake n v = drop (countFromStart v n) v
479
480revDrop :: CountOf ty -> Array ty -> Array ty
481revDrop n v = take (countFromStart v n) v
482
483revSplitAt :: CountOf ty -> Array ty -> (Array ty, Array ty)
484revSplitAt n v = (drop idx v, take idx v) where idx = countFromStart v n
485
486splitOn ::  (ty -> Bool) -> Array ty -> [Array ty]
487splitOn predicate vec
488    | len == CountOf 0 = [mempty]
489    | otherwise     = loop (Offset 0) (Offset 0)
490  where
491    !len = length vec
492    !endIdx = Offset 0 `offsetPlusE` len
493    loop prevIdx idx
494        | idx == endIdx = [sub vec prevIdx idx]
495        | otherwise     =
496            let e = unsafeIndex vec idx
497                idx' = idx + 1
498             in if predicate e
499                    then sub vec prevIdx idx : loop idx' idx'
500                    else loop prevIdx idx'
501
502sub :: Array ty -> Offset ty -> Offset ty -> Array ty
503sub (Array start len a) startIdx expectedEndIdx
504    | startIdx == endIdx           = empty
505    | otherwise                    = Array (start + startIdx) newLen a
506  where
507    newLen = endIdx - startIdx
508    endIdx = min expectedEndIdx (sizeAsOffset len)
509
510break ::  (ty -> Bool) -> Array ty -> (Array ty, Array ty)
511break predicate v = findBreak 0
512  where
513    !len = length v
514    findBreak i
515        | i .==# len  = (v, empty)
516        | otherwise   =
517            if predicate (unsafeIndex v i)
518                then splitAt (offsetAsSize i) v
519                else findBreak (i+1)
520
521breakEnd ::  (ty -> Bool) -> Array ty -> (Array ty, Array ty)
522breakEnd predicate v = findBreak (sizeAsOffset len)
523  where
524    !len = length v
525    findBreak !i
526        | i == 0      = (v, empty)
527        | predicate e = splitAt (offsetAsSize i) v
528        | otherwise   = findBreak i'
529      where
530        e = unsafeIndex v i'
531        i' = i `offsetSub` 1
532
533intersperse :: ty -> Array ty -> Array ty
534intersperse sep v = case len - 1 of
535    Nothing -> v
536    Just 0 -> v
537    Just more -> runST $ unsafeCopyFrom v (len + more) (go (Offset 0 `offsetPlusE` more) sep)
538  where len = length v
539        -- terminate 1 before the end
540
541        go :: Offset ty -> ty -> Array ty -> Offset ty -> MArray ty s -> ST s ()
542        go endI sep' oldV oldI newV
543            | oldI == endI = unsafeWrite newV dst e
544            | otherwise    = do
545                unsafeWrite newV dst e
546                unsafeWrite newV (dst + 1) sep'
547          where
548            e = unsafeIndex oldV oldI
549            dst = oldI + oldI
550
551span ::  (ty -> Bool) -> Array ty -> (Array ty, Array ty)
552span p = break (not . p)
553
554spanEnd ::  (ty -> Bool) -> Array ty -> (Array ty, Array ty)
555spanEnd p = breakEnd (not . p)
556
557map :: (a -> b) -> Array a -> Array b
558map f a = create (sizeCast Proxy $ length a) (\i -> f $ unsafeIndex a (offsetCast Proxy i))
559
560mapFromUnboxed :: PrimType a => (a -> b) -> UArray a -> Array b
561mapFromUnboxed f arr = vFromListN (sizeCast Proxy $ UArray.length arr) . fmap f . toList $ arr
562
563mapToUnboxed :: PrimType b => (a -> b) -> Array a -> UArray b
564mapToUnboxed f arr = UArray.vFromListN (sizeCast Proxy $ length arr) . fmap f . toList $ arr
565
566{-
567mapIndex :: (Int -> a -> b) -> Array a -> Array b
568mapIndex f a = create (length a) (\i -> f i $ unsafeIndex a i)
569-}
570
571singleton :: ty -> Array ty
572singleton e = runST $ do
573    a <- new 1
574    unsafeWrite a 0 e
575    unsafeFreeze a
576
577replicate :: CountOf ty -> ty -> Array ty
578replicate sz ty = create sz (const ty)
579
580cons :: ty -> Array ty -> Array ty
581cons e vec
582    | len == CountOf 0 = singleton e
583    | otherwise     = runST $ do
584        mv <- new (len + CountOf 1)
585        unsafeWrite mv 0 e
586        unsafeCopyAtRO mv (Offset 1) vec (Offset 0) len
587        unsafeFreeze mv
588  where
589    !len = length vec
590
591snoc ::  Array ty -> ty -> Array ty
592snoc vec e
593    | len == 0  = singleton e
594    | otherwise = runST $ do
595        mv <- new (len + 1)
596        unsafeCopyAtRO mv 0 vec 0 len
597        unsafeWrite mv (sizeAsOffset len) e
598        unsafeFreeze mv
599  where
600    !len = length vec
601
602uncons :: Array ty -> Maybe (ty, Array ty)
603uncons vec
604    | len == 0  = Nothing
605    | otherwise = Just (unsafeIndex vec 0, drop 1 vec)
606  where
607    !len = length vec
608
609unsnoc :: Array ty -> Maybe (Array ty, ty)
610unsnoc vec = case len - 1 of
611    Nothing -> Nothing
612    Just newLen -> Just (take newLen vec, unsafeIndex vec (sizeLastOffset len))
613  where
614    !len = length vec
615
616elem :: Eq ty => ty -> Array ty -> Bool
617elem !ty arr = loop 0
618  where
619    !sz = length arr
620    loop !i | i .==# sz = False
621            | t == ty   = True
622            | otherwise = loop (i+1)
623      where t = unsafeIndex arr i
624
625find :: (ty -> Bool) -> Array ty -> Maybe ty
626find predicate vec = loop 0
627  where
628    !len = length vec
629    loop i
630        | i .==# len = Nothing
631        | otherwise  =
632            let e = unsafeIndex vec i
633             in if predicate e then Just e else loop (i+1)
634
635instance (PrimMonad prim, st ~ PrimState prim)
636         => Alg.RandomAccess (MArray ty st) prim ty where
637    read (MArray _ _ mba) = primMutableArrayRead mba
638    write (MArray _ _ mba) = primMutableArrayWrite mba
639
640sortBy :: forall ty . (ty -> ty -> Ordering) -> Array ty -> Array ty
641sortBy xford vec
642    | len == 0  = empty
643    | otherwise = runST (thaw vec >>= doSort xford)
644  where
645    len = length vec
646    doSort :: PrimMonad prim => (ty -> ty -> Ordering) -> MArray ty (PrimState prim) -> prim (Array ty)
647    doSort ford ma = Alg.inplaceSortBy ford 0 len ma >> unsafeFreeze ma
648
649filter :: forall ty . (ty -> Bool) -> Array ty -> Array ty
650filter predicate vec = runST (new len >>= copyFilterFreeze predicate (unsafeIndex vec))
651  where
652    !len = length vec
653    copyFilterFreeze :: PrimMonad prim => (ty -> Bool) -> (Offset ty -> ty) -> MArray ty (PrimState prim) -> prim (Array ty)
654    copyFilterFreeze predi getVec mvec = loop (Offset 0) (Offset 0) >>= freezeUntilIndex mvec
655      where
656        loop d s
657            | s .==# len  = pure d
658            | predi v     = unsafeWrite mvec d v >> loop (d+1) (s+1)
659            | otherwise   = loop d (s+1)
660          where
661            v = getVec s
662
663freezeUntilIndex :: PrimMonad prim => MArray ty (PrimState prim) -> Offset ty -> prim (Array ty)
664freezeUntilIndex mvec d = do
665    m <- new (offsetAsSize d)
666    copyAt m (Offset 0) mvec (Offset 0) (offsetAsSize d)
667    unsafeFreeze m
668
669unsafeFreezeShrink :: PrimMonad prim => MArray ty (PrimState prim) -> CountOf ty -> prim (Array ty)
670unsafeFreezeShrink (MArray start _ ma) n = unsafeFreeze (MArray start n ma)
671
672reverse :: Array ty -> Array ty
673reverse a = create len toEnd
674  where
675    len@(CountOf s) = length a
676    toEnd (Offset i) = unsafeIndex a (Offset (s - 1 - i))
677
678foldr :: (ty -> a -> a) -> a -> Array ty -> a
679foldr f initialAcc vec = loop 0
680  where
681    len = length vec
682    loop !i
683        | i .==# len = initialAcc
684        | otherwise  = unsafeIndex vec i `f` loop (i+1)
685
686foldl' :: (a -> ty -> a) -> a -> Array ty -> a
687foldl' f initialAcc vec = loop 0 initialAcc
688  where
689    len = length vec
690    loop !i !acc
691        | i .==# len = acc
692        | otherwise  = loop (i+1) (f acc (unsafeIndex vec i))
693
694foldl1' :: (ty -> ty -> ty) -> NonEmpty (Array ty) -> ty
695foldl1' f arr = let (initialAcc, rest) = splitAt 1 $ getNonEmpty arr
696               in foldl' f (unsafeIndex initialAcc 0) rest
697
698foldr1 :: (ty -> ty -> ty) -> NonEmpty (Array ty) -> ty
699foldr1 f arr = let (initialAcc, rest) = revSplitAt 1 $ getNonEmpty arr
700               in foldr f (unsafeIndex initialAcc 0) rest
701
702all :: (ty -> Bool) -> Array ty -> Bool
703all p ba = loop 0
704  where
705    len = length ba
706    loop !i
707      | i .==# len = True
708      | not $ p (unsafeIndex ba i) = False
709      | otherwise = loop (i + 1)
710
711any :: (ty -> Bool) -> Array ty -> Bool
712any p ba = loop 0
713  where
714    len = length ba
715    loop !i
716      | i .==# len = False
717      | p (unsafeIndex ba i) = True
718      | otherwise = loop (i + 1)
719
720isPrefixOf :: Eq ty => Array ty -> Array ty -> Bool
721isPrefixOf pre arr
722    | pLen > pArr = False
723    | otherwise   = pre == take pLen arr
724  where
725    !pLen = length pre
726    !pArr = length arr
727
728isSuffixOf :: Eq ty => Array ty -> Array ty -> Bool
729isSuffixOf suffix arr
730    | pLen > pArr = False
731    | otherwise   = suffix == revTake pLen arr
732  where
733    !pLen = length suffix
734    !pArr = length arr
735
736builderAppend :: PrimMonad state => ty -> Builder (Array ty) (MArray ty) ty state err ()
737builderAppend v = Builder $ State $ \(i, st, e) ->
738    if i .==# chunkSize st
739        then do
740            cur      <- unsafeFreeze (curChunk st)
741            newChunk <- new (chunkSize st)
742            unsafeWrite newChunk 0 v
743            pure ((), (Offset 1, st { prevChunks     = cur : prevChunks st
744                                      , prevChunksSize = chunkSize st + prevChunksSize st
745                                      , curChunk       = newChunk
746                                      }, e))
747        else do
748            unsafeWrite (curChunk st) i v
749            pure ((), (i+1, st, e))
750
751builderBuild :: PrimMonad m => Int -> Builder (Array ty) (MArray ty) ty m err () -> m (Either err (Array ty))
752builderBuild sizeChunksI ab
753    | sizeChunksI <= 0 = builderBuild 64 ab
754    | otherwise        = do
755        first      <- new sizeChunks
756        (i, st, e) <- snd <$> runState (runBuilder ab) (Offset 0, BuildingState [] (CountOf 0) first sizeChunks, Nothing)
757        case e of
758          Just err -> pure (Left err)
759          Nothing -> do
760            cur <- unsafeFreezeShrink (curChunk st) (offsetAsSize i)
761            -- Build final array
762            let totalSize = prevChunksSize st + offsetAsSize i
763            bytes <- new totalSize >>= fillFromEnd totalSize (cur : prevChunks st) >>= unsafeFreeze
764            pure (Right bytes)
765  where
766    sizeChunks = CountOf sizeChunksI
767
768    fillFromEnd _    []     mua = pure mua
769    fillFromEnd !end (x:xs) mua = do
770        let sz = length x
771        let start = end `sizeSub` sz
772        unsafeCopyAtRO mua (sizeAsOffset start) x (Offset 0) sz
773        fillFromEnd start xs mua
774
775builderBuild_ :: PrimMonad m => Int -> Builder (Array ty) (MArray ty) ty m () () -> m (Array ty)
776builderBuild_ sizeChunksI ab = either (\() -> internalError "impossible output") id <$> builderBuild sizeChunksI ab
777