@@ -83,6 +83,7 @@ module Streamly.Internal.Data.Array.Foreign.Mut.Type
8383 , toList
8484
8585 -- * Combining
86+ , spliceWith
8687 , splice
8788 , spliceExp
8889 , spliceTwo
@@ -155,6 +156,7 @@ plusForeignPtr (ForeignPtr addr c) (I# d) = ForeignPtr (plusAddr# addr d) c
155156
156157-- $setup
157158-- >>> import qualified Streamly.Internal.Data.Stream.StreamD as StreamD
159+-- >>> import qualified Streamly.Internal.Data.Array.Foreign.Mut.Type as Array
158160
159161-------------------------------------------------------------------------------
160162-- Array Data Type
@@ -470,8 +472,6 @@ snocWith allocSize arr x = liftIO $ do
470472--
471473-- Note that the returned array may be a mutated version of the original array.
472474--
473−-- >>> snocLinear = snocWith (+ arrayChunkSize)
474−--
475475-- Performs O(n^2) copies to grow but is thrifty on memory.
476476--
477477-- /Pre-release/
@@ -1235,29 +1235,32 @@ spliceTwo arr1 arr2 = do
12351235 touchForeignPtr (aStart arr2)
12361236 return arr { aEnd = dst `plusPtr` (len1 + len2) }
12371237
1238−-- | Splice an array into a pre-reserved mutable array. The user must ensure
1239−-- that there is enough space in the mutable array, otherwise the splicing
1240−-- fails.
1241−{-# INLINE splice #-}
1242−splice :: (MonadIO m) => Array a -> Array a -> m (Array a)
1243−splice dst@(Array _ end bound) src =
1238+-- | Really really unsafe, appends the second array into the first array. If
1239+-- the first array does not have enough space it may cause silent data
1240+-- corruption or if you are lucky a segfault.
1241+{-# INLINE spliceUnsafe #-}
1242+spliceUnsafe :: MonadIO m => Array a -> (Array a, Int) -> m (Array a)
1243+spliceUnsafe dst (src, srcLen) =
12441244 liftIO $ do
1245− let srcLen = byteLength src
1246− if end `plusPtr` srcLen > bound
1247− then error
1248− "Bug: splice: Not enough space in the target array"
1249− else unsafeWithForeignPtr (aStart dst) $ \_ ->
1250− unsafeWithForeignPtr (aStart src) $ \psrc -> do
1251− let pdst = aEnd dst
1252− memcpy (castPtr pdst) (castPtr psrc) srcLen
1253− return $ dst {aEnd = pdst `plusPtr` srcLen}
1254−
1255−-- | Splice a new array into a preallocated mutable array, doubling the space
1256−-- if there is no space in the target array.
1257−{-# INLINE spliceExp #-}
1258−spliceExp :: (MonadIO m, Storable a)
1259− => Array a -> Array a -> m (Array a)
1260−spliceExp dst@(Array start end bound) src = do
1245+ unsafeWithForeignPtr (aStart dst) $ \_ ->
1246+ unsafeWithForeignPtr (aStart src) $ \psrc -> do
1247+ let pdst = aEnd dst
1248+ assert (pdst `plusPtr` srcLen <= aBound dst) (return ())
1249+ memcpy (castPtr pdst) (castPtr psrc) srcLen
1250+ return $ dst {aEnd = pdst `plusPtr` srcLen}
1251+
1252+-- | @spliceWith sizer dst src@ mutates @dst@ to append @src@. If there is no
1253+-- reserved space available in the @dst@ it is reallocated to a size determined
1254+-- by @sizer dstSize srcSize@ function, where @dstSize@ is the number of bytes
1255+-- in the first array and @srcSize@ is the number of bytes in the second array.
1256+--
1257+-- Note that the returned array may be a mutated version of first array.
1258+--
1259+-- /Pre-release/
1260+{-# INLINE spliceWith #-}
1261+spliceWith :: forall m a. (MonadIO m, Storable a) =>
1262+ (Int -> Int -> Int) -> Array a -> Array a -> m (Array a)
1263+spliceWith allocSize dst@(Array start end bound) src = do
12611264 assert (end <= bound) (return ())
12621265 let srcLen = aEnd src `minusPtr` unsafeForeignPtrToPtr (aStart src)
12631266
@@ -1266,7 +1269,7 @@ spliceExp dst@(Array start end bound) src = do
12661269 then do
12671270 let oldStart = unsafeForeignPtrToPtr start
12681271 oldSize = end `minusPtr` oldStart
1269− newSize = max (oldSize * 2) (oldSize + srcLen)
1272+ newSize = allocSize oldSize srcLen
12701273 when (newSize < oldSize + srcLen)
12711274 $ error
12721275 $ "splice: newSize is less than the total size "
@@ -1274,7 +1277,34 @@ spliceExp dst@(Array start end bound) src = do
12741277 ++ "newSize function passed."
12751278 liftIO $ realloc newSize dst
12761279 else return dst
1277− splice dst1 src
1280+ spliceUnsafe dst1 (src, srcLen)
1281+
1282+-- | The first array is mutated to append the second array. If there is no
1283+-- reserved space available in the first array a new allocation of exact
1284+-- required size is done.
1285+--
1286+-- Note that the returned array may be a mutated version of first array.
1287+--
1288+-- >>> splice = Array.spliceWith (+)
1289+--
1290+-- /Pre-release/
1291+{-# INLINE splice #-}
1292+splice :: (MonadIO m, Storable a) => Array a -> Array a -> m (Array a)
1293+splice = spliceWith (+)
1294+
1295+-- | Like 'append' but the growth of the array is exponential. Whenever a new
1296+-- allocation is required the previous array size is at least doubled.
1297+--
1298+-- This is useful to reduce allocations when folding many arrays together.
1299+--
1300+-- Note that the returned array may be a mutated version of first array.
1301+--
1302+-- >>> spliceExp = Array.spliceWith (\l1 l2 -> max (l1 * 2) (l1 + l2))
1303+--
1304+-- /Pre-release/
1305+{-# INLINE spliceExp #-}
1306+spliceExp :: (MonadIO m, Storable a) => Array a -> Array a -> m (Array a)
1307+spliceExp = spliceWith (\l1 l2 -> max (l1 * 2) (l1 + l2))
12781308
12791309-------------------------------------------------------------------------------
12801310-- Splitting