7cf3784 · streamly · git.iko.soy

git.iko.soy/streamly

on master help sign in

refactor to allow using array write fold in fromStreamD

7cf3784Harendra Kumar7 years ago

src/Streamly/Memory/Array.hs+3−51

compare with master

unifiedsplit

@@ -81,7 +81,7 @@ module Streamly.Memory.Array

81 -- Monadic APIs

82 -- , newArray

83 , A.writeN -- drop new

84− , write -- full buffer

84+ , A.write -- full buffer

85 -- , writeLastN -- drop old (ring buffer)

87 -- Stream Folds

@@ -137,12 +137,9 @@ module Streamly.Memory.Array

137 )

138where

139

140−import Control.Monad (when)

141import Control.Monad.IO.Class (MonadIO(..))

140import Control.Monad.IO.Class (MonadIO(..))

142-- import Data.Functor.Identity (Identity)

141-- import Data.Functor.Identity (Identity)

143import Foreign.ForeignPtr (withForeignPtr)

142import Foreign.ForeignPtr (withForeignPtr)

144−import Foreign.ForeignPtr.Unsafe (unsafeForeignPtrToPtr)

145−import Foreign.Ptr (minusPtr, plusPtr)

146import Foreign.Storable (Storable(..))

143import Foreign.Storable (Storable(..))

147import Prelude hiding (length, null, last, map, (!!), read, concat)

144import Prelude hiding (length, null, last, map, (!!), read, concat)

148

145

@@ -180,51 +177,6 @@ _writeNS n m = do

180 if n < 0 then error "writeN: negative write count specified" else return ()

177 if n < 0 then error "writeN: negative write count specified" else return ()

181 A.fromStreamDN n $ D.toStreamD m

178 A.fromStreamDN n $ D.toStreamD m

182

179

183−-- XXX The realloc based implementation needs to make one extra copy if we use

184−-- shrinkToFit. On the other hand, the stream of arrays implementation may

185−-- buffer the array chunk pointers in memory but it does not have to shrink as

186−-- we know the exact size in the end. However, memory copying does not seems to

187−-- be as expensive as the allocations. Therefore, we need to reduce the number

188−-- of allocations instead. Also, the size of allocations matters, right sizing

189−-- an allocation even at the cost of copying sems to help. Should be measured

190−-- on a big stream with heavy calls to toArray to see the effect.

191−--

192−-- XXX check if GHC's memory allocator is efficient enough. We can try the C

193−-- malloc to compare against.

194−

195−{-# INLINE_NORMAL toArrayMinChunk #-}

196−toArrayMinChunk :: forall m a. (MonadIO m, Storable a)

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

198−-- toArrayMinChunk n = FL.mapM spliceArrays $ toArraysOf n

199−toArrayMinChunk elemCount = Fold step initial extract

200−

201− where

202−

203− insertElem (Array start end bound) x = do

204− liftIO $ poke end x

205− return $ Array start (end `plusPtr` sizeOf (undefined :: a)) bound

206−

207− initial = do

208− when (elemCount < 0) $ error "toArrayMinChunk: elemCount is negative"

209− liftIO $ A.newArray elemCount

210− step arr@(Array start end bound) x | end == bound = do

211− let p = unsafeForeignPtrToPtr start

212− oldSize = end `minusPtr` p

213− newSize = max (oldSize * 2) 1

214− arr1 <- liftIO $ A.realloc newSize arr

215− insertElem arr1 x

216− step arr x = insertElem arr x

217− extract = liftIO . A.shrinkToFit

218−

219−-- | Fold the whole input to a single array.

220−--

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

222−--

223−-- @since 0.7.0

224−{-# INLINE write #-}

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

226−write = toArrayMinChunk (A.bytesToElemCount (undefined :: a) (A.mkChunkSize 1024))

227−

228-- | Create an 'Array' from a stream. This is useful when we want to create a

180-- | Create an 'Array' from a stream. This is useful when we want to create a

229-- single array from a stream of unknown size. 'writeN' is at least twice

181-- single array from a stream of unknown size. 'writeN' is at least twice

230-- as efficient when the size is already known.

182-- as efficient when the size is already known.

@@ -235,7 +187,7 @@ write = toArrayMinChunk (A.bytesToElemCount (undefined :: a) (A.mkChunkSize 1024

235--

187--

236{-# INLINE _writeS #-}

188{-# INLINE _writeS #-}

237_writeS :: (MonadIO m, Storable a) => SerialT m a -> m (Array a)

189_writeS :: (MonadIO m, Storable a) => SerialT m a -> m (Array a)

238−_writeS = P.runFold write

190+_writeS = P.runFold A.write

239-- write m = A.fromStreamD $ D.toStreamD m

191-- write m = A.fromStreamD $ D.toStreamD m

240

192

241-------------------------------------------------------------------------------

193-------------------------------------------------------------------------------

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

458{-# INLINE _runTransform #-}

410{-# INLINE _runTransform #-}

459_runTransform :: (MonadIO m, Storable a, Storable b)

411_runTransform :: (MonadIO m, Storable a, Storable b)

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

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

461−_runTransform f arr = P.runFold (toArrayMinChunk (length arr)) $ f (A.read arr)

413+_runTransform f arr = P.runFold (A.toArrayMinChunk (length arr)) $ f (A.read arr)

462

414

463-- | Fold an array using a 'Fold'.

415-- | Fold an array using a 'Fold'.

464--

416--

src/Streamly/Memory/Array/Types.hs+55−2

@@ -61,8 +61,10 @@ module Streamly.Memory.Array.Types

61 , toStreamK

62 , toStreamKRev

63 , toList

64+ , toArrayMinChunk

64 , writeN

65 , writeN

65 , writeNUnsafe

66 , writeNUnsafe

67+ , write

66 , read

68 , read

67 , readU

69 , readU

@@ -551,17 +553,63 @@ writeNUnsafe n = Fold step initial extract

551 return $ (ArrayUnsafe start (end `plusPtr` sizeOf (undefined :: a)))

553 return $ (ArrayUnsafe start (end `plusPtr` sizeOf (undefined :: a)))

552 extract (ArrayUnsafe start end) = return $ Array start end end -- liftIO . shrinkToFit

554 extract (ArrayUnsafe start end) = return $ Array start end end -- liftIO . shrinkToFit

553

555

556+-- XXX The realloc based implementation needs to make one extra copy if we use

557+-- shrinkToFit. On the other hand, the stream of arrays implementation may

558+-- buffer the array chunk pointers in memory but it does not have to shrink as

559+-- we know the exact size in the end. However, memory copying does not seems to

560+-- be as expensive as the allocations. Therefore, we need to reduce the number

561+-- of allocations instead. Also, the size of allocations matters, right sizing

562+-- an allocation even at the cost of copying sems to help. Should be measured

563+-- on a big stream with heavy calls to toArray to see the effect.

564+--

565+-- XXX check if GHC's memory allocator is efficient enough. We can try the C

566+-- malloc to compare against.

567+

568+{-# INLINE_NORMAL toArrayMinChunk #-}

569+toArrayMinChunk :: forall m a. (MonadIO m, Storable a)

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

571+-- toArrayMinChunk n = FL.mapM spliceArrays $ toArraysOf n

572+toArrayMinChunk elemCount = Fold step initial extract

573+

574+ where

575+

576+ insertElem (Array start end bound) x = do

577+ liftIO $ poke end x

578+ return $ Array start (end `plusPtr` sizeOf (undefined :: a)) bound

579+

580+ initial = do

581+ when (elemCount < 0) $ error "toArrayMinChunk: elemCount is negative"

582+ liftIO $ newArray elemCount

583+ step arr@(Array start end bound) x | end == bound = do

584+ let p = unsafeForeignPtrToPtr start

585+ oldSize = end `minusPtr` p

586+ newSize = max (oldSize * 2) 1

587+ arr1 <- liftIO $ realloc newSize arr

588+ insertElem arr1 x

589+ step arr x = insertElem arr x

590+ extract = liftIO . shrinkToFit

591+

592+-- | Fold the whole input to a single array.

593+--

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

595+--

596+-- @since 0.7.0

597+{-# INLINE write #-}

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

599+write = toArrayMinChunk (bytesToElemCount (undefined :: a)

600+ (mkChunkSize 1024))

601+

554{-# INLINE_NORMAL fromStreamDN #-}

602{-# INLINE_NORMAL fromStreamDN #-}

555fromStreamDN :: forall m a. (MonadIO m, Storable a)

603fromStreamDN :: forall m a. (MonadIO m, Storable a)

556 => Int -> D.Stream m a -> m (Array a)

604 => Int -> D.Stream m a -> m (Array a)

557fromStreamDN limit str = do

605fromStreamDN limit str = do

558 arr <- liftIO $ newArray limit

606 arr <- liftIO $ newArray limit

559− end <- D.foldlM' write (aEnd arr) $ D.take limit str

607+ end <- D.foldlM' fwrite (aEnd arr) $ D.take limit str

560 return $ arr {aEnd = end}

608 return $ arr {aEnd = end}

561

609

562 where

610 where

563

611

564− write ptr x = do

612+ fwrite ptr x = do

565 liftIO $ poke ptr x

613 liftIO $ poke ptr x

566 return $ ptr `plusPtr` sizeOf (undefined :: a)

614 return $ ptr `plusPtr` sizeOf (undefined :: a)

567

615

@@ -688,6 +736,11 @@ fromStreamD m = do

688 buffered <- D.foldr K.cons K.nil s

736 buffered <- D.foldr K.cons K.nil s

689 len <- K.foldl' (+) 0 (K.map length buffered)

737 len <- K.foldl' (+) 0 (K.map length buffered)

690 fromStreamDN len $ flattenArrays $ D.fromStreamK buffered

738 fromStreamDN len $ flattenArrays $ D.fromStreamK buffered

739+{-

740+fromStreamD m = runFold write m

741+ where

742+ runFold (Fold step begin done) = D.foldlMx' step begin done

743+-}

691

744

692-- | Convert an 'Array' into a list.

745-- | Convert an 'Array' into a list.

693--

746--