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
28 , createWith -- createOfMin/createMin/createGE?
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
176import Control.Monad (when)
177import Control.Monad.IO.Class (MonadIO(..))
178import Data.Functor.Identity (Identity(..))
187import GHC.IO (IO(..))
188import GHC.Int (Int(..))
189import Streamly.Internal.Data.Fold.Type (Fold(..))
190import Streamly.Internal.Data.Producer.Type (Producer (..))
191import Streamly.Internal.Data.Unfold.Type (Unfold(..))
192import Streamly.Internal.Data.Stream.Type (Stream)
193import Streamly.Internal.Data.SVar.Type (adaptState)
195import qualified Streamly.Internal.Data.Fold.Type as FL
196import qualified Streamly.Internal.Data.Producer as Producer
197import qualified Streamly.Internal.Data.Stream.Type as D
198import qualified Streamly.Internal.Data.Stream.Generate as D
199import qualified Streamly.Internal.Data.Stream.Lift as D
200import qualified Streamly.Internal.Data.StreamK.Type as K
202import Prelude hiding (read, length)
204#include "DocTestDataMutArrayGeneric.hs"
206-------------------------------------------------------------------------------
208-------------------------------------------------------------------------------
212 { arrContents# :: MutableArray# RealWorld a
213 -- ^ The internal contents of the array representing the entire array.
215 , arrStart :: {-# UNPACK #-}!Int
216 -- ^ The starting index of this slice.
218 , arrLen :: {-# UNPACK #-}!Int
219 -- ^ The length of this slice.
221 , arrTrueLen :: {-# UNPACK #-}!Int
222 -- ^ This is the true length of the array. Coincidentally, this also
223 -- represents the first index beyond the maximum acceptable index of
224 -- the array. This is specific to the array contents itself and not
225 -- dependent on the slice. This value should not change and is shared
226 -- across all the slices.
229{-# INLINE bottomElement #-}
235 , "This is the bottom element of the array."
236 , "This is a place holder and should never be reached!"
241 funcName = "Streamly.Internal.Data.MutArray.Generic.bottomElement:"
243-- XXX Would be nice if GHC can provide something like newUninitializedArray# so
244-- that we do not have to write undefined or error in the whole array.
246-- | @emptyOf count@ allocates a zero length array that can be extended to hold
247-- up to 'count' items without reallocating.
250{-# INLINE emptyOf #-}
251emptyOf :: MonadIO m => Int -> m (MutArray a)
256 case newArray# n# bottomElement s# of
258 let ma = MutArray arr# 0 0 n
261-- XXX Deprecate in major
262-- {-# DEPRECATED new "Please use emptyOf instead." #-}
264new :: MonadIO m => Int -> m (MutArray a)
267-- XXX This could be pure?
272-- >>> nil = MutArray.new 0
274nil :: MonadIO m => m (MutArray a)
277-------------------------------------------------------------------------------
279-------------------------------------------------------------------------------
281-- | Write the given element to the given index of the 'MutableArray#'. Does not
282-- check if the index is out of bounds of the array.
285{-# INLINE putIndexUnsafeWith #-}
286putIndexUnsafeWith :: MonadIO m => Int -> MutableArray# RealWorld a -> a -> m ()
287putIndexUnsafeWith n _arrContents# x =
293 let s1# = writeArray# _arrContents# n# x s#
296-- | Write the given element to the given index of the array. Does not check if
297-- the index is out of bounds of the array.
300{-# INLINE putIndexUnsafe #-}
301putIndexUnsafe :: forall m a. MonadIO m => Int -> MutArray a -> a -> m ()
302putIndexUnsafe i MutArray {..} x =
303 assert (i >= 0 && i < arrLen)
304 putIndexUnsafeWith (i + arrStart) arrContents# x
306invalidIndex :: String -> Int -> a
307invalidIndex label i =
308 error $ label ++ ": invalid array index " ++ show i
310-- | /O(1)/ Write the given element at the given index in the array.
311-- Performs in-place mutation of the array.
313-- >>> putIndex ix arr val = MutArray.modifyIndex ix arr (const (val, ()))
316{-# INLINE putIndex #-}
317putIndex :: MonadIO m => Int -> MutArray a -> a -> m ()
318putIndex i arr@MutArray {..} x =
319 if i >= 0 && i < arrLen
320 then putIndexUnsafe i arr x
321 else invalidIndex "putIndex" i
323-- | Write an input stream of (index, value) pairs to an array. Throws an
324-- error if any index is out of bounds.
327{-# INLINE putIndices #-}
328putIndices :: MonadIO m
329 => MutArray a -> Fold m (Int, a) ()
330putIndices arr = FL.foldlM' step (return ())
334 step () (i, x) = putIndex i arr x
336-- | Modify a given index of an array using a modifier function without checking
339-- Unsafe because it does not check the bounds of the array.
342modifyIndexUnsafe :: MonadIO m => Int -> MutArray a -> (a -> (a, b)) -> m b
343modifyIndexUnsafe i MutArray {..} f = do
349 case readArray# arrContents# n# s# of
352 s2# = writeArray# arrContents# n# a1 s1#
355-- | Modify a given index of an array using a modifier function.
358modifyIndex :: MonadIO m => Int -> MutArray a -> (a -> (a, b)) -> m b
359modifyIndex i arr@MutArray {..} f = do
360 if i >= 0 && i < arrLen
361 then modifyIndexUnsafe i arr f
362 else invalidIndex "modifyIndex" i
364-------------------------------------------------------------------------------
366-------------------------------------------------------------------------------
368-- | Reallocates the array according to the new size. This is a safe function
369-- that always creates a new array and copies the old array into the new one.
370-- If the reallocated size is less than the original array it results in a
371-- truncated version of the original array.
373realloc :: MonadIO m => Int -> MutArray a -> m (MutArray a)
376 let !newLen@(I# newLen#) = min n (arrLen arr)
377 !(I# arrS#) = arrStart arr
378 !(I# arr1S#) = arrStart arr1
379 arrC# = arrContents# arr
380 arr1C# = arrContents# arr1
384 let s1# = copyMutableArray# arrC# arrS# arr1C# arr1S# newLen# s#
385 in (# s1#, arr1 {arrLen = newLen, arrTrueLen = n} #)
388 MonadIO m => String -> (Int -> Int) -> Int -> MutArray a -> m (MutArray a)
389reallocWith label sizer reqSize arr = do
390 let oldSize = arrLen arr
391 newSize = sizer oldSize
392 safeSize = max newSize (oldSize + reqSize)
393 assert (newSize >= oldSize + reqSize || error badSize) (return ())
400 , ": new array size is less than required size "
402 , ". Please check the sizing function passed."
405-------------------------------------------------------------------------------
407-------------------------------------------------------------------------------
409-- XXX Not sure of the behavior of writeArray# if we specify an index which is
410-- out of bounds. This comment should be rewritten based on that.
411-- | Really really unsafe, appends the element into the first array, may
412-- cause silent data corruption or if you are lucky a segfault if the index
416{-# INLINE snocUnsafe #-}
417snocUnsafe :: MonadIO m => MutArray a -> a -> m (MutArray a)
418snocUnsafe arr@MutArray {..} a = do
419 assert (arrStart + arrLen < arrTrueLen) (return ())
420 let arr1 = arr {arrLen = arrLen + 1}
421 putIndexUnsafe arrLen arr1 a
424-- NOINLINE to move it out of the way and not pollute the instruction cache.
425{-# NOINLINE snocWithRealloc #-}
426snocWithRealloc :: MonadIO m => (Int -> Int) -> MutArray a -> a -> m (MutArray a)
427snocWithRealloc sizer arr x = do
428 arr1 <- reallocWith "snocWithRealloc" sizer 1 arr
431-- | @snocWith sizer arr elem@ mutates @arr@ to append @elem@. The length of
432-- the array increases by 1.
434-- If there is no reserved space available in @arr@ it is reallocated to a size
435-- in bytes determined by the @sizer oldSize@ function, where @oldSize@ is the
436-- original size of the array.
438-- Note that the returned array may be a mutated version of the original array.
441{-# INLINE snocWith #-}
442snocWith :: MonadIO m => (Int -> Int) -> MutArray a -> a -> m (MutArray a)
443snocWith sizer arr@MutArray {..} x = do
444 if arrStart + arrLen < arrTrueLen
445 then snocUnsafe arr x
446 else snocWithRealloc sizer arr x
448-- XXX round it to next power of 2.
450-- | The array is mutated to append an additional element to it. If there is no
451-- reserved space available in the array then it is reallocated to double the
454-- This is useful to reduce allocations when appending unknown number of
457-- Note that the returned array may be a mutated version of the original array.
459-- >>> snoc = MutArray.snocWith (* 2)
461-- Performs O(n * log n) copies to grow, but is liberal with memory allocation.
465snoc :: MonadIO m => MutArray a -> a -> m (MutArray a)
468-- | Make the uninitialized memory in the array available for use extending it
469-- by the supplied length beyond the current length of the array. The array may
473uninit :: MonadIO m => MutArray a -> Int -> m (MutArray a)
474uninit arr@MutArray{..} len =
475 if arrStart + arrLen + len <= arrTrueLen
476 then return $ arr {arrLen = arrLen + len}
477 else realloc (arrLen + len) arr
479-------------------------------------------------------------------------------
481-------------------------------------------------------------------------------
483-- | Return the element at the specified index without checking the bounds from
484-- a @MutableArray# RealWorld@.
486-- Unsafe because it does not check the bounds of the array.
487{-# INLINE getIndexUnsafeWith #-}
488getIndexUnsafeWith :: MonadIO m => MutableArray# RealWorld a -> Int -> m a
489getIndexUnsafeWith _arrContents# n =
494 in readArray# _arrContents# i# s#
496-- | Return the element at the specified index without checking the bounds.
498-- Unsafe because it does not check the bounds of the array.
499{-# INLINE_NORMAL getIndexUnsafe #-}
500getIndexUnsafe :: MonadIO m => Int -> MutArray a -> m a
501getIndexUnsafe n MutArray {..} = getIndexUnsafeWith arrContents# (n + arrStart)
503-- | /O(1)/ Lookup the element at the given index. Index starts from 0.
505{-# INLINE getIndex #-}
506getIndex :: MonadIO m => Int -> MutArray a -> m (Maybe a)
507getIndex i arr@MutArray {..} =
508 if i >= 0 && i < arrLen
509 then Just <$> getIndexUnsafe i arr
512-------------------------------------------------------------------------------
514-------------------------------------------------------------------------------
516-- XXX We can also get immutable slices.
518-- | /O(1)/ Slice an array in constant time.
520-- Unsafe: The bounds of the slice are not checked.
525{-# INLINE getSliceUnsafe #-}
527 :: Int -- ^ from index
528 -> Int -- ^ length of the slice
531getSliceUnsafe index len arr@MutArray {..} =
532 assert (index >= 0 && len >= 0 && index + len <= arrLen)
533 $ arr {arrStart = arrStart + index, arrLen = len}
535-- | /O(1)/ Slice an array in constant time. Throws an error if the slice
536-- extends out of the array bounds.
539{-# INLINE getSlice #-}
541 :: Int -- ^ from index
542 -> Int -- ^ length of the slice
545getSlice index len arr@MutArray{..} =
546 if index >= 0 && len >= 0 && index + len <= arrLen
547 then arr {arrStart = arrStart + index, arrLen = len}
549 $ "getSlice: invalid slice, index "
550 ++ show index ++ " length " ++ show len
552-------------------------------------------------------------------------------
553-- to Lists and streams
554-------------------------------------------------------------------------------
556-- XXX Maybe faster to create a list explicitly instead of mapM, if list fusion
557-- does not work well.
559-- | Convert an 'Array' into a list.
563toList :: MonadIO m => MutArray a -> m [a]
564toList arr@MutArray{..} = mapM (`getIndexUnsafe` arr) [0 .. (arrLen - 1)]
566-- | Generates a stream from the elements of a @MutArray@.
568-- >>> read = Stream.unfold MutArray.reader
570{-# INLINE_NORMAL read #-}
571read :: MonadIO m => MutArray a -> D.Stream m a
572read arr@MutArray{..} =
573 D.mapM (`getIndexUnsafe` arr) $ D.enumerateFromToIntegral 0 (arrLen - 1)
575-- Check equivalence with StreamK.fromStream . toStreamD and remove
576{-# INLINE toStreamK #-}
577toStreamK :: MonadIO m => MutArray a -> K.StreamK m a
578toStreamK arr@MutArray{..} = K.unfoldrM step 0
583 | i == arrLen = return Nothing
585 x <- getIndexUnsafe i arr
586 return $ Just (x, i + 1)
588{-# INLINE_NORMAL readRev #-}
589readRev :: MonadIO m => MutArray a -> D.Stream m a
590readRev arr@MutArray{..} =
591 D.mapM (`getIndexUnsafe` arr)
592 $ D.enumerateFromThenToIntegral (arrLen - 1) (arrLen - 2) 0
594-------------------------------------------------------------------------------
596-------------------------------------------------------------------------------
598-- XXX deduplicate this across unboxed array and this module?
600-- | The default chunk size by which the array creation routines increase the
601-- size of the array when the array is grown linearly.
605-- | Like 'createOf' but does not check the array bounds when writing. The fold
606-- driver must not call the step function more than 'n' times otherwise it will
607-- corrupt the memory and crash. This function exists mainly because any
608-- conditional in the step function blocks fusion causing 10x performance
612{-# INLINE_NORMAL unsafeCreateOf #-}
613unsafeCreateOf :: MonadIO m => Int -> Fold m a (MutArray a)
614unsafeCreateOf n = Fold step initial return return
618 initial = FL.Partial <$> new (max n 0)
620 step arr x = FL.Partial <$> snocUnsafe arr x
622{-# DEPRECATED writeNUnsafe "Please use unsafeCreateOf instead." #-}
623{-# INLINE writeNUnsafe #-}
624writeNUnsafe :: MonadIO m => Int -> Fold m a (MutArray a)
625writeNUnsafe = unsafeCreateOf
627-- | @createOf n@ folds a maximum of @n@ elements from the input stream to an
630-- >>> createOf n = Fold.take n (MutArray.unsafeCreateOf n)
633{-# INLINE_NORMAL createOf #-}
634createOf :: MonadIO m => Int -> Fold m a (MutArray a)
635createOf n = FL.take n $ unsafeCreateOf n
637-- XXX Deprecate in major
638-- {-# DEPRECATED writeN "Please use createOf instead." #-}
640writeN :: MonadIO m => Int -> Fold m a (MutArray a)
643-- >>> f n = MutArray.writeAppendWith (* 2) (MutArray.pinnedNew n)
644-- >>> writeWith n = Fold.rmapM MutArray.rightSize (f n)
645-- >>> writeWith n = Fold.rmapM MutArray.fromArrayStreamK (MutArray.writeChunks n)
647-- | @createWith minCount@ folds the whole input to a single array. The array
648-- starts at a size big enough to hold minCount elements, the size is doubled
649-- every time the array needs to be grown.
651-- /Caution! Do not use this on infinite streams./
654{-# INLINE_NORMAL createWith #-}
655createWith :: MonadIO m => Int -> Fold m a (MutArray a)
656-- writeWith n = FL.rmapM rightSize $ writeAppendWith (* 2) (pinnedNew n)
657createWith elemCount = FL.rmapM extract $ FL.foldlM' step initial
662 when (elemCount < 0) $ error "createWith: elemCount is negative"
665 step arr@(MutArray _ start end bound) x
667 let oldSize = end - start
668 newSize = max (oldSize * 2) 1
669 arr1 <- realloc newSize arr
671 step arr x = snocUnsafe arr x
673 -- extract = rightSize
676{-# DEPRECATED writeWith "Please use createWith instead." #-}
677{-# INLINE writeWith #-}
678writeWith :: MonadIO m => Int -> Fold m a (MutArray a)
679writeWith = createWith
681-- | Fold the whole input to a single array.
683-- Same as 'createWith' using an initial array size of 'arrayChunkSize' bytes
684-- rounded up to the element size.
686-- /Caution! Do not use this on infinite streams./
689create :: MonadIO m => Fold m a (MutArray a)
690create = writeWith arrayChunkSize
692-- XXX Deprecate in major
693-- {-# DEPRECATED write "Please use create instead." #-}
695write :: MonadIO m => Fold m a (MutArray a)
698-- | Create a 'MutArray' from the first @n@ elements of a stream. The
699-- array is allocated to size @n@, if the stream terminates before @n@
700-- elements then the array may hold less than @n@ elements.
702{-# INLINE fromStreamN #-}
703fromStreamN :: MonadIO m => Int -> Stream m a -> m (MutArray a)
704fromStreamN n = D.fold (writeN n)
706{-# INLINE fromStream #-}
707fromStream :: MonadIO m => Stream m a -> m (MutArray a)
708fromStream = D.fold write
710{-# INLINABLE fromListN #-}
711fromListN :: MonadIO m => Int -> [a] -> m (MutArray a)
712fromListN n xs = fromStreamN n $ D.fromList xs
714{-# INLINABLE fromList #-}
715fromList :: MonadIO m => [a] -> m (MutArray a)
716fromList xs = fromStream $ D.fromList xs
718{-# INLINABLE fromPureStream #-}
719fromPureStream :: MonadIO m => Stream Identity a -> m (MutArray a)
721 D.fold write $ D.morphInner (return . runIdentity) xs
723-------------------------------------------------------------------------------
725-------------------------------------------------------------------------------
727data GroupState s a start end bound
729 | GroupBuffer s (MutableArray# RealWorld a) start end bound
731 (MutableArray# RealWorld a)
735 (GroupState s a start end bound)
738-- | @chunksOf n stream@ groups the input stream into a stream of
741-- @chunksOf n = foldMany (MutArray.writeN n)@
744{-# INLINE_NORMAL chunksOf #-}
745chunksOf :: forall m a. MonadIO m
746 => Int -> D.Stream m a -> D.Stream m (MutArray a)
747-- XXX the idiomatic implementation leads to large regression in the D.reverse'
748-- benchmark. It seems it has difficulty producing optimized code when
749-- converting to StreamK. Investigate GHC optimizations.
750-- chunksOf n = D.foldMany (writeN n)
751chunksOf n (D.Stream step state) =
752 D.Stream step' (GroupStart state)
757 -- end and len are always equal
759 {-# INLINE_LATE step' #-}
760 step' _ (GroupStart st) = do
762 -- XXX we can pass the module string from the higher level API
763 error $ "Streamly.Internal.Data.Array.Generic.Mut.Type.chunksOf: "
764 ++ "the size of arrays [" ++ show n
765 ++ "] must be a natural number"
766 (MutArray contents start end bound :: MutArray a) <- new n
767 return $ D.Skip (GroupBuffer st contents start end bound)
769 step' gst (GroupBuffer st contents start end bound) = do
770 r <- step (adaptState gst) st
773 putIndexUnsafeWith end contents x
779 contents start end1 bound (GroupStart s))
780 else D.Skip (GroupBuffer s contents start end1 bound)
782 return $ D.Skip (GroupBuffer s contents start end bound)
785 $ D.Skip (GroupYield contents start end bound GroupFinish)
787 step' _ (GroupYield contents start end bound next) =
788 return $ D.Yield (MutArray contents start end bound) next
790 step' _ GroupFinish = return D.Stop
792-------------------------------------------------------------------------------
794-------------------------------------------------------------------------------
796-- | Resumable unfold of an array.
798{-# INLINE_NORMAL producerWith #-}
799producerWith :: Monad m => (forall b. IO b -> m b) -> Producer m (MutArray a) a
800producerWith liftio = Producer step inject extract
804 {-# INLINE inject #-}
805 inject arr = return (arr, 0)
807 {-# INLINE extract #-}
809 return $ arr {arrStart = arrStart arr + i, arrLen = arrLen arr - i}
811 {-# INLINE_LATE step #-}
813 | assert (arrLen arr >= 0) (i == arrLen arr) = return D.Stop
815 x <- liftio $ getIndexUnsafe i arr
816 return $ D.Yield x (arr, i + 1)
818-- | Resumable unfold of an array.
820{-# INLINE_NORMAL producer #-}
821producer :: MonadIO m => Producer m (MutArray a) a
822producer = producerWith liftIO
824-- | Unfold an array into a stream.
826{-# INLINE_NORMAL reader #-}
827reader :: MonadIO m => Unfold m (MutArray a) a
828reader = Producer.simplify producer
830--------------------------------------------------------------------------------
832--------------------------------------------------------------------------------
834-- | Put a sub range of a source array into a subrange of a destination array.
835-- This is not safe as it does not check the bounds.
836{-# INLINE putSliceUnsafe #-}
837putSliceUnsafe :: MonadIO m =>
838 MutArray a -> Int -> MutArray a -> Int -> Int -> m ()
839putSliceUnsafe src srcStart dst dstStart len = liftIO $ do
840 assertM(len <= arrLen dst)
841 assertM(len <= arrLen src)
842 let !(I# srcStart#) = srcStart + arrStart src
843 !(I# dstStart#) = dstStart + arrStart dst
845 let arrS# = arrContents# src
846 arrD# = arrContents# dst
847 IO $ \s# -> (# copyMutableArray#
848 arrS# srcStart# arrD# dstStart# len# s#
852clone :: MonadIO m => MutArray a -> m (MutArray a)
856 putSliceUnsafe src 0 dst 0 len
859-------------------------------------------------------------------------------
861-------------------------------------------------------------------------------
864length :: MutArray a -> Int
867-------------------------------------------------------------------------------
869-------------------------------------------------------------------------------
871-- | Compare the length of the arrays. If the length is equal, compare the
872-- lexicographical ordering of two underlying byte arrays otherwise return the
873-- result of length comparison.
877cmp :: (MonadIO m, Ord a) => MutArray a -> MutArray a -> m Ordering
879 case compare lenA1 lenA2 of
880 EQ -> loop (lenA1 - 1)
891 v1 <- getIndexUnsafe i a1
892 v2 <- getIndexUnsafe i a2
893 case compare v1 v2 of
898eq :: (MonadIO m, Eq a) => MutArray a -> MutArray a -> m Bool
901 then loop (lenA1 - 1)
910 | i < 0 = return True
912 v1 <- getIndexUnsafe i a1
913 v2 <- getIndexUnsafe i a2
919strip :: MonadIO m => (a -> Bool) -> MutArray a -> m (MutArray a)
920strip p arr = liftIO $ do
921 let lastIndex = length arr - 1
922 indexR <- getIndexR lastIndex -- last predicate failing index
926 indexL <- getIndexL 0 -- first predicate failing index
927 if indexL == 0 && indexR == lastIndex
930 let newLen = indexR - indexL + 1
931 in return $ getSliceUnsafe indexL newLen arr
936 | idx < 0 = return idx
938 r <- getIndexUnsafe idx arr
940 then getIndexR (idx - 1)
944 r <- getIndexUnsafe idx arr
946 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