2{-# LANGUAGE UnboxedTuples #-}
4-- Module : Streamly.Internal.Data.MutArray.Generic
5-- Copyright : (c) 2020 Composewell Technologies
6-- License : BSD3-3-Clause
7-- Maintainer : streamly@composewell.com
8-- Stability : experimental
11module Streamly.Internal.Data.MutArray.Generic
17 -- * Constructing and Writing
21 -- *** Uninitialized Arrays
46 -- , putFrom -- start writing at the given position
47 -- , putUpto -- write from beginning up to the given position
57 -- * Growing and Shrinking
58 -- Arrays grow only at the end, though it is possible to grow on both sides
59 -- and therefore have a cons as well as snoc. But that will require two
60 -- bounds in the array representation.
66 -- ** Appending elements
73 -- ** Appending streams
74 -- , writeAppendNUnsafe
80 -- These are not the same as slicing the array at the beginning, they may
81 -- reduce the length as well as the capacity of the array.
86 -- * Eliminating and Reading
91 , producerWith -- experimental
92 , producer -- experimental
112 -- * In-place Mutation Algorithms
127 -- * Arrays of arrays
128 -- We can add dimensionality parameter to the array type to get
129 -- multidimensional arrays. Multidimensional arrays would just be a
130 -- convenience wrapper on top of single dimensional arrays.
132 -- | Operations dealing with multiple arrays, streams of arrays or
133 -- multidimensional array representations.
135 -- ** Construct from streams
137 -- , arrayStreamKFromStreamD
140 -- ** Eliminate to streams
142 -- , flattenArraysRev
143 -- , fromArrayStreamK
145 -- ** Construct from arrays
146 -- get chunks without copying
149 -- , getSlicesFromLenN
150 -- , splitAt -- XXX should be able to express using getSlice
153 -- ** Appending arrays
169import Control.Monad (when)
170import Control.Monad.IO.Class (MonadIO(..))
171import Data.Functor.Identity (Identity(..))
180import GHC.IO (IO(..))
181import GHC.Int (Int(..))
182import Streamly.Internal.Data.Fold.Type (Fold(..))
183import Streamly.Internal.Data.Producer.Type (Producer (..))
184import Streamly.Internal.Data.Unfold.Type (Unfold(..))
185import Streamly.Internal.Data.Stream.Type (Stream)
186import Streamly.Internal.Data.SVar.Type (adaptState)
188import qualified Streamly.Internal.Data.Fold.Type as FL
189import qualified Streamly.Internal.Data.Producer as Producer
190import qualified Streamly.Internal.Data.Stream.Type as D
191import qualified Streamly.Internal.Data.Stream.Generate as D
192import qualified Streamly.Internal.Data.Stream.Lift as D
193import qualified Streamly.Internal.Data.StreamK.Type as K
195import Prelude hiding (read, length)
197#include "DocTestDataMutArrayGeneric.hs"
199-------------------------------------------------------------------------------
201-------------------------------------------------------------------------------
205 { arrContents# :: MutableArray# RealWorld a
206 -- ^ The internal contents of the array representing the entire array.
208 , arrStart :: {-# UNPACK #-}!Int
209 -- ^ The starting index of this slice.
211 , arrLen :: {-# UNPACK #-}!Int
212 -- ^ The length of this slice.
214 , arrTrueLen :: {-# UNPACK #-}!Int
215 -- ^ This is the true length of the array. Coincidentally, this also
216 -- represents the first index beyond the maximum acceptable index of
217 -- the array. This is specific to the array contents itself and not
218 -- dependent on the slice. This value should not change and is shared
219 -- across all the slices.
222{-# INLINE bottomElement #-}
228 , "This is the bottom element of the array."
229 , "This is a place holder and should never be reached!"
234 funcName = "Streamly.Internal.Data.MutArray.Generic.bottomElement:"
236-- XXX Would be nice if GHC can provide something like newUninitializedArray# so
237-- that we do not have to write undefined or error in the whole array.
239-- | @new count@ allocates a zero length array that can be extended to hold
240-- up to 'count' items without reallocating.
244new :: MonadIO m => Int -> m (MutArray a)
249 case newArray# n# bottomElement s# of
251 let ma = MutArray arr# 0 0 n
254-- XXX This could be pure?
259-- >>> nil = MutArray.new 0
261nil :: MonadIO m => m (MutArray a)
264-------------------------------------------------------------------------------
266-------------------------------------------------------------------------------
268-- | Write the given element to the given index of the 'MutableArray#'. Does not
269-- check if the index is out of bounds of the array.
272{-# INLINE putIndexUnsafeWith #-}
273putIndexUnsafeWith :: MonadIO m => Int -> MutableArray# RealWorld a -> a -> m ()
274putIndexUnsafeWith n _arrContents# x =
280 let s1# = writeArray# _arrContents# n# x s#
283-- | Write the given element to the given index of the array. Does not check if
284-- the index is out of bounds of the array.
287{-# INLINE putIndexUnsafe #-}
288putIndexUnsafe :: forall m a. MonadIO m => Int -> MutArray a -> a -> m ()
289putIndexUnsafe i MutArray {..} x =
290 assert (i >= 0 && i < arrLen)
291 putIndexUnsafeWith (i + arrStart) arrContents# x
293invalidIndex :: String -> Int -> a
294invalidIndex label i =
295 error $ label ++ ": invalid array index " ++ show i
297-- | /O(1)/ Write the given element at the given index in the array.
298-- Performs in-place mutation of the array.
300-- >>> putIndex ix arr val = MutArray.modifyIndex ix arr (const (val, ()))
303{-# INLINE putIndex #-}
304putIndex :: MonadIO m => Int -> MutArray a -> a -> m ()
305putIndex i arr@MutArray {..} x =
306 if i >= 0 && i < arrLen
307 then putIndexUnsafe i arr x
308 else invalidIndex "putIndex" i
310-- | Write an input stream of (index, value) pairs to an array. Throws an
311-- error if any index is out of bounds.
314{-# INLINE putIndices #-}
315putIndices :: MonadIO m
316 => MutArray a -> Fold m (Int, a) ()
317putIndices arr = FL.foldlM' step (return ())
321 step () (i, x) = liftIO (putIndex i arr x)
323-- | Modify a given index of an array using a modifier function without checking
326-- Unsafe because it does not check the bounds of the array.
329modifyIndexUnsafe :: MonadIO m => Int -> MutArray a -> (a -> (a, b)) -> m b
330modifyIndexUnsafe i MutArray {..} f = do
336 case readArray# arrContents# n# s# of
339 s2# = writeArray# arrContents# n# a1 s1#
342-- | Modify a given index of an array using a modifier function.
345modifyIndex :: MonadIO m => Int -> MutArray a -> (a -> (a, b)) -> m b
346modifyIndex i arr@MutArray {..} f = do
347 if i >= 0 && i < arrLen
348 then modifyIndexUnsafe i arr f
349 else invalidIndex "modifyIndex" i
351-------------------------------------------------------------------------------
353-------------------------------------------------------------------------------
355-- | Reallocates the array according to the new size. This is a safe function
356-- that always creates a new array and copies the old array into the new one.
357-- If the reallocated size is less than the original array it results in a
358-- truncated version of the original array.
360realloc :: MonadIO m => Int -> MutArray a -> m (MutArray a)
363 let !newLen@(I# newLen#) = min n (arrLen arr)
364 !(I# arrS#) = arrStart arr
365 !(I# arr1S#) = arrStart arr1
366 arrC# = arrContents# arr
367 arr1C# = arrContents# arr1
371 let s1# = copyMutableArray# arrC# arrS# arr1C# arr1S# newLen# s#
372 in (# s1#, arr1 {arrLen = newLen, arrTrueLen = n} #)
375 MonadIO m => String -> (Int -> Int) -> Int -> MutArray a -> m (MutArray a)
376reallocWith label sizer reqSize arr = do
377 let oldSize = arrLen arr
378 newSize = sizer oldSize
379 safeSize = max newSize (oldSize + reqSize)
380 assert (newSize >= oldSize + reqSize || error badSize) (return ())
387 , ": new array size is less than required size "
389 , ". Please check the sizing function passed."
392-------------------------------------------------------------------------------
394-------------------------------------------------------------------------------
396-- XXX Not sure of the behavior of writeArray# if we specify an index which is
397-- out of bounds. This comment should be rewritten based on that.
398-- | Really really unsafe, appends the element into the first array, may
399-- cause silent data corruption or if you are lucky a segfault if the index
403{-# INLINE snocUnsafe #-}
404snocUnsafe :: MonadIO m => MutArray a -> a -> m (MutArray a)
405snocUnsafe arr@MutArray {..} a = do
406 assert (arrStart + arrLen < arrTrueLen) (return ())
407 let arr1 = arr {arrLen = arrLen + 1}
408 putIndexUnsafe arrLen arr1 a
411-- NOINLINE to move it out of the way and not pollute the instruction cache.
412{-# NOINLINE snocWithRealloc #-}
413snocWithRealloc :: MonadIO m => (Int -> Int) -> MutArray a -> a -> m (MutArray a)
414snocWithRealloc sizer arr x = do
415 arr1 <- reallocWith "snocWithRealloc" sizer 1 arr
418-- | @snocWith sizer arr elem@ mutates @arr@ to append @elem@. The length of
419-- the array increases by 1.
421-- If there is no reserved space available in @arr@ it is reallocated to a size
422-- in bytes determined by the @sizer oldSize@ function, where @oldSize@ is the
423-- original size of the array.
425-- Note that the returned array may be a mutated version of the original array.
428{-# INLINE snocWith #-}
429snocWith :: MonadIO m => (Int -> Int) -> MutArray a -> a -> m (MutArray a)
430snocWith sizer arr@MutArray {..} x = do
431 if arrStart + arrLen < arrTrueLen
432 then snocUnsafe arr x
433 else snocWithRealloc sizer arr x
435-- XXX round it to next power of 2.
437-- | The array is mutated to append an additional element to it. If there is no
438-- reserved space available in the array then it is reallocated to double the
441-- This is useful to reduce allocations when appending unknown number of
444-- Note that the returned array may be a mutated version of the original array.
446-- >>> snoc = MutArray.snocWith (* 2)
448-- Performs O(n * log n) copies to grow, but is liberal with memory allocation.
452snoc :: MonadIO m => MutArray a -> a -> m (MutArray a)
455-- | Make the uninitialized memory in the array available for use extending it
456-- by the supplied length beyond the current length of the array. The array may
460uninit :: MonadIO m => MutArray a -> Int -> m (MutArray a)
461uninit arr@MutArray{..} len =
462 if arrStart + arrLen + len <= arrTrueLen
463 then return $ arr {arrLen = arrLen + len}
464 else realloc (arrLen + len) arr
466-------------------------------------------------------------------------------
468-------------------------------------------------------------------------------
470-- | Return the element at the specified index without checking the bounds from
471-- a @MutableArray# RealWorld@.
473-- Unsafe because it does not check the bounds of the array.
474{-# INLINE getIndexUnsafeWith #-}
475getIndexUnsafeWith :: MonadIO m => MutableArray# RealWorld a -> Int -> m a
476getIndexUnsafeWith _arrContents# n =
481 in readArray# _arrContents# i# s#
483-- | Return the element at the specified index without checking the bounds.
485-- Unsafe because it does not check the bounds of the array.
486{-# INLINE_NORMAL getIndexUnsafe #-}
487getIndexUnsafe :: MonadIO m => Int -> MutArray a -> m a
488getIndexUnsafe n MutArray {..} = getIndexUnsafeWith arrContents# (n + arrStart)
490-- | /O(1)/ Lookup the element at the given index. Index starts from 0.
492{-# INLINE getIndex #-}
493getIndex :: MonadIO m => Int -> MutArray a -> m (Maybe a)
494getIndex i arr@MutArray {..} =
495 if i >= 0 && i < arrLen
496 then Just <$> getIndexUnsafe i arr
499-------------------------------------------------------------------------------
501-------------------------------------------------------------------------------
503-- XXX We can also get immutable slices.
505-- | /O(1)/ Slice an array in constant time.
507-- Unsafe: The bounds of the slice are not checked.
512{-# INLINE getSliceUnsafe #-}
514 :: Int -- ^ from index
515 -> Int -- ^ length of the slice
518getSliceUnsafe index len arr@MutArray {..} =
519 assert (index >= 0 && len >= 0 && index + len <= arrLen)
520 $ arr {arrStart = arrStart + index, arrLen = len}
522-- | /O(1)/ Slice an array in constant time. Throws an error if the slice
523-- extends out of the array bounds.
526{-# INLINE getSlice #-}
528 :: Int -- ^ from index
529 -> Int -- ^ length of the slice
532getSlice index len arr@MutArray{..} =
533 if index >= 0 && len >= 0 && index + len <= arrLen
534 then arr {arrStart = arrStart + index, arrLen = len}
536 $ "getSlice: invalid slice, index "
537 ++ show index ++ " length " ++ show len
539-------------------------------------------------------------------------------
540-- to Lists and streams
541-------------------------------------------------------------------------------
543-- XXX Maybe faster to create a list explicitly instead of mapM, if list fusion
544-- does not work well.
546-- | Convert an 'Array' into a list.
550toList :: MonadIO m => MutArray a -> m [a]
551toList arr@MutArray{..} = mapM (`getIndexUnsafe` arr) [0 .. (arrLen - 1)]
553-- | Generates a stream from the elements of a @MutArray@.
555-- >>> read = Stream.unfold MutArray.reader
557{-# INLINE_NORMAL read #-}
558read :: MonadIO m => MutArray a -> D.Stream m a
559read arr@MutArray{..} =
560 D.mapM (`getIndexUnsafe` arr) $ D.enumerateFromToIntegral 0 (arrLen - 1)
562-- Check equivalence with StreamK.fromStream . toStreamD and remove
563{-# INLINE toStreamK #-}
564toStreamK :: MonadIO m => MutArray a -> K.StreamK m a
565toStreamK arr@MutArray{..} = K.unfoldrM step 0
570 | i == arrLen = return Nothing
572 x <- getIndexUnsafe i arr
573 return $ Just (x, i + 1)
575{-# INLINE_NORMAL readRev #-}
576readRev :: MonadIO m => MutArray a -> D.Stream m a
577readRev arr@MutArray{..} =
578 D.mapM (`getIndexUnsafe` arr)
579 $ D.enumerateFromThenToIntegral (arrLen - 1) (arrLen - 2) 0
581-------------------------------------------------------------------------------
583-------------------------------------------------------------------------------
585-- XXX deduplicate this across unboxed array and this module?
587-- | The default chunk size by which the array creation routines increase the
588-- size of the array when the array is grown linearly.
592-- | Like 'writeN' but does not check the array bounds when writing. The fold
593-- driver must not call the step function more than 'n' times otherwise it will
594-- corrupt the memory and crash. This function exists mainly because any
595-- conditional in the step function blocks fusion causing 10x performance
599{-# INLINE_NORMAL writeNUnsafe #-}
600writeNUnsafe :: MonadIO m => Int -> Fold m a (MutArray a)
601writeNUnsafe n = Fold step initial return return
605 initial = FL.Partial <$> new (max n 0)
607 step arr x = FL.Partial <$> snocUnsafe arr x
609-- | @writeN n@ folds a maximum of @n@ elements from the input stream to an
612-- >>> writeN n = Fold.take n (MutArray.writeNUnsafe n)
615{-# INLINE_NORMAL writeN #-}
616writeN :: MonadIO m => Int -> Fold m a (MutArray a)
617writeN n = FL.take n $ writeNUnsafe n
619-- >>> f n = MutArray.writeAppendWith (* 2) (MutArray.pinnedNew n)
620-- >>> writeWith n = Fold.rmapM MutArray.rightSize (f n)
621-- >>> writeWith n = Fold.rmapM MutArray.fromArrayStreamK (MutArray.writeChunks n)
623-- | @writeWith minCount@ folds the whole input to a single array. The array
624-- starts at a size big enough to hold minCount elements, the size is doubled
625-- every time the array needs to be grown.
627-- /Caution! Do not use this on infinite streams./
630{-# INLINE_NORMAL writeWith #-}
631writeWith :: MonadIO m => Int -> Fold m a (MutArray a)
632-- writeWith n = FL.rmapM rightSize $ writeAppendWith (* 2) (pinnedNew n)
633writeWith elemCount = FL.rmapM extract $ FL.foldlM' step initial
638 when (elemCount < 0) $ error "writeWith: elemCount is negative"
639 liftIO $ new elemCount
641 step arr@(MutArray _ start end bound) x
643 let oldSize = end - start
644 newSize = max (oldSize * 2) 1
645 arr1 <- liftIO $ realloc newSize arr
647 step arr x = snocUnsafe arr x
649 -- extract = liftIO . rightSize
652-- | Fold the whole input to a single array.
654-- Same as 'writeWith' using an initial array size of 'arrayChunkSize' bytes
655-- rounded up to the element size.
657-- /Caution! Do not use this on infinite streams./
660write :: MonadIO m => Fold m a (MutArray a)
661write = writeWith arrayChunkSize
663-- | Create a 'MutArray' from the first @n@ elements of a stream. The
664-- array is allocated to size @n@, if the stream terminates before @n@
665-- elements then the array may hold less than @n@ elements.
667{-# INLINE fromStreamN #-}
668fromStreamN :: MonadIO m => Int -> Stream m a -> m (MutArray a)
669fromStreamN n = D.fold (writeN n)
671{-# INLINE fromStream #-}
672fromStream :: MonadIO m => Stream m a -> m (MutArray a)
673fromStream = D.fold write
675{-# INLINABLE fromListN #-}
676fromListN :: MonadIO m => Int -> [a] -> m (MutArray a)
677fromListN n xs = fromStreamN n $ D.fromList xs
679{-# INLINABLE fromList #-}
680fromList :: MonadIO m => [a] -> m (MutArray a)
681fromList xs = fromStream $ D.fromList xs
683{-# INLINABLE fromPureStream #-}
684fromPureStream :: MonadIO m => Stream Identity a -> m (MutArray a)
686 liftIO $ D.fold write $ D.morphInner (return . runIdentity) xs
688-------------------------------------------------------------------------------
690-------------------------------------------------------------------------------
692data GroupState s a start end bound
694 | GroupBuffer s (MutableArray# RealWorld a) start end bound
696 (MutableArray# RealWorld a)
700 (GroupState s a start end bound)
703-- | @chunksOf n stream@ groups the input stream into a stream of
706-- @chunksOf n = foldMany (MutArray.writeN n)@
709{-# INLINE_NORMAL chunksOf #-}
710chunksOf :: forall m a. MonadIO m
711 => Int -> D.Stream m a -> D.Stream m (MutArray a)
712-- XXX the idiomatic implementation leads to large regression in the D.reverse'
713-- benchmark. It seems it has difficulty producing optimized code when
714-- converting to StreamK. Investigate GHC optimizations.
715-- chunksOf n = D.foldMany (writeN n)
716chunksOf n (D.Stream step state) =
717 D.Stream step' (GroupStart state)
722 -- end and len are always equal
724 {-# INLINE_LATE step' #-}
725 step' _ (GroupStart st) = do
727 -- XXX we can pass the module string from the higher level API
728 error $ "Streamly.Internal.Data.Array.Generic.Mut.Type.chunksOf: "
729 ++ "the size of arrays [" ++ show n
730 ++ "] must be a natural number"
731 (MutArray contents start end bound :: MutArray a) <- new n
732 return $ D.Skip (GroupBuffer st contents start end bound)
734 step' gst (GroupBuffer st contents start end bound) = do
735 r <- step (adaptState gst) st
738 putIndexUnsafeWith end contents x
744 contents start end1 bound (GroupStart s))
745 else D.Skip (GroupBuffer s contents start end1 bound)
747 return $ D.Skip (GroupBuffer s contents start end bound)
750 $ D.Skip (GroupYield contents start end bound GroupFinish)
752 step' _ (GroupYield contents start end bound next) =
753 return $ D.Yield (MutArray contents start end bound) next
755 step' _ GroupFinish = return D.Stop
757-------------------------------------------------------------------------------
759-------------------------------------------------------------------------------
761-- | Resumable unfold of an array.
763{-# INLINE_NORMAL producerWith #-}
764producerWith :: Monad m => (forall b. IO b -> m b) -> Producer m (MutArray a) a
765producerWith liftio = Producer step inject extract
769 {-# INLINE inject #-}
770 inject arr = return (arr, 0)
772 {-# INLINE extract #-}
774 return $ arr {arrStart = arrStart arr + i, arrLen = arrLen arr - i}
776 {-# INLINE_LATE step #-}
778 | assert (arrLen arr >= 0) (i == arrLen arr) = return D.Stop
780 x <- liftio $ getIndexUnsafe i arr
781 return $ D.Yield x (arr, i + 1)
783-- | Resumable unfold of an array.
785{-# INLINE_NORMAL producer #-}
786producer :: MonadIO m => Producer m (MutArray a) a
787producer = producerWith liftIO
789-- | Unfold an array into a stream.
791{-# INLINE_NORMAL reader #-}
792reader :: MonadIO m => Unfold m (MutArray a) a
793reader = Producer.simplify producer
795--------------------------------------------------------------------------------
797--------------------------------------------------------------------------------
799-- | Put a sub range of a source array into a subrange of a destination array.
800-- This is not safe as it does not check the bounds.
801{-# INLINE putSliceUnsafe #-}
802putSliceUnsafe :: MonadIO m =>
803 MutArray a -> Int -> MutArray a -> Int -> Int -> m ()
804putSliceUnsafe src srcStart dst dstStart len = liftIO $ do
805 assertM(len <= arrLen dst)
806 assertM(len <= arrLen src)
807 let !(I# srcStart#) = srcStart + arrStart src
808 !(I# dstStart#) = dstStart + arrStart dst
810 let arrS# = arrContents# src
811 arrD# = arrContents# dst
812 IO $ \s# -> (# copyMutableArray#
813 arrS# srcStart# arrD# dstStart# len# s#
817clone :: MonadIO m => MutArray a -> m (MutArray a)
818clone src = liftIO $ do
821 putSliceUnsafe src 0 dst 0 len
824-------------------------------------------------------------------------------
826-------------------------------------------------------------------------------
829length :: MutArray a -> Int
832-------------------------------------------------------------------------------
834-------------------------------------------------------------------------------
836-- | Compare the length of the arrays. If the length is equal, compare the
837-- lexicographical ordering of two underlying byte arrays otherwise return the
838-- result of length comparison.
842cmp :: (MonadIO m, Ord a) => MutArray a -> MutArray a -> m Ordering
844 case compare lenA1 lenA2 of
845 EQ -> loop (lenA1 - 1)
856 v1 <- getIndexUnsafe i a1
857 v2 <- getIndexUnsafe i a2
858 case compare v1 v2 of
863eq :: (MonadIO m, Eq a) => MutArray a -> MutArray a -> m Bool
866 then loop (lenA1 - 1)
875 | i < 0 = return True
877 v1 <- getIndexUnsafe i a1
878 v2 <- getIndexUnsafe i a2
884strip :: MonadIO m => (a -> Bool) -> MutArray a -> m (MutArray a)
885strip p arr = liftIO $ do
886 let lastIndex = length arr - 1
887 indexR <- getIndexR lastIndex -- last predicate failing index
891 indexL <- getIndexL 0 -- first predicate failing index
892 if indexL == 0 && indexR == lastIndex
895 let newLen = indexR - indexL + 1
896 in return $ getSliceUnsafe indexL newLen arr
901 | idx < 0 = return idx
903 r <- getIndexUnsafe idx arr
905 then getIndexR (idx - 1)
909 r <- getIndexUnsafe idx arr
911 then getIndexL (idx + 1)
keysesc to close
- j k
- move down, up
- enter
- open
- esc
- back, or up a level
- 1 2 3 4
- readme, files, commits, refs
- /
- filter
- y
- copy the clone command
- c
- compare the chosen commit with another (or shift-click both)
- ?
- show or hide this list