@@ -25,9 +25,8 @@ module Streamly.Internal.Data.Array.Generic.Mut.Type
2525 -- *** From streams
2626 , writeNUnsafe
2727 , writeN
28−
29− -- , writeWith
30− -- , write
28+ , writeWith
29+ , write
3130
3231 -- , writeRevN
3332 -- , writeRev
@@ -160,6 +159,7 @@ where
160159#include "inline.hs"
161160#include "assert.hs"
162161
162+import Control.Monad (when)
163163import Control.Monad.IO.Class (MonadIO(..))
164164import GHC.Base
165165 ( MutableArray#
@@ -329,8 +329,8 @@ modifyIndex i arr@MutArray {..} f = do
329329-------------------------------------------------------------------------------
330330
331331-- | Reallocates the array according to the new size. This is a safe function
332−-- that always creates a new array and copies the old array into the new one. If
333−-- the reallocated size is less than the original array it results in a
332+-- that always creates a new array and copies the old array into the new one.
333+-- If the reallocated size is less than the original array it results in a
334334-- truncated version of the original array.
335335--
336336realloc :: MonadIO m => Int -> MutArray a -> m (MutArray a)
@@ -409,7 +409,7 @@ snocWith sizer arr@MutArray {..} x = do
409409 else snocWithRealloc sizer arr x
410410
411411-- XXX round it to next power of 2.
412−--
412+
413413-- | The array is mutated to append an additional element to it. If there is no
414414-- reserved space available in the array then it is reallocated to double the
415415-- original size.
@@ -487,7 +487,6 @@ getSliceUnsafe index len arr@MutArray {..} =
487487 assert (index >= 0 && len >= 0 && index + len <= arrLen)
488488 $ arr {arrStart = arrStart + index, arrLen = len}
489489
490−
491490-- | /O(1)/ Slice an array in constant time. Throws an error if the slice
492491-- extends out of the array bounds.
493492--
@@ -509,6 +508,9 @@ getSlice index len arr@MutArray{..} =
509508-- to Lists and streams
510509-------------------------------------------------------------------------------
511510
511+-- XXX Maybe faster to create a list explicitly instead of mapM, if list fusion
512+-- does not work well.
513+
512514-- | Convert an 'Array' into a list.
513515--
514516-- /Pre-release/
@@ -538,11 +540,17 @@ toStreamK arr@MutArray{..} = K.unfoldrM step 0
538540 x <- getIndexUnsafe i arr
539541 return $ Just (x, i + 1)
540542
541−
542543-------------------------------------------------------------------------------
543544-- Folds
544545-------------------------------------------------------------------------------
545546
547+-- XXX deduplicate this across unboxed array and this module?
548+
549+-- | The default chunk size by which the array creation routines increase the
550+-- size of the array when the array is grown linearly.
551+arrayChunkSize :: Int
552+arrayChunkSize = 1024
553+
546554-- | Like 'writeN' but does not check the array bounds when writing. The fold
547555-- driver must not call the step function more than 'n' times otherwise it will
548556-- corrupt the memory and crash. This function exists mainly because any
@@ -570,6 +578,50 @@ writeNUnsafe n = Fold step initial return
570578writeN :: MonadIO m => Int -> Fold m a (MutArray a)
571579writeN n = FL.take n $ writeNUnsafe n
572580
581+-- >>> f n = MutArray.writeAppendWith (* 2) (MutArray.newPinned n)
582+-- >>> writeWith n = Fold.rmapM MutArray.rightSize (f n)
583+-- >>> writeWith n = Fold.rmapM MutArray.fromArrayStreamK (MutArray.writeChunks n)
584+
585+-- | @writeWith minCount@ folds the whole input to a single array. The array
586+-- starts at a size big enough to hold minCount elements, the size is doubled
587+-- every time the array needs to be grown.
588+--
589+-- /Caution! Do not use this on infinite streams./
590+--
591+-- /Pre-release/
592+{-# INLINE_NORMAL writeWith #-}
593+writeWith :: MonadIO m => Int -> Fold m a (MutArray a)
594+-- writeWith n = FL.rmapM rightSize $ writeAppendWith (* 2) (newPinned n)
595+writeWith elemCount = FL.rmapM extract $ FL.foldlM' step initial
596+
597+ where
598+
599+ initial = do
600+ when (elemCount < 0) $ error "writeWith: elemCount is negative"
601+ liftIO $ new elemCount
602+
603+ step arr@(MutArray _ start end bound) x
604+ | end == bound = do
605+ let oldSize = end - start
606+ newSize = max (oldSize * 2) 1
607+ arr1 <- liftIO $ realloc newSize arr
608+ snocUnsafe arr1 x
609+ step arr x = snocUnsafe arr x
610+
611+ -- extract = liftIO . rightSize
612+ extract = return
613+
614+-- | Fold the whole input to a single array.
615+--
616+-- Same as 'writeWith' using an initial array size of 'arrayChunkSize' bytes
617+-- rounded up to the element size.
618+--
619+-- /Caution! Do not use this on infinite streams./
620+--
621+{-# INLINE write #-}
622+write :: MonadIO m => Fold m a (MutArray a)
623+write = writeWith arrayChunkSize
624+
573625-------------------------------------------------------------------------------
574626-- Unfolds
575627-------------------------------------------------------------------------------