@@ -0,0 +1,569 @@
1+{-# LANGUAGE UnboxedTuples #-}
2+
3+-- |
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
9+-- Portability : GHC
10+--
11+module Streamly.Internal.Data.Array.Mut.Type
12+(
13+ -- * Type
14+ -- $arrayNotes
15+ Array (..)
16+
17+ -- * Constructing and Writing
18+ -- ** Construction
19+ -- , nil
20+
21+ -- *** Uninitialized Arrays
22+ , newArray
23+ -- , newArrayWith
24+
25+ -- *** From streams
26+ , writeNUnsafe
27+ , writeN
28+
29+ -- , writeWith
30+ -- , write
31+
32+ -- , writeRevN
33+ -- , writeRev
34+
35+ -- ** From containers
36+ -- , fromListN
37+ -- , fromList
38+ -- , fromStreamDN
39+ -- , fromStreamD
40+
41+ -- * Random writes
42+ , putIndex
43+ , putIndexUnsafe
44+ -- , putIndices
45+ -- , putFromThenTo
46+ -- , putFrom -- start writing at the given position
47+ -- , putUpto -- write from beginning up to the given position
48+ -- , putFromTo
49+ -- , putFromRev
50+ -- , putUptoRev
51+ , modifyIndexUnsafe
52+ , modifyIndex
53+ -- , modifyIndices
54+ -- , modify
55+ -- , swapIndices
56+
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.
61+
62+ -- ** Appending elements
63+ , snocWith
64+ , snoc
65+ -- , snocLinear
66+ -- , snocMay
67+ , snocUnsafe
68+
69+ -- ** Appending streams
70+ -- , appendNUnsafe
71+ -- , appendN
72+ -- , appendWith
73+ -- , append
74+
75+ -- ** Truncation
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.
78+ -- , truncateWith
79+ -- , truncate
80+ -- , truncateExp
81+
82+ -- * Eliminating and Reading
83+
84+ -- ** To streams
85+ , read
86+ -- , readRev
87+
88+ -- ** To containers
89+ , toStreamD
90+ -- , toStreamDRev
91+ , toStreamK
92+ -- , toStreamKRev
93+ , toList
94+
95+ -- experimental
96+ , producer
97+
98+ -- ** Random reads
99+ , getIndex
100+ , getIndexUnsafe
101+ -- , getIndices
102+ -- , getFromThenTo
103+ -- , getIndexRev
104+
105+ -- * In-place Mutation Algorithms
106+ -- , reverse
107+ -- , permute
108+ -- , partitionBy
109+ -- , shuffleBy
110+ -- , divideBy
111+ -- , mergeBy
112+
113+ -- * Folding
114+ -- , foldl'
115+ -- , foldr
116+ -- , cmp
117+
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.
122+
123+ -- | Operations dealing with multiple arrays, streams of arrays or
124+ -- multidimensional array representations.
125+
126+ -- ** Construct from streams
127+ -- , arraysOf
128+ -- , arrayStreamKFromStreamD
129+ -- , writeChunks
130+
131+ -- ** Eliminate to streams
132+ -- , flattenArrays
133+ -- , flattenArraysRev
134+ -- , fromArrayStreamK
135+
136+ -- ** Construct from arrays
137+ -- get chunks without copying
138+ , getSliceUnsafe
139+ , getSlice
140+ -- , getSlicesFromLenN
141+ -- , splitAt -- XXX should be able to express using getSlice
142+ -- , breakOn
143+
144+ -- ** Appending arrays
145+ -- , spliceCopy
146+ -- , spliceWith
147+ -- , splice
148+ -- , spliceExp
149+ -- , putSlice
150+ -- , appendSlice
151+ -- , appendSliceFrom
152+ )
153+where
154+
155+#include "inline.hs"
156+
157+import Control.Exception (assert)
158+import Control.Monad.IO.Class (MonadIO(..))
159+import GHC.Base
160+ ( MutableArray#
161+ , RealWorld
162+ , copyMutableArray#
163+ , newArray#
164+ , readArray#
165+ , writeArray#
166+ )
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(..))
172+
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
177+
178+import Prelude hiding (read)
179+
180+-- $setup
181+-- >>> :m
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
187+
188+-------------------------------------------------------------------------------
189+-- Array Data Type
190+-------------------------------------------------------------------------------
191+
192+data Array a =
193+ Array
194+ { arrContents# :: MutableArray# RealWorld a
195+ -- ^ The internal contents of the array representing the entire array.
196+
197+ , arrStart :: {-# UNPACK #-}!Int
198+ -- ^ The starting index of this slice.
199+
200+ , arrLen :: {-# UNPACK #-}!Int
201+ -- ^ The length of this slice.
202+
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.
209+ }
210+
211+{-# INLINE bottomElement #-}
212+bottomElement :: a
213+bottomElement =
214+ error
215+ $ unwords
216+ [ funcName
217+ , "This is the bottom element of the array."
218+ , "This is a place holder and should never be reached!"
219+ ]
220+
221+ where
222+
223+ funcName = "Streamly.Internal.Data.Array.Mut.Type.bottomElement:"
224+
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.
228+--
229+-- /Pre-release/
230+{-# INLINE newArray #-}
231+newArray :: forall m a. MonadIO m => Int -> m (Array a)
232+newArray n@(I# n#) =
233+ liftIO
234+ $ IO
235+ $ \s# ->
236+ case newArray# n# bottomElement s# of
237+ (# s1#, arr# #) ->
238+ let ma = Array arr# 0 0 n
239+ in (# s1#, ma #)
240+
241+-------------------------------------------------------------------------------
242+-- Random writes
243+-------------------------------------------------------------------------------
244+
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.
247+--
248+-- /Pre-release/
249+{-# INLINE putIndexUnsafe #-}
250+putIndexUnsafe :: forall m a. MonadIO m => Array a -> Int -> a -> m ()
251+putIndexUnsafe Array {..} i x =
252+ liftIO
253+ $ IO
254+ $ \s# ->
255+ case i + arrStart of
256+ I# n# ->
257+ let s1# = writeArray# arrContents# n# x s#
258+ in (# s1#, () #)
259+
260+invalidIndex :: String -> Int -> a
261+invalidIndex label i =
262+ error $ label ++ ": invalid array index " ++ show i
263+
264+-- | /O(1)/ Write the given element at the given index in the array.
265+-- Performs in-place mutation of the array.
266+--
267+-- >>> putIndex arr ix val = Array.modifyIndex arr ix (const (val, ()))
268+--
269+-- /Pre-release/
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
276+
277+-- | Modify a given index of an array using a modifier function without checking
278+-- the bounds.
279+--
280+-- Unsafe because it does not check the bounds of the array.
281+--
282+-- /Pre-release/
283+modifyIndexUnsafe :: MonadIO m => Array a -> Int -> (a -> (a, b)) -> m b
284+modifyIndexUnsafe Array {..} i f = do
285+ liftIO
286+ $ IO
287+ $ \s# ->
288+ case i + arrStart of
289+ I# n# ->
290+ case readArray# arrContents# n# s# of
291+ (# s1#, a #) ->
292+ let (a1, b) = f a
293+ s2# = writeArray# arrContents# n# a1 s1#
294+ in (# s2#, b #)
295+
296+-- | Modify a given index of an array using a modifier function.
297+--
298+-- /Pre-release/
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
304+
305+-------------------------------------------------------------------------------
306+-- Resizing
307+-------------------------------------------------------------------------------
308+
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.
313+--
314+realloc :: MonadIO m => Int -> Array a -> m (Array a)
315+realloc n arr = do
316+ arr1 <- newArray n
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
322+ liftIO
323+ $ IO
324+ $ \s# ->
325+ let s1# = copyMutableArray# arrC# arrS# arr1C# arr1S# newLen# s#
326+ in (# s1#, arr1 {arrLen = newLen, arrTrueLen = n} #)
327+
328+reallocWith ::
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
336+
337+ where
338+
339+ badSize = concat
340+ [ label
341+ , ": new array size is less than required size "
342+ , show reqSize
343+ , ". Please check the sizing function passed."
344+ ]
345+
346+-------------------------------------------------------------------------------
347+-- Snoc
348+-------------------------------------------------------------------------------
349+
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
354+-- is out of bounds.
355+--
356+-- /Internal/
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}
363+
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
369+ snocUnsafe arr1 x
370+
371+-- | @snocWith sizer arr elem@ mutates @arr@ to append @elem@. The length of
372+-- the array increases by 1.
373+--
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.
377+--
378+-- Note that the returned array may be a mutated version of the original array.
379+--
380+-- /Pre-release/
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
387+
388+-- XXX round it to next power of 2.
389+--
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
392+-- original size.
393+--
394+-- This is useful to reduce allocations when appending unknown number of
395+-- elements.
396+--
397+-- Note that the returned array may be a mutated version of the original array.
398+--
399+-- >>> snoc = Array.snocWith (* 2)
400+--
401+-- Performs O(n * log n) copies to grow, but is liberal with memory allocation.
402+--
403+-- /Pre-release/
404+{-# INLINE snoc #-}
405+snoc :: MonadIO m => Array a -> a -> m (Array a)
406+snoc = snocWith (* 2)
407+
408+-------------------------------------------------------------------------------
409+-- Random reads
410+-------------------------------------------------------------------------------
411+
412+-- | Return the element at the specified index without checking the bounds.
413+--
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 =
418+ liftIO
419+ $ IO
420+ $ \s# ->
421+ let !(I# i#) = arrStart + n
422+ in readArray# arrContents# i# s#
423+
424+-- | /O(1)/ Lookup the element at the given index. Index starts from 0.
425+--
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
432+
433+-------------------------------------------------------------------------------
434+-- Subarrays
435+-------------------------------------------------------------------------------
436+
437+-- XXX We can also get immutable slices.
438+
439+-- | /O(1)/ Slice an array in constant time.
440+--
441+-- Unsafe: The bounds of the slice are not checked.
442+--
443+-- /Unsafe/
444+--
445+-- /Pre-release/
446+{-# INLINE getSliceUnsafe #-}
447+getSliceUnsafe
448+ :: Int -- ^ from index
449+ -> Int -- ^ length of the slice
450+ -> Array a
451+ -> Array a
452+getSliceUnsafe index len arr@(Array {..}) =
453+ assert (index >= 0 && len >= 0 && index + len <= arrLen)
454+ $ arr {arrStart = arrStart + index, arrLen = len}
455+
456+
457+-- | /O(1)/ Slice an array in constant time. Throws an error if the slice
458+-- extends out of the array bounds.
459+--
460+-- /Pre-release/
461+{-# INLINE getSlice #-}
462+getSlice
463+ :: Int -- ^ from index
464+ -> Int -- ^ length of the slice
465+ -> Array a
466+ -> Array a
467+getSlice index len arr@(Array{..}) =
468+ if index >= 0 && len >= 0 && index + len <= arrLen
469+ then arr {arrStart = arrStart + index, arrLen = len}
470+ else error
471+ $ "getSlice: invalid slice, index "
472+ ++ show index ++ " length " ++ show len
473+
474+-------------------------------------------------------------------------------
475+-- to Lists and streams
476+-------------------------------------------------------------------------------
477+
478+-- | Convert an 'Array' into a list.
479+--
480+-- /Pre-release/
481+{-# INLINE toList #-}
482+toList :: MonadIO m => Array a -> m [a]
483+toList arr@(Array {..}) = mapM (getIndexUnsafe arr) [0 .. (arrLen - 1)]
484+
485+-- | Use the 'read' unfold instead.
486+--
487+-- @toStreamD = D.unfold read@
488+--
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)
494+
495+{-# INLINE toStreamK #-}
496+toStreamK :: MonadIO m => Array a -> K.Stream m a
497+toStreamK arr@(Array {..}) = K.unfoldrM step 0
498+
499+ where
500+
501+ step i
502+ | i == arrLen = return Nothing
503+ | otherwise = do
504+ x <- getIndexUnsafe arr i
505+ return $ Just (x, i + 1)
506+
507+
508+-------------------------------------------------------------------------------
509+-- Folds
510+-------------------------------------------------------------------------------
511+
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
516+-- slowdown.
517+--
518+-- /Pre-release/
519+{-# INLINE_NORMAL writeNUnsafe #-}
520+writeNUnsafe :: MonadIO m => Int -> Fold m a (Array a)
521+writeNUnsafe n = Fold step initial return
522+
523+ where
524+
525+ initial = FL.Partial <$> newArray (max n 0)
526+
527+ step arr x = FL.Partial <$> snocUnsafe arr x
528+
529+-- | @writeN n@ folds a maximum of @n@ elements from the input stream to an
530+-- 'Array'.
531+--
532+-- >>> writeN n = Fold.take n (Array.writeNUnsafe n)
533+--
534+-- /Pre-release/
535+{-# INLINE_NORMAL writeN #-}
536+writeN :: MonadIO m => Int -> Fold m a (Array a)
537+writeN n = FL.take n $ writeNUnsafe n
538+
539+-------------------------------------------------------------------------------
540+-- Unfolds
541+-------------------------------------------------------------------------------
542+
543+-- | Resumable unfold of an array.
544+--
545+{-# INLINE_NORMAL producer #-}
546+producer :: MonadIO m => Producer m (Array a) a
547+producer = Producer step inject extract
548+
549+ where
550+
551+ {-# INLINE inject #-}
552+ inject arr = return (arr, 0)
553+
554+ {-# INLINE extract #-}
555+ extract (arr, i) =
556+ return $ arr {arrStart = arrStart arr + i, arrLen = arrLen arr - i}
557+
558+ {-# INLINE_LATE step #-}
559+ step (arr, i)
560+ | assert (arrLen arr >= 0) (i == arrLen arr) = return D.Stop
561+ step (arr, i) = do
562+ x <- getIndexUnsafe arr i
563+ return $ D.Yield x (arr, i + 1)
564+
565+-- | Unfold an array into a stream.
566+--
567+{-# INLINE_NORMAL read #-}
568+read :: MonadIO m => Unfold m (Array a) a
569+read = Producer.simplify producer