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
111 -- * In-place Mutation Algorithms
126 -- * Arrays of arrays
127 -- We can add dimensionality parameter to the array type to get
128 -- multidimensional arrays. Multidimensional arrays would just be a
129 -- convenience wrapper on top of single dimensional arrays.
131 -- | Operations dealing with multiple arrays, streams of arrays or
132 -- multidimensional array representations.
134 -- ** Construct from streams
136 -- , arrayStreamKFromStreamD
139 -- ** Eliminate to streams
141 -- , flattenArraysRev
142 -- , fromArrayStreamK
144 -- ** Construct from arrays
145 -- get chunks without copying
148 -- , getSlicesFromLenN
149 -- , splitAt -- XXX should be able to express using getSlice
152 -- ** Appending arrays
168import Control.Monad (when)
169import Control.Monad.IO.Class (MonadIO(..))
170import Data.Functor.Identity (Identity(..))
179import GHC.IO (IO(..))
180import GHC.Int (Int(..))
181import Streamly.Internal.Data.Fold.Type (Fold(..))
182import Streamly.Internal.Data.Producer.Type (Producer (..))
183import Streamly.Internal.Data.Unfold.Type (Unfold(..))
184import Streamly.Internal.Data.Stream.StreamD.Type (Stream)
186import qualified Streamly.Internal.Data.Fold.Type as FL
187import qualified Streamly.Internal.Data.Producer as Producer
188import qualified Streamly.Internal.Data.Stream.StreamD.Type as D
189import qualified Streamly.Internal.Data.Stream.StreamD.Generate as D
190import qualified Streamly.Internal.Data.Stream.StreamD.Lift as D
191import qualified Streamly.Internal.Data.Stream.StreamK.Type as K
193import Prelude hiding (read, length)
195#include "DocTestDataMutArrayGeneric.hs"
197-------------------------------------------------------------------------------
199-------------------------------------------------------------------------------
203 { arrContents# :: MutableArray# RealWorld a
204 -- ^ The internal contents of the array representing the entire array.
206 , arrStart :: {-# UNPACK #-}!Int
207 -- ^ The starting index of this slice.
209 , arrLen :: {-# UNPACK #-}!Int
210 -- ^ The length of this slice.
212 , arrTrueLen :: {-# UNPACK #-}!Int
213 -- ^ This is the true length of the array. Coincidentally, this also
214 -- represents the first index beyond the maximum acceptable index of
215 -- the array. This is specific to the array contents itself and not
216 -- dependent on the slice. This value should not change and is shared
217 -- across all the slices.
220{-# INLINE bottomElement #-}
226 , "This is the bottom element of the array."
227 , "This is a place holder and should never be reached!"
232 funcName = "Streamly.Internal.Data.MutArray.Generic.bottomElement:"
234-- XXX Would be nice if GHC can provide something like newUninitializedArray# so
235-- that we do not have to write undefined or error in the whole array.
237-- | @new count@ allocates a zero length array that can be extended to hold
238-- up to 'count' items without reallocating.
242new :: MonadIO m => Int -> m (MutArray a)
247 case newArray# n# bottomElement s# of
249 let ma = MutArray arr# 0 0 n
252-- XXX This could be pure?
257-- >>> nil = MutArray.new 0
259nil :: MonadIO m => m (MutArray a)
262-------------------------------------------------------------------------------
264-------------------------------------------------------------------------------
266-- | Write the given element to the given index of the array. Does not check if
267-- the index is out of bounds of the array.
270{-# INLINE putIndexUnsafe #-}
271putIndexUnsafe :: forall m a. MonadIO m => Int -> MutArray a -> a -> m ()
272putIndexUnsafe i MutArray {..} x =
273 assert (i >= 0 && i < arrLen)
279 let s1# = writeArray# arrContents# n# x s#
282invalidIndex :: String -> Int -> a
283invalidIndex label i =
284 error $ label ++ ": invalid array index " ++ show i
286-- | /O(1)/ Write the given element at the given index in the array.
287-- Performs in-place mutation of the array.
289-- >>> putIndex ix arr val = MutArray.modifyIndex ix arr (const (val, ()))
292{-# INLINE putIndex #-}
293putIndex :: MonadIO m => Int -> MutArray a -> a -> m ()
294putIndex i arr@MutArray {..} x =
295 if i >= 0 && i < arrLen
296 then putIndexUnsafe i arr x
297 else invalidIndex "putIndex" i
299-- | Write an input stream of (index, value) pairs to an array. Throws an
300-- error if any index is out of bounds.
303{-# INLINE putIndices #-}
304putIndices :: MonadIO m
305 => MutArray a -> Fold m (Int, a) ()
306putIndices arr = FL.foldlM' step (return ())
310 step () (i, x) = liftIO (putIndex i arr x)
312-- | Modify a given index of an array using a modifier function without checking
315-- Unsafe because it does not check the bounds of the array.
318modifyIndexUnsafe :: MonadIO m => Int -> MutArray a -> (a -> (a, b)) -> m b
319modifyIndexUnsafe i MutArray {..} f = do
325 case readArray# arrContents# n# s# of
328 s2# = writeArray# arrContents# n# a1 s1#
331-- | Modify a given index of an array using a modifier function.
334modifyIndex :: MonadIO m => Int -> MutArray a -> (a -> (a, b)) -> m b
335modifyIndex i arr@MutArray {..} f = do
336 if i >= 0 && i < arrLen
337 then modifyIndexUnsafe i arr f
338 else invalidIndex "modifyIndex" i
340-------------------------------------------------------------------------------
342-------------------------------------------------------------------------------
344-- | Reallocates the array according to the new size. This is a safe function
345-- that always creates a new array and copies the old array into the new one.
346-- If the reallocated size is less than the original array it results in a
347-- truncated version of the original array.
349realloc :: MonadIO m => Int -> MutArray a -> m (MutArray a)
352 let !newLen@(I# newLen#) = min n (arrLen arr)
353 !(I# arrS#) = arrStart arr
354 !(I# arr1S#) = arrStart arr1
355 arrC# = arrContents# arr
356 arr1C# = arrContents# arr1
360 let s1# = copyMutableArray# arrC# arrS# arr1C# arr1S# newLen# s#
361 in (# s1#, arr1 {arrLen = newLen, arrTrueLen = n} #)
364 MonadIO m => String -> (Int -> Int) -> Int -> MutArray a -> m (MutArray a)
365reallocWith label sizer reqSize arr = do
366 let oldSize = arrLen arr
367 newSize = sizer oldSize
368 safeSize = max newSize (oldSize + reqSize)
369 assert (newSize >= oldSize + reqSize || error badSize) (return ())
376 , ": new array size is less than required size "
378 , ". Please check the sizing function passed."
381-------------------------------------------------------------------------------
383-------------------------------------------------------------------------------
385-- XXX Not sure of the behavior of writeArray# if we specify an index which is
386-- out of bounds. This comment should be rewritten based on that.
387-- | Really really unsafe, appends the element into the first array, may
388-- cause silent data corruption or if you are lucky a segfault if the index
392{-# INLINE snocUnsafe #-}
393snocUnsafe :: MonadIO m => MutArray a -> a -> m (MutArray a)
394snocUnsafe arr@MutArray {..} a = do
395 assert (arrStart + arrLen < arrTrueLen) (return ())
396 let arr1 = arr {arrLen = arrLen + 1}
397 putIndexUnsafe arrLen arr1 a
400-- NOINLINE to move it out of the way and not pollute the instruction cache.
401{-# NOINLINE snocWithRealloc #-}
402snocWithRealloc :: MonadIO m => (Int -> Int) -> MutArray a -> a -> m (MutArray a)
403snocWithRealloc sizer arr x = do
404 arr1 <- reallocWith "snocWithRealloc" sizer 1 arr
407-- | @snocWith sizer arr elem@ mutates @arr@ to append @elem@. The length of
408-- the array increases by 1.
410-- If there is no reserved space available in @arr@ it is reallocated to a size
411-- in bytes determined by the @sizer oldSize@ function, where @oldSize@ is the
412-- original size of the array.
414-- Note that the returned array may be a mutated version of the original array.
417{-# INLINE snocWith #-}
418snocWith :: MonadIO m => (Int -> Int) -> MutArray a -> a -> m (MutArray a)
419snocWith sizer arr@MutArray {..} x = do
420 if arrStart + arrLen < arrTrueLen
421 then snocUnsafe arr x
422 else snocWithRealloc sizer arr x
424-- XXX round it to next power of 2.
426-- | The array is mutated to append an additional element to it. If there is no
427-- reserved space available in the array then it is reallocated to double the
430-- This is useful to reduce allocations when appending unknown number of
433-- Note that the returned array may be a mutated version of the original array.
435-- >>> snoc = MutArray.snocWith (* 2)
437-- Performs O(n * log n) copies to grow, but is liberal with memory allocation.
441snoc :: MonadIO m => MutArray a -> a -> m (MutArray a)
444-- | Make the uninitialized memory in the array available for use extending it
445-- by the supplied length beyond the current length of the array. The array may
449uninit :: MonadIO m => MutArray a -> Int -> m (MutArray a)
450uninit arr@MutArray{..} len =
451 if arrStart + arrLen + len <= arrTrueLen
452 then return $ arr {arrLen = arrLen + len}
453 else realloc (arrLen + len) arr
455-------------------------------------------------------------------------------
457-------------------------------------------------------------------------------
459-- | Return the element at the specified index without checking the bounds.
461-- Unsafe because it does not check the bounds of the array.
462{-# INLINE_NORMAL getIndexUnsafe #-}
463getIndexUnsafe :: MonadIO m => Int -> MutArray a -> m a
464getIndexUnsafe n MutArray {..} =
468 let !(I# i#) = arrStart + n
469 in readArray# arrContents# i# s#
471-- | /O(1)/ Lookup the element at the given index. Index starts from 0.
473{-# INLINE getIndex #-}
474getIndex :: MonadIO m => Int -> MutArray a -> m (Maybe a)
475getIndex i arr@MutArray {..} =
476 if i >= 0 && i < arrLen
477 then Just <$> getIndexUnsafe i arr
480-------------------------------------------------------------------------------
482-------------------------------------------------------------------------------
484-- XXX We can also get immutable slices.
486-- | /O(1)/ Slice an array in constant time.
488-- Unsafe: The bounds of the slice are not checked.
493{-# INLINE getSliceUnsafe #-}
495 :: Int -- ^ from index
496 -> Int -- ^ length of the slice
499getSliceUnsafe index len arr@MutArray {..} =
500 assert (index >= 0 && len >= 0 && index + len <= arrLen)
501 $ arr {arrStart = arrStart + index, arrLen = len}
503-- | /O(1)/ Slice an array in constant time. Throws an error if the slice
504-- extends out of the array bounds.
507{-# INLINE getSlice #-}
509 :: Int -- ^ from index
510 -> Int -- ^ length of the slice
513getSlice index len arr@MutArray{..} =
514 if index >= 0 && len >= 0 && index + len <= arrLen
515 then arr {arrStart = arrStart + index, arrLen = len}
517 $ "getSlice: invalid slice, index "
518 ++ show index ++ " length " ++ show len
520-------------------------------------------------------------------------------
521-- to Lists and streams
522-------------------------------------------------------------------------------
524-- XXX Maybe faster to create a list explicitly instead of mapM, if list fusion
525-- does not work well.
527-- | Convert an 'Array' into a list.
531toList :: MonadIO m => MutArray a -> m [a]
532toList arr@MutArray{..} = mapM (`getIndexUnsafe` arr) [0 .. (arrLen - 1)]
534-- | Generates a stream from the elements of a @MutArray@.
536-- >>> read = Stream.unfold MutArray.reader
538{-# INLINE_NORMAL read #-}
539read :: MonadIO m => MutArray a -> D.Stream m a
540read arr@MutArray{..} =
541 D.mapM (`getIndexUnsafe` arr) $ D.enumerateFromToIntegral 0 (arrLen - 1)
543-- Check equivalence with StreamK.fromStream . toStreamD and remove
544{-# INLINE toStreamK #-}
545toStreamK :: MonadIO m => MutArray a -> K.StreamK m a
546toStreamK arr@MutArray{..} = K.unfoldrM step 0
551 | i == arrLen = return Nothing
553 x <- getIndexUnsafe i arr
554 return $ Just (x, i + 1)
556{-# INLINE_NORMAL readRev #-}
557readRev :: MonadIO m => MutArray a -> D.Stream m a
558readRev arr@MutArray{..} =
559 D.mapM (`getIndexUnsafe` arr)
560 $ D.enumerateFromThenToIntegral (arrLen - 1) (arrLen - 2) 0
562-------------------------------------------------------------------------------
564-------------------------------------------------------------------------------
566-- XXX deduplicate this across unboxed array and this module?
568-- | The default chunk size by which the array creation routines increase the
569-- size of the array when the array is grown linearly.
573-- | Like 'writeN' but does not check the array bounds when writing. The fold
574-- driver must not call the step function more than 'n' times otherwise it will
575-- corrupt the memory and crash. This function exists mainly because any
576-- conditional in the step function blocks fusion causing 10x performance
580{-# INLINE_NORMAL writeNUnsafe #-}
581writeNUnsafe :: MonadIO m => Int -> Fold m a (MutArray a)
582writeNUnsafe n = Fold step initial return
586 initial = FL.Partial <$> new (max n 0)
588 step arr x = FL.Partial <$> snocUnsafe arr x
590-- | @writeN n@ folds a maximum of @n@ elements from the input stream to an
593-- >>> writeN n = Fold.take n (MutArray.writeNUnsafe n)
596{-# INLINE_NORMAL writeN #-}
597writeN :: MonadIO m => Int -> Fold m a (MutArray a)
598writeN n = FL.take n $ writeNUnsafe n
600-- >>> f n = MutArray.writeAppendWith (* 2) (MutArray.pinnedNew n)
601-- >>> writeWith n = Fold.rmapM MutArray.rightSize (f n)
602-- >>> writeWith n = Fold.rmapM MutArray.fromArrayStreamK (MutArray.writeChunks n)
604-- | @writeWith minCount@ folds the whole input to a single array. The array
605-- starts at a size big enough to hold minCount elements, the size is doubled
606-- every time the array needs to be grown.
608-- /Caution! Do not use this on infinite streams./
611{-# INLINE_NORMAL writeWith #-}
612writeWith :: MonadIO m => Int -> Fold m a (MutArray a)
613-- writeWith n = FL.rmapM rightSize $ writeAppendWith (* 2) (pinnedNew n)
614writeWith elemCount = FL.rmapM extract $ FL.foldlM' step initial
619 when (elemCount < 0) $ error "writeWith: elemCount is negative"
620 liftIO $ new elemCount
622 step arr@(MutArray _ start end bound) x
624 let oldSize = end - start
625 newSize = max (oldSize * 2) 1
626 arr1 <- liftIO $ realloc newSize arr
628 step arr x = snocUnsafe arr x
630 -- extract = liftIO . rightSize
633-- | Fold the whole input to a single array.
635-- Same as 'writeWith' using an initial array size of 'arrayChunkSize' bytes
636-- rounded up to the element size.
638-- /Caution! Do not use this on infinite streams./
641write :: MonadIO m => Fold m a (MutArray a)
642write = writeWith arrayChunkSize
644-- | Create a 'MutArray' from the first @n@ elements of a stream. The
645-- array is allocated to size @n@, if the stream terminates before @n@
646-- elements then the array may hold less than @n@ elements.
648{-# INLINE fromStreamN #-}
649fromStreamN :: MonadIO m => Int -> Stream m a -> m (MutArray a)
650fromStreamN n = D.fold (writeN n)
652{-# INLINE fromStream #-}
653fromStream :: MonadIO m => Stream m a -> m (MutArray a)
654fromStream = D.fold write
656{-# INLINABLE fromListN #-}
657fromListN :: MonadIO m => Int -> [a] -> m (MutArray a)
658fromListN n xs = fromStreamN n $ D.fromList xs
660{-# INLINABLE fromList #-}
661fromList :: MonadIO m => [a] -> m (MutArray a)
662fromList xs = fromStream $ D.fromList xs
664{-# INLINABLE fromPureStream #-}
665fromPureStream :: MonadIO m => Stream Identity a -> m (MutArray a)
667 liftIO $ D.fold write $ D.morphInner (return . runIdentity) xs
669-------------------------------------------------------------------------------
671-------------------------------------------------------------------------------
673-- | Resumable unfold of an array.
675{-# INLINE_NORMAL producerWith #-}
676producerWith :: Monad m => (forall b. IO b -> m b) -> Producer m (MutArray a) a
677producerWith liftio = Producer step inject extract
681 {-# INLINE inject #-}
682 inject arr = return (arr, 0)
684 {-# INLINE extract #-}
686 return $ arr {arrStart = arrStart arr + i, arrLen = arrLen arr - i}
688 {-# INLINE_LATE step #-}
690 | assert (arrLen arr >= 0) (i == arrLen arr) = return D.Stop
692 x <- liftio $ getIndexUnsafe i arr
693 return $ D.Yield x (arr, i + 1)
695-- | Resumable unfold of an array.
697{-# INLINE_NORMAL producer #-}
698producer :: MonadIO m => Producer m (MutArray a) a
699producer = producerWith liftIO
701-- | Unfold an array into a stream.
703{-# INLINE_NORMAL reader #-}
704reader :: MonadIO m => Unfold m (MutArray a) a
705reader = Producer.simplify producer
707--------------------------------------------------------------------------------
709--------------------------------------------------------------------------------
711-- | Put a sub range of a source array into a subrange of a destination array.
712-- This is not safe as it does not check the bounds.
713{-# INLINE putSliceUnsafe #-}
714putSliceUnsafe :: MonadIO m =>
715 MutArray a -> Int -> MutArray a -> Int -> Int -> m ()
716putSliceUnsafe src srcStart dst dstStart len = liftIO $ do
717 assertM(len <= arrLen dst)
718 assertM(len <= arrLen src)
719 let !(I# srcStart#) = srcStart + arrStart src
720 !(I# dstStart#) = dstStart + arrStart dst
722 let arrS# = arrContents# src
723 arrD# = arrContents# dst
724 IO $ \s# -> (# copyMutableArray#
725 arrS# srcStart# arrD# dstStart# len# s#
729clone :: MonadIO m => MutArray a -> m (MutArray a)
730clone src = liftIO $ do
733 putSliceUnsafe src 0 dst 0 len
736-------------------------------------------------------------------------------
738-------------------------------------------------------------------------------
741length :: MutArray a -> Int
744-------------------------------------------------------------------------------
746-------------------------------------------------------------------------------
748-- | Compare the length of the arrays. If the length is equal, compare the
749-- lexicographical ordering of two underlying byte arrays otherwise return the
750-- result of length comparison.
754cmp :: (MonadIO m, Ord a) => MutArray a -> MutArray a -> m Ordering
756 case compare lenA1 lenA2 of
757 EQ -> loop (lenA1 - 1)
768 v1 <- getIndexUnsafe i a1
769 v2 <- getIndexUnsafe i a2
770 case compare v1 v2 of
775eq :: (MonadIO m, Eq a) => MutArray a -> MutArray a -> m Bool
778 then loop (lenA1 - 1)
787 | i < 0 = return True
789 v1 <- getIndexUnsafe i a1
790 v2 <- getIndexUnsafe i a2
796strip :: MonadIO m => (a -> Bool) -> MutArray a -> m (MutArray a)
797strip p arr = liftIO $ do
798 let lastIndex = length arr - 1
799 indexR <- getIndexR lastIndex -- last predicate failing index
803 indexL <- getIndexL 0 -- first predicate failing index
804 if indexL == 0 && indexR == lastIndex
807 let newLen = indexR - indexL + 1
808 in return $ getSliceUnsafe indexL newLen arr
813 | idx < 0 = return idx
815 r <- getIndexUnsafe idx arr
817 then getIndexR (idx - 1)
821 r <- getIndexUnsafe idx arr
823 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