21fff61 · streamly · git.iko.soy

git.iko.soy/streamly

on master help sign in

Refactor writeWith and spliceWith

21fff61Harendra Kumar5 years ago

src/Streamly/Internal/Data/Array/Foreign.hs+1−2

compare with master

unifiedsplit

@@ -467,8 +467,7 @@ runPipe f arr = P.runPipe (toArrayMinChunk (length arr)) $ f (A.read arr)

467streamTransform :: forall m a b. (MonadIO m, Storable a, Storable b)

468 => (SerialT m a -> SerialT m b) -> Array a -> m (Array b)

469streamTransform f arr =

470− P.fold (A.writeWith (alignment (undefined :: a)) (length arr))

471− $ getSerialT $ f (A.toStream arr)

470+ P.fold (A.writeWith (length arr)) $ getSerialT $ f (A.toStream arr)

472

471

473-------------------------------------------------------------------------------

472-------------------------------------------------------------------------------

474-- Casts

473-- Casts

src/Streamly/Internal/Data/Array/Foreign/Mut/Type.hs+23−13

@@ -106,6 +106,7 @@ module Streamly.Internal.Data.Array.Foreign.Mut.Type

106 -- * Stream of arrays

107 , arraysOf

108 , writeChunks

109+ , fromArrayStreamK

109

110

110 -- * Utilities

111 -- * Utilities

111 , bytesToElemCount

112 , bytesToElemCount

@@ -1244,12 +1245,16 @@ writeChunks n = FL.many (writeN n) FL.toStreamK

1244--

1245--

1245-- /Caution! Do not use this on infinite streams./

1246-- /Caution! Do not use this on infinite streams./

1246--

1247--

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+--

1247-- /Pre-release/

1252-- /Pre-release/

1248{-# INLINE_NORMAL writeWith #-}

1253{-# INLINE_NORMAL writeWith #-}

1249writeWith :: forall m a. (MonadIO m, Storable a)

1254writeWith :: forall m a. (MonadIO m, Storable a)

1250− => Int -> Int -> Fold m a (Array a)

1255+ => Int -> Fold m a (Array a)

1251−-- writeWith n = FL.rmapM spliceArrays $ toArraysOf n

1256+-- writeWith n = FL.rmapM rightSize $ appendWith (* 2) (newArray n)

1252−writeWith alignSize elemCount =

1257+writeWith elemCount =

1253 FL.rmapM extract $ FL.foldlM' step initial

1258 FL.rmapM extract $ FL.foldlM' step initial

1254

1259

1255 where

1260 where

@@ -1260,7 +1265,7 @@ writeWith alignSize elemCount =

1260

1265

1261 initial = do

1266 initial = do

1262 when (elemCount < 0) $ error "writeWith: elemCount is negative"

1267 when (elemCount < 0) $ error "writeWith: elemCount is negative"

1263− liftIO $ newArrayAligned alignSize elemCount

1268+ liftIO $ newArrayAligned (alignment (undefined :: a)) elemCount

1264 step arr@(Array start end bound) x

1269 step arr@(Array start end bound) x

1265 | end `plusPtr` sizeOf (undefined :: a) > bound = do

1270 | end `plusPtr` sizeOf (undefined :: a) > bound = do

1266 let p = unsafeForeignPtrToPtr start

1271 let p = unsafeForeignPtrToPtr start

@@ -1269,7 +1274,10 @@ writeWith alignSize elemCount =

1269 arr1 <-

1274 arr1 <-

1270 liftIO

1275 liftIO

1271 $ reallocAligned

1276 $ reallocAligned

1272− (sizeOf (undefined :: a)) alignSize newSize arr

1277+ (sizeOf (undefined :: a))

1278+ (alignment (undefined :: a))

1279+ newSize

1280+ arr

1273 insertElem arr1 x

1281 insertElem arr1 x

1274 step arr x = insertElem arr x

1282 step arr x = insertElem arr x

1275 extract = liftIO . rightSize

1283 extract = liftIO . rightSize

@@ -1284,9 +1292,7 @@ writeWith alignSize elemCount =

1284-- @since 0.7.0

1292-- @since 0.7.0

1285{-# INLINE write #-}

1293{-# INLINE write #-}

1286write :: forall m a. (MonadIO m, Storable a) => Fold m a (Array a)

1294write :: forall m a. (MonadIO m, Storable a) => Fold m a (Array a)

1287−write = writeWith (alignment (undefined :: a))

1288− (bytesToElemCount (undefined :: a)

1295+write = writeWith (bytesToElemCount (undefined :: a) arrayChunkSize)

1289− (mkChunkSize 1024))

1290

1296

1291-------------------------------------------------------------------------------

1297-------------------------------------------------------------------------------

1292-- construct from streams, known size

1298-- construct from streams, known size

@@ -1402,9 +1408,9 @@ spliceUnsafe dst (src, srcLen) =

1402 return $ dst {aEnd = pdst `plusPtr` srcLen}

1408 return $ dst {aEnd = pdst `plusPtr` srcLen}

1403

1409

1404-- | @spliceWith sizer dst src@ mutates @dst@ to append @src@. If there is no

1410-- | @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

1411+-- reserved space available in @dst@ it is reallocated to a size determined by

1406−-- by @sizer dstSize srcSize@ function, where @dstSize@ is the number of bytes

1412+-- the @sizer dstBytesn srcBytes@ function, where @dstBytes@ is the size of the

1407−-- in the first array and @srcSize@ is the number of bytes in the second array.

1413+-- first array and @srcBytes@ is the size of the second array, in bytes.

1408--

1414--

1409-- Note that the returned array may be a mutated version of first array.

1415-- Note that the returned array may be a mutated version of first array.

1410--

1416--

@@ -1412,7 +1418,11 @@ spliceUnsafe dst (src, srcLen) =

1412{-# INLINE spliceWith #-}

1418{-# INLINE spliceWith #-}

1413spliceWith :: forall m a. (MonadIO m, Storable a) =>

1419spliceWith :: forall m a. (MonadIO m, Storable a) =>

1414 (Int -> Int -> Int) -> Array a -> Array a -> m (Array a)

1420 (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+-}

1416 assert (end <= bound) (return ())

1426 assert (end <= bound) (return ())

1417 let srcLen = aEnd src `minusPtr` unsafeForeignPtrToPtr (aStart src)

1427 let srcLen = aEnd src `minusPtr` unsafeForeignPtrToPtr (aStart src)

1418

1428

@@ -1421,7 +1431,7 @@ spliceWith allocSize dst@(Array start end bound) src = do

1421 then do

1431 then do

1422 let oldStart = unsafeForeignPtrToPtr start

1432 let oldStart = unsafeForeignPtrToPtr start

1423 oldSize = end `minusPtr` oldStart

1433 oldSize = end `minusPtr` oldStart

1424− newSize = allocSize oldSize srcLen

1434+ newSize = sizer oldSize srcLen

1425 when (newSize < oldSize + srcLen)

1435 when (newSize < oldSize + srcLen)

1426 $ error

1436 $ error

1427 $ "splice: newSize is less than the total size "

1437 $ "splice: newSize is less than the total size "

src/Streamly/Internal/Data/Array/Foreign/Type.hs+2−3

@@ -551,10 +551,9 @@ writeNUnsafe n = unsafeFreeze <$> MA.writeNUnsafe n

551

552{-# INLINE_NORMAL writeWith #-}

553writeWith :: forall m a. (MonadIO m, Storable a)

554− => Int -> Int -> Fold m a (Array a)

554+ => Int -> Fold m a (Array a)

555-- writeWith n = FL.rmapM spliceArrays $ toArraysOf n

556−writeWith alignSize elemCount =

557− unsafeFreeze <$> MA.writeWith alignSize elemCount

556+writeWith elemCount = unsafeFreeze <$> MA.writeWith elemCount

558

557

559-- | Fold the whole input to a single array.

558-- | Fold the whole input to a single array.

560--

559--