@@ -0,0 +1,569 @@
1+{-# LANGUAGE UnboxedTuples #-}
4+-- Module : Streamly.Internal.Data.Array.Mut.Type
5+-- Copyright : (c) 2020 Composewell Technologies
6+-- License : BSD3-3-Clause
7+-- Maintainer : streamly@composewell.com
8+-- Stability : experimental
11+module Streamly.Internal.Data.Array.Mut.Type
17+ -- * Constructing and Writing
21+ -- *** Uninitialized Arrays
35+ -- ** From containers
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.
62+ -- ** Appending elements
69+ -- ** Appending streams
76+ -- These are not the same as slicing the array at the beginning, they may
77+ -- reduce the length as well as the capacity of the array.
82+ -- * Eliminating and Reading
105+ -- * In-place Mutation Algorithms
118+ -- * Arrays of arrays
119+ -- We can add dimensionality parameter to the array type to get
120+ -- multidimensional arrays. Multidimensional arrays would just be a
121+ -- convenience wrapper on top of single dimensional arrays.
123+ -- | Operations dealing with multiple arrays, streams of arrays or
124+ -- multidimensional array representations.
126+ -- ** Construct from streams
128+ -- , arrayStreamKFromStreamD
131+ -- ** Eliminate to streams
133+ -- , flattenArraysRev
134+ -- , fromArrayStreamK
136+ -- ** Construct from arrays
137+ -- get chunks without copying
140+ -- , getSlicesFromLenN
141+ -- , splitAt -- XXX should be able to express using getSlice
144+ -- ** Appending arrays
151+ -- , appendSliceFrom
157+import Control.Exception (assert)
158+import Control.Monad.IO.Class (MonadIO(..))
167+import GHC.IO (IO(..))
168+import GHC.Int (Int(..))
169+import Streamly.Internal.Data.Fold.Type (Fold(..))
170+import Streamly.Internal.Data.Producer.Type (Producer (..))
171+import Streamly.Internal.Data.Unfold.Type (Unfold(..))
173+import qualified Streamly.Internal.Data.Fold.Type as FL
174+import qualified Streamly.Internal.Data.Producer as Producer
175+import qualified Streamly.Internal.Data.Stream.StreamD as D
176+import qualified Streamly.Internal.Data.Stream.StreamK as K
178+import Prelude hiding (read)
182+-- >>> import qualified Streamly.Internal.Data.Array.Mut.Type as Array
183+-- >>> import qualified Streamly.Internal.Data.Stream.IsStream as Stream
184+-- >>> import qualified Streamly.Internal.Data.Stream.StreamD as StreamD
185+-- >>> import qualified Streamly.Internal.Data.Fold as Fold
186+-- >>> import qualified Streamly.Internal.Data.Fold.Type as Fold
188+-------------------------------------------------------------------------------
190+-------------------------------------------------------------------------------
194+ { arrContents# :: MutableArray# RealWorld a
195+ -- ^ The internal contents of the array representing the entire array.
197+ , arrStart :: {-# UNPACK #-}!Int
198+ -- ^ The starting index of this slice.
200+ , arrLen :: {-# UNPACK #-}!Int
201+ -- ^ The length of this slice.
203+ , arrTrueLen :: {-# UNPACK #-}!Int
204+ -- ^ This is the true length of the array. Coincidentally, this also
205+ -- represents the first index beyond the maximum acceptable index of
206+ -- the array. This is specific to the array contents itself and not
207+ -- dependent on the slice. This value should not change and is shared
208+ -- across all the slices.
211+{-# INLINE bottomElement #-}
217+ , "This is the bottom element of the array."
218+ , "This is a place holder and should never be reached!"
223+ funcName = "Streamly.Internal.Data.Array.Mut.Type.bottomElement:"
225+-- XXX Would be nice if GHC can provide something like newUninitializedArray# so
226+-- that we do not have to write undefined or error in the whole array.
227+-- | @newArray count@ allocates an empty array that can hold 'count' items.
230+{-# INLINE newArray #-}
231+newArray :: forall m a. MonadIO m => Int -> m (Array a)
236+ case newArray# n# bottomElement s# of
238+ let ma = Array arr# 0 0 n
241+-------------------------------------------------------------------------------
243+-------------------------------------------------------------------------------
245+-- | Write the given element to the given index of the array. Does not check if
246+-- the index is out of bounds of the array.
249+{-# INLINE putIndexUnsafe #-}
250+putIndexUnsafe :: forall m a. MonadIO m => Array a -> Int -> a -> m ()
251+putIndexUnsafe Array {..} i x =
255+ case i + arrStart of
257+ let s1# = writeArray# arrContents# n# x s#
260+invalidIndex :: String -> Int -> a
261+invalidIndex label i =
262+ error $ label ++ ": invalid array index " ++ show i
264+-- | /O(1)/ Write the given element at the given index in the array.
265+-- Performs in-place mutation of the array.
267+-- >>> putIndex arr ix val = Array.modifyIndex arr ix (const (val, ()))
270+{-# INLINE putIndex #-}
271+putIndex :: MonadIO m => Array a -> Int -> a -> m ()
272+putIndex arr@(Array {..}) i x =
273+ if i >= 0 && i < arrLen
274+ then putIndexUnsafe arr i x
275+ else invalidIndex "putIndex" i
277+-- | Modify a given index of an array using a modifier function without checking
280+-- Unsafe because it does not check the bounds of the array.
283+modifyIndexUnsafe :: MonadIO m => Array a -> Int -> (a -> (a, b)) -> m b
284+modifyIndexUnsafe Array {..} i f = do
288+ case i + arrStart of
290+ case readArray# arrContents# n# s# of
293+ s2# = writeArray# arrContents# n# a1 s1#
296+-- | Modify a given index of an array using a modifier function.
299+modifyIndex :: MonadIO m => Array a -> Int -> (a -> (a, b)) -> m b
300+modifyIndex arr@(Array {..}) i f = do
301+ if i >= 0 && i < arrLen
302+ then modifyIndexUnsafe arr i f
303+ else invalidIndex "modifyIndex" i
305+-------------------------------------------------------------------------------
307+-------------------------------------------------------------------------------
309+-- | Reallocates the array according to the new size. This is a safe function
310+-- that always creates a new array and copies the old array into the new one. If
311+-- the reallocated size is less than the original array it results in a
312+-- truncated version of the original array.
314+realloc :: MonadIO m => Int -> Array a -> m (Array a)
317+ let !newLen@(I# newLen#) = min n (arrLen arr)
318+ !(I# arrS#) = arrStart arr
319+ !(I# arr1S#) = arrStart arr1
320+ arrC# = arrContents# arr
321+ arr1C# = arrContents# arr1
325+ let s1# = copyMutableArray# arrC# arrS# arr1C# arr1S# newLen# s#
326+ in (# s1#, arr1 {arrLen = newLen, arrTrueLen = n} #)
329+ MonadIO m => String -> (Int -> Int) -> Int -> Array a -> m (Array a)
330+reallocWith label sizer reqSize arr = do
331+ let oldSize = arrLen arr
332+ newSize = sizer oldSize
333+ safeSize = max newSize (oldSize + reqSize)
334+ assert (newSize >= oldSize + reqSize || error badSize) (return ())
335+ realloc safeSize arr
341+ , ": new array size is less than required size "
343+ , ". Please check the sizing function passed."
346+-------------------------------------------------------------------------------
348+-------------------------------------------------------------------------------
350+-- XXX Not sure of the behavior of writeArray# if we specify an index which is
351+-- out of bounds. This comment should be rewritten based on that.
352+-- | Really really unsafe, appends the element into the first array, may
353+-- cause silent data corruption or if you are lucky a segfault if the index
357+{-# INLINE snocUnsafe #-}
358+snocUnsafe :: MonadIO m => Array a -> a -> m (Array a)
359+snocUnsafe arr@Array {..} a = do
360+ assert (arrStart + arrLen < arrTrueLen) (return ())
361+ putIndexUnsafe arr arrLen a
362+ return $ arr {arrLen = arrLen + 1}
364+-- NOINLINE to move it out of the way and not pollute the instruction cache.
365+{-# NOINLINE snocWithRealloc #-}
366+snocWithRealloc :: MonadIO m => (Int -> Int) -> Array a -> a -> m (Array a)
367+snocWithRealloc sizer arr x = do
368+ arr1 <- reallocWith "snocWithRealloc" sizer 1 arr
371+-- | @snocWith sizer arr elem@ mutates @arr@ to append @elem@. The length of
372+-- the array increases by 1.
374+-- If there is no reserved space available in @arr@ it is reallocated to a size
375+-- in bytes determined by the @sizer oldSize@ function, where @oldSize@ is the
376+-- original size of the array.
378+-- Note that the returned array may be a mutated version of the original array.
381+{-# INLINE snocWith #-}
382+snocWith :: MonadIO m => (Int -> Int) -> Array a -> a -> m (Array a)
383+snocWith sizer arr@(Array {..}) x = do
384+ if arrStart + arrLen < arrTrueLen
385+ then snocUnsafe arr x
386+ else snocWithRealloc sizer arr x
388+-- XXX round it to next power of 2.
390+-- | The array is mutated to append an additional element to it. If there is no
391+-- reserved space available in the array then it is reallocated to double the
394+-- This is useful to reduce allocations when appending unknown number of
397+-- Note that the returned array may be a mutated version of the original array.
399+-- >>> snoc = Array.snocWith (* 2)
401+-- Performs O(n * log n) copies to grow, but is liberal with memory allocation.
405+snoc :: MonadIO m => Array a -> a -> m (Array a)
406+snoc = snocWith (* 2)
408+-------------------------------------------------------------------------------
410+-------------------------------------------------------------------------------
412+-- | Return the element at the specified index without checking the bounds.
414+-- Unsafe because it does not check the bounds of the array.
415+{-# INLINE_NORMAL getIndexUnsafe #-}
416+getIndexUnsafe :: MonadIO m => Array a -> Int -> m a
417+getIndexUnsafe Array {..} n =
421+ let !(I# i#) = arrStart + n
422+ in readArray# arrContents# i# s#
424+-- | /O(1)/ Lookup the element at the given index. Index starts from 0.
426+{-# INLINE getIndex #-}
427+getIndex :: MonadIO m => Array a -> Int -> m a
428+getIndex arr@(Array {..}) i =
429+ if i >= 0 && i < arrLen
430+ then getIndexUnsafe arr i
431+ else invalidIndex "getIndex" i
433+-------------------------------------------------------------------------------
435+-------------------------------------------------------------------------------
437+-- XXX We can also get immutable slices.
439+-- | /O(1)/ Slice an array in constant time.
441+-- Unsafe: The bounds of the slice are not checked.
446+{-# INLINE getSliceUnsafe #-}
448+ :: Int -- ^ from index
449+ -> Int -- ^ length of the slice
452+getSliceUnsafe index len arr@(Array {..}) =
453+ assert (index >= 0 && len >= 0 && index + len <= arrLen)
454+ $ arr {arrStart = arrStart + index, arrLen = len}
457+-- | /O(1)/ Slice an array in constant time. Throws an error if the slice
458+-- extends out of the array bounds.
461+{-# INLINE getSlice #-}
463+ :: Int -- ^ from index
464+ -> Int -- ^ length of the slice
467+getSlice index len arr@(Array{..}) =
468+ if index >= 0 && len >= 0 && index + len <= arrLen
469+ then arr {arrStart = arrStart + index, arrLen = len}
471+ $ "getSlice: invalid slice, index "
472+ ++ show index ++ " length " ++ show len
474+-------------------------------------------------------------------------------
475+-- to Lists and streams
476+-------------------------------------------------------------------------------
478+-- | Convert an 'Array' into a list.
481+{-# INLINE toList #-}
482+toList :: MonadIO m => Array a -> m [a]
483+toList arr@(Array {..}) = mapM (getIndexUnsafe arr) [0 .. (arrLen - 1)]
485+-- | Use the 'read' unfold instead.
487+-- @toStreamD = D.unfold read@
489+-- We can try this if the unfold has any performance issues.
490+{-# INLINE_NORMAL toStreamD #-}
491+toStreamD :: MonadIO m => Array a -> D.Stream m a
492+toStreamD arr@(Array {..}) =
493+ D.mapM (getIndexUnsafe arr) $ D.enumerateFromToIntegral 0 (arrLen - 1)
495+{-# INLINE toStreamK #-}
496+toStreamK :: MonadIO m => Array a -> K.Stream m a
497+toStreamK arr@(Array {..}) = K.unfoldrM step 0
502+ | i == arrLen = return Nothing
504+ x <- getIndexUnsafe arr i
505+ return $ Just (x, i + 1)
508+-------------------------------------------------------------------------------
510+-------------------------------------------------------------------------------
512+-- | Like 'writeN' but does not check the array bounds when writing. The fold
513+-- driver must not call the step function more than 'n' times otherwise it will
514+-- corrupt the memory and crash. This function exists mainly because any
515+-- conditional in the step function blocks fusion causing 10x performance
519+{-# INLINE_NORMAL writeNUnsafe #-}
520+writeNUnsafe :: MonadIO m => Int -> Fold m a (Array a)
521+writeNUnsafe n = Fold step initial return
525+ initial = FL.Partial <$> newArray (max n 0)
527+ step arr x = FL.Partial <$> snocUnsafe arr x
529+-- | @writeN n@ folds a maximum of @n@ elements from the input stream to an
532+-- >>> writeN n = Fold.take n (Array.writeNUnsafe n)
535+{-# INLINE_NORMAL writeN #-}
536+writeN :: MonadIO m => Int -> Fold m a (Array a)
537+writeN n = FL.take n $ writeNUnsafe n
539+-------------------------------------------------------------------------------
541+-------------------------------------------------------------------------------
543+-- | Resumable unfold of an array.
545+{-# INLINE_NORMAL producer #-}
546+producer :: MonadIO m => Producer m (Array a) a
547+producer = Producer step inject extract
551+ {-# INLINE inject #-}
552+ inject arr = return (arr, 0)
554+ {-# INLINE extract #-}
556+ return $ arr {arrStart = arrStart arr + i, arrLen = arrLen arr - i}
558+ {-# INLINE_LATE step #-}
560+ | assert (arrLen arr >= 0) (i == arrLen arr) = return D.Stop
562+ x <- getIndexUnsafe arr i
563+ return $ D.Yield x (arr, i + 1)
565+-- | Unfold an array into a stream.
567+{-# INLINE_NORMAL read #-}
568+read :: MonadIO m => Unfold m (Array a) a
569+read = Producer.simplify producer