d0bba76 · streamly · git.iko.soy

git.iko.soy/streamly

on master help sign in

Add toArray fold to fold a stream into a single array

d0bba76Harendra Kumar7 years ago

examples/FileSinkServer.hs+3−3

compare with master

unifiedsplit

@@ -22,12 +22,12 @@ main = do

22 File.append file

23 $ encodeChar8Unchecked

24 $ S.concatMap A.read

25− $ S.concatMapBy parallel recv

25+ $ S.concatMapBy parallel (flip NS.withSocketS recv)

26 $ NS.connectionsOnAllAddrs 8090

28 where

30 recv =

31− FL.splitSuffixBy' (== '\n') (A.toArrayN 1024)

31+ FL.splitSuffixBy' (== '\n') A.toArray

32 . decodeChar8

33− . NS.fromSocket

33+ . NS.read

src/Streamly/Mem/Array.hs+101−14

@@ -99,7 +99,7 @@ module Streamly.Mem.Array

99 -- Folds

100 , toArrayN

101 -- , toArrays

102− -- , toArray

102+ , toArray

103

104 -- Streams

105 , arraysOf

@@ -132,8 +132,9 @@ module Streamly.Mem.Array

132where

133

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

135−import Data.Functor.Identity

135+import Data.Functor.Identity (Identity)

136import Foreign.ForeignPtr (withForeignPtr)

137+import Foreign.ForeignPtr.Unsafe (unsafeForeignPtrToPtr)

137import Foreign.Ptr (minusPtr, plusPtr, castPtr)

138import Foreign.Ptr (minusPtr, plusPtr, castPtr)

138import Foreign.Storable (Storable(..))

139import Foreign.Storable (Storable(..))

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

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

@@ -141,11 +142,13 @@ import Prelude hiding (length, null, last, map, (!!), read)

141import Streamly.Mem.Array.Types hiding (flattenArrays, newArray)

142import Streamly.Mem.Array.Types hiding (flattenArrays, newArray)

142import Streamly.Streams.Serial (SerialT)

143import Streamly.Streams.Serial (SerialT)

143import Streamly.Streams.StreamK.Type (IsStream)

144import Streamly.Streams.StreamK.Type (IsStream)

144−import Streamly.Fold (Fold, lchunksOf, toStream)

145+import Streamly.Fold.Types (Fold(..))

145

146

147+import qualified Streamly.Fold as FL

146import qualified Streamly.Mem.Array.Types as A

148import qualified Streamly.Mem.Array.Types as A

147import qualified Streamly.Prelude as S

149import qualified Streamly.Prelude as S

148import qualified Streamly.Streams.StreamD as D

150import qualified Streamly.Streams.StreamD as D

151+import qualified Streamly.Streams.Prelude as P

149

152

150-------------------------------------------------------------------------------

153-------------------------------------------------------------------------------

151-- Construction

154-- Construction

@@ -353,11 +356,63 @@ writeSliceRev arr i len s = undefined

353-- Streams of Arrays

356-- Streams of Arrays

354-------------------------------------------------------------------------------

357-------------------------------------------------------------------------------

355

358

359+-- exponentially increasing sizes of the chunks upto the max limit.

360+-- XXX this will be easier to implement with parsers/terminating folds

361+-- With this we should be able to reduce the number of chunks/allocations.

362+-- The reallocation/copy based toArray can also be implemented using this.

363+--

364+{-

365+{-# INLINE toArraysInRange #-}

366+toArraysInRange :: (IsStream t, MonadIO m, Storable a)

367+ => Int -> Int -> Fold m (Array a) b -> Fold m a b

368+toArraysInRange low high (Fold step initial extract) =

369+-}

370+

356-- | Fold the input to a pure buffered stream (List) of arrays.

371-- | Fold the input to a pure buffered stream (List) of arrays.

357−{-# INLINE _toArrays #-}

372+{-# INLINE _toArraysOf #-}

358−_toArrays :: (MonadIO m, Storable a)

373+_toArraysOf :: (MonadIO m, Storable a)

359 => Int -> Fold m a (SerialT Identity (Array a))

374 => Int -> Fold m a (SerialT Identity (Array a))

360−_toArrays n = lchunksOf n (toArrayN n) toStream

375+_toArraysOf n = FL.lchunksOf n (toArrayN n) FL.toStream

376+

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

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

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

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

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

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

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

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

385+--

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

387+-- malloc to compare against.

388+

389+{-# INLINE toArrayMinChunk #-}

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

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

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

393+toArrayMinChunk chunkSize = Fold step initial extract

394+

395+ where

396+

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

398+ liftIO $ poke end x

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

400+

401+ initial = liftIO $ A.newArray (mkChunkSize chunkSize)

402+ step arr@(Array _ end bound) x | end == bound = do

403+ arr1 <- liftIO $ reallocDouble 1 arr

404+ insertElem arr1 x

405+ step arr x = insertElem arr x

406+ extract = liftIO . shrinkToFit

407+

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

409+--

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

411+--

412+-- @since 0.7.0

413+{-# INLINE toArray #-}

414+toArray :: forall m a. (MonadIO m, Storable a) => Fold m a (Array a)

415+toArray = toArrayMinChunk 1024

361

416

362-- | Convert a stream of arrays into a stream of their elements.

417-- | Convert a stream of arrays into a stream of their elements.

363--

418--

@@ -389,14 +444,14 @@ arraysOf :: (IsStream t, MonadIO m, Storable a)

389arraysOf n str =

444arraysOf n str =

390 D.fromStreamD $ fromStreamDArraysOf n (D.toStreamD str)

445 D.fromStreamD $ fromStreamDArraysOf n (D.toStreamD str)

391

446

392−-- | Given a stream of arrays, splice them all together to generate a single

447+-- XXX Both of these implementations of splicing seem to perform equally well.

393−-- array. The stream must be /finite/.

448+-- We need to perform benchmarks over a range of sizes though.

394−--

449+

395−-- @since 0.7.0

450+{-# INLINE _spliceArraysRealloced #-}

396−{-# INLINABLE spliceArrays #-}

451+_spliceArraysRealloced :: (MonadIO m, Storable a)

397−spliceArrays :: (MonadIO m, Storable a) => SerialT m (Array a) -> m (Array a)

452+ => SerialT m (Array a) -> m (Array a)

398−spliceArrays s = do

453+_spliceArraysRealloced s = do

399− buffered <- S.foldr S.cons S.nil s

454+ buffered <- P.foldr S.cons S.nil s

400 len <- S.sum (S.map length buffered)

455 len <- S.sum (S.map length buffered)

401

456

402 arr <- liftIO $ A.newArray len

457 arr <- liftIO $ A.newArray len

@@ -411,6 +466,38 @@ spliceArrays s = do

411 memcpy (castPtr dst) (castPtr src) len

466 memcpy (castPtr dst) (castPtr src) len

412 return $ dst `plusPtr` len

467 return $ dst `plusPtr` len

413

468

469+{-# INLINE spliceArraysBuffered #-}

470+spliceArraysBuffered :: (MonadIO m, Storable a)

471+ => SerialT m (Array a) -> m (Array a)

472+spliceArraysBuffered s = do

473+ idst <- liftIO $ A.newArray (mkChunkSizeKB 4)

474+ arr <- S.foldlM' appendArr idst s

475+ liftIO $ shrinkToFit arr

476+

477+ where

478+

479+ appendArr dst@(Array _ end bound) src = liftIO $ do

480+ let srcLen = aEnd src `minusPtr` unsafeForeignPtrToPtr (aStart src)

481+ dst1 <-

482+ if end `plusPtr` srcLen >= bound

483+ then reallocDouble srcLen dst

484+ else return dst

485+

486+ withForeignPtr (aStart dst1) $ \_ -> do

487+ withForeignPtr (aStart src) $ \psrc -> do

488+ let pdst = aEnd dst1

489+ memcpy (castPtr pdst) (castPtr psrc) srcLen

490+ return $ dst1 { aEnd = pdst `plusPtr` srcLen }

491+

492+-- | Given a stream of arrays, splice them all together to generate a single

493+-- array. The stream must be /finite/.

494+--

495+-- @since 0.7.0

496+{-# INLINE spliceArrays #-}

497+spliceArrays :: (MonadIO m, Storable a) => SerialT m (Array a) -> m (Array a)

498+spliceArrays = spliceArraysBuffered

499+-- spliceArrays = _spliceArraysRealloced

500+

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

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

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

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

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

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

src/Streamly/Mem/Array/Types.hs+31−15

@@ -54,7 +54,9 @@ module Streamly.Mem.Array.Types

55 -- * Utilities

56 , defaultChunkSize

57+ , mkChunkSize

57 , mkChunkSizeKB

58 , mkChunkSizeKB

59+ , reallocDouble

58 )

60 )

59where

61where

@@ -241,6 +243,26 @@ shrinkToFit arr@Array{..} = do

241 }

243 }

242 else return arr

244 else return arr

243

245

246+-- | Expand the free space in the array doubling the size with a minimum

247+-- exapnsion of the specified amount of bytes.

248+{-# NOINLINE reallocDouble #-}

249+reallocDouble :: forall a. Storable a => Int -> Array a -> IO (Array a)

250+reallocDouble minIncrease Array{..} = do

251+ assert (aEnd <= aBound) (return ())

252+ let oldStart = unsafeForeignPtrToPtr aStart

253+ let size = aEnd `minusPtr` oldStart

254+ newSize = max (size * 2) (size + minIncrease)

255+ newPtr <- mallocPlainForeignPtrAlignedBytes

256+ newSize (alignment (undefined :: a))

257+ withForeignPtr newPtr $ \pNew -> do

258+ memcpy (castPtr pNew) (castPtr oldStart) size

259+ touchForeignPtr aStart

260+ return $ Array

261+ { aStart = newPtr

262+ , aEnd = pNew `plusPtr` size

263+ , aBound = pNew `plusPtr` newSize

264+ }

265+

244-- XXX when converting an array of Word8 from a literal string we can simply

266-- XXX when converting an array of Word8 from a literal string we can simply

245-- refer to the literal string. Is it possible to write rules such that

267-- refer to the literal string. Is it possible to write rules such that

246-- fromList Word8 can be rewritten so that GHC does not first convert the

268-- fromList Word8 can be rewritten so that GHC does not first convert the

@@ -367,16 +389,7 @@ toArrayN n = Fold step initial extract

367 step (Array start end bound) x = do

389 step (Array start end bound) x = do

368 liftIO $ poke end x

390 liftIO $ poke end x

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

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

370− extract = return

392+ extract = liftIO . shrinkToFit

371−

372−{-

373−-- This can be implemented by combining the List of arrays from toArrays into a

374−-- single array.

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

376−{-# INLINE toArray #-}

377−toArray :: forall m a. (Monad m, Storable a) => Fold m a (Array a)

378−toArray = Fold step begin done

379−-}

380

393

381{-# INLINE_NORMAL fromStreamDN #-}

394{-# INLINE_NORMAL fromStreamDN #-}

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

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

@@ -638,13 +651,16 @@ read = D.fromStreamD . toStreamD

638allocOverhead :: Int

651allocOverhead :: Int

639allocOverhead = 2 * sizeOf (undefined :: Int)

652allocOverhead = 2 * sizeOf (undefined :: Int)

640

653

654+mkChunkSize :: Int -> Int

655+mkChunkSize n = n - allocOverhead

656+

657+mkChunkSizeKB :: Int -> Int

658+mkChunkSizeKB n = mkChunkSize (n * k)

659+ where k = 1024

660+

641-- | Default maximum buffer size in bytes, for reading from and writing to IO

661-- | Default maximum buffer size in bytes, for reading from and writing to IO

642-- devices, the value is 32KB minus GHC allocation overhead, which is a few

662-- devices, the value is 32KB minus GHC allocation overhead, which is a few

643-- bytes, so that the actual allocation is 32KB.

663-- bytes, so that the actual allocation is 32KB.

644defaultChunkSize :: Int

664defaultChunkSize :: Int

645−defaultChunkSize = 32 * k - allocOverhead

665+defaultChunkSize = mkChunkSizeKB 32

646− where k = 1024

647

666

648−mkChunkSizeKB :: Int -> Int

649−mkChunkSizeKB n = n * k - allocOverhead

650− where k = 1024