@@ -106,6 +106,7 @@ module Streamly.Internal.Data.Array.Foreign.Mut.Type
106106 -- * Stream of arrays
107107 , arraysOf
108108 , writeChunks
109+ , fromArrayStreamK
109110
110111 -- * Utilities
111112 , bytesToElemCount
@@ -1244,12 +1245,16 @@ writeChunks n = FL.many (writeN n) FL.toStreamK
12441245--
12451246-- /Caution! Do not use this on infinite streams./
12461247--
1248+-- >>> f n = Array.appendWith (* 2) (Array.newArray n)
1249+-- >>> writeWith n = Fold.rmapM Array.rightSize (f n)
1250+-- >>> writeWith n = Fold.rmapM Array.fromArrayStreamK (Array.writeChunks n)
1251+--
12471252-- /Pre-release/
12481253{-# INLINE_NORMAL writeWith #-}
12491254writeWith :: forall m a. (MonadIO m, Storable a)
1250− => Int -> Int -> Fold m a (Array a)
1251−-- writeWith n = FL.rmapM spliceArrays $ toArraysOf n
1252−writeWith alignSize elemCount =
1255+ => Int -> Fold m a (Array a)
1256+-- writeWith n = FL.rmapM rightSize $ appendWith (* 2) (newArray n)
1257+writeWith elemCount =
12531258 FL.rmapM extract $ FL.foldlM' step initial
12541259
12551260 where
@@ -1260,7 +1265,7 @@ writeWith alignSize elemCount =
12601265
12611266 initial = do
12621267 when (elemCount < 0) $ error "writeWith: elemCount is negative"
1263− liftIO $ newArrayAligned alignSize elemCount
1268+ liftIO $ newArrayAligned (alignment (undefined :: a)) elemCount
12641269 step arr@(Array start end bound) x
12651270 | end `plusPtr` sizeOf (undefined :: a) > bound = do
12661271 let p = unsafeForeignPtrToPtr start
@@ -1269,7 +1274,10 @@ writeWith alignSize elemCount =
12691274 arr1 <-
12701275 liftIO
12711276 $ reallocAligned
1272− (sizeOf (undefined :: a)) alignSize newSize arr
1277+ (sizeOf (undefined :: a))
1278+ (alignment (undefined :: a))
1279+ newSize
1280+ arr
12731281 insertElem arr1 x
12741282 step arr x = insertElem arr x
12751283 extract = liftIO . rightSize
@@ -1284,9 +1292,7 @@ writeWith alignSize elemCount =
12841292-- @since 0.7.0
12851293{-# INLINE write #-}
12861294write :: forall m a. (MonadIO m, Storable a) => Fold m a (Array a)
1287−write = writeWith (alignment (undefined :: a))
1288− (bytesToElemCount (undefined :: a)
1289− (mkChunkSize 1024))
1295+write = writeWith (bytesToElemCount (undefined :: a) arrayChunkSize)
12901296
12911297-------------------------------------------------------------------------------
12921298-- construct from streams, known size
@@ -1402,9 +1408,9 @@ spliceUnsafe dst (src, srcLen) =
14021408 return $ dst {aEnd = pdst `plusPtr` srcLen}
14031409
14041410-- | @spliceWith sizer dst src@ mutates @dst@ to append @src@. If there is no
1405−-- reserved space available in the @dst@ it is reallocated to a size determined
1406−-- by @sizer dstSize srcSize@ function, where @dstSize@ is the number of bytes
1407−-- in the first array and @srcSize@ is the number of bytes in the second array.
1411+-- reserved space available in @dst@ it is reallocated to a size determined by
1412+-- the @sizer dstBytesn srcBytes@ function, where @dstBytes@ is the size of the
1413+-- first array and @srcBytes@ is the size of the second array, in bytes.
14081414--
14091415-- Note that the returned array may be a mutated version of first array.
14101416--
@@ -1412,7 +1418,11 @@ spliceUnsafe dst (src, srcLen) =
14121418{-# INLINE spliceWith #-}
14131419spliceWith :: forall m a. (MonadIO m, Storable a) =>
14141420 (Int -> Int -> Int) -> Array a -> Array a -> m (Array a)
1415−spliceWith allocSize dst@(Array start end bound) src = do
1421+spliceWith sizer dst@(Array start end bound) src = do
1422+{-
1423+ let f = appendWith (`sizer` byteLength src) (return dst)
1424+ in D.fold f (toStreamD src)
1425+-}
14161426 assert (end <= bound) (return ())
14171427 let srcLen = aEnd src `minusPtr` unsafeForeignPtrToPtr (aStart src)
14181428
@@ -1421,7 +1431,7 @@ spliceWith allocSize dst@(Array start end bound) src = do
14211431 then do
14221432 let oldStart = unsafeForeignPtrToPtr start
14231433 oldSize = end `minusPtr` oldStart
1424− newSize = allocSize oldSize srcLen
1434+ newSize = sizer oldSize srcLen
14251435 when (newSize < oldSize + srcLen)
14261436 $ error
14271437 $ "splice: newSize is less than the total size "