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