6b3a8c4 · streamly · git.iko.soy

git.iko.soy/streamly

on master help sign in

Move the ArrayStream splicing code to MutArray module

6b3a8c4Harendra Kumar3 years ago

core/src/Streamly/Internal/Data/Array/Stream.hs+10−78

compare with master

unifiedsplit

@@ -236,6 +236,15 @@ compact :: (MonadIO m, Unbox a)

236 => Int -> Stream m (Array a) -> Stream m (Array a)

237compact = packArraysChunksOf

238

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

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

241+--

242+-- @since 0.7.0

243+{-# INLINE toArray #-}

244+toArray :: (MonadIO m, Unbox a) => Stream m (Array a) -> m (Array a)

245+toArray s =

246+ fmap A.unsafeFreeze $ MA.fromArrayStreamRealloced (fmap A.unsafeThaw s)

247+

239-------------------------------------------------------------------------------

248-------------------------------------------------------------------------------

240-- Split

249-- Split

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

250-------------------------------------------------------------------------------

@@ -439,7 +448,7 @@ foldBreak = foldBreakK

439-- foldBreak f = runArrayFoldBreak (ChunkFold.fromFold f)

448-- foldBreak f = runArrayFoldBreak (ChunkFold.fromFold f)

440

449

441-------------------------------------------------------------------------------

450-------------------------------------------------------------------------------

442−-- Fold to a single Array

451+-- Elimination - running element parsers

443-------------------------------------------------------------------------------

452-------------------------------------------------------------------------------

444

453

445-- When we have to take an array partially, take the last part of the array.

454-- When we have to take an array partially, take the last part of the array.

@@ -485,83 +494,6 @@ splitAtArrayListRev n ls

485 arr1 = Array contents end1 end

494 arr1 = Array contents end1 end

486 in ([arr1], arr2:xs)

495 in ([arr1], arr2:xs)

487

496

488−-------------------------------------------------------------------------------

489−-- Fold to a single Array

490−-------------------------------------------------------------------------------

491−

492−{-

493−

494−-- XXX Both of these implementations of splicing seem to perform equally well.

495−-- We need to perform benchmarks over a range of sizes though.

496−

497−-- CAUTION! length must more than equal to lengths of all the arrays in the

498−-- stream.

499−{-# INLINE spliceArraysLenUnsafe #-}

500−spliceArraysLenUnsafe :: (MonadIO m, Unbox a)

501− => Int -> Stream m (MutArray a) -> m (MutArray a)

502−spliceArraysLenUnsafe len buffered = do

503− -- XXX The new array's pinned state should depend on the first element

504− -- of the stream. We should uncons the stream and do the required.

505− arr <- undefined

506− D.foldlM' MA.spliceUnsafe (return arr) buffered

507−

508−{-# INLINE _spliceArrays #-}

509−_spliceArrays :: (MonadIO m, Unbox a)

510− => Stream m (Array a) -> m (Array a)

511−_spliceArrays s = do

512− buffered <- D.foldr K.cons K.nil s

513− len <- K.fold FL.sum (fmap Array.length buffered)

514− -- XXX The new array's pinned state should depend on the first element

515− -- of the stream. We should uncons the stream and do the required.

516− arr <- undefined

517− final <- D.foldlM' writeArr (return arr) (toStream buffered)

518− return $ A.unsafeFreeze final

519−

520− where

521−

522− writeArr dst arr = MA.spliceUnsafe dst (A.unsafeThaw arr)

523−

524−{-# INLINE _spliceArraysBuffered #-}

525−_spliceArraysBuffered :: (MonadIO m, Unbox a)

526− => Stream m (Array a) -> m (Array a)

527−_spliceArraysBuffered s = do

528− buffered <- D.foldr K.cons K.nil s

529− len <- K.fold FL.sum (fmap Array.length buffered)

530− A.unsafeFreeze <$>

531− spliceArraysLenUnsafe len (fmap A.unsafeThaw (toStream buffered))

532−

533−-}

534−

535−{-# INLINE spliceArraysRealloced #-}

536−spliceArraysRealloced :: forall m a. (MonadIO m, Unbox a)

537− => Stream m (Array a) -> m (Array a)

538−spliceArraysRealloced s = do

539− res <- D.uncons s

540− case res of

541− Just (a, strm) -> do

542− arr <-

543− D.foldlM'

544− MA.spliceExp

545− (pure (A.unsafeThaw a))

546− (fmap A.unsafeThaw strm)

547− liftIO $ A.unsafeFreeze <$> MA.rightSize arr

548− Nothing -> pure A.nil

549−

550−-- XXX This should just be "fold A.write"

551−--

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

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

554−--

555−-- @since 0.7.0

556−{-# INLINE toArray #-}

557−toArray :: (MonadIO m, Unbox a) => Stream m (Array a) -> m (Array a)

558−toArray = spliceArraysRealloced

559−-- spliceArrays = _spliceArraysBuffered

560−

561−-------------------------------------------------------------------------------

562−-- Elimination - running element parsers

563−-------------------------------------------------------------------------------

564−

565-- GHC parser does not accept {-# ANN type [] NoSpecConstr #-}, so we need

497-- GHC parser does not accept {-# ANN type [] NoSpecConstr #-}, so we need

566-- to make a newtype.

498-- to make a newtype.

567{-# ANN type List NoSpecConstr #-}

499{-# ANN type List NoSpecConstr #-}

core/src/Streamly/Internal/Data/MutArray/Type.hs+23−2

@@ -78,6 +78,8 @@ module Streamly.Internal.Data.MutArray.Type

78 , fromPureStream

79 , fromByteStr#

80 , fromPtrN

81+ , fromArrayStreamK

82+ , fromArrayStreamRealloced

82 -- ** Random writes

84 -- ** Random writes

83 , putIndex

85 , putIndex

@@ -200,7 +202,6 @@ module Streamly.Internal.Data.MutArray.Type

200 -- *** Eliminate to streams

202 -- *** Eliminate to streams

201 , flattenArrays

203 , flattenArrays

202 , flattenArraysRev

204 , flattenArraysRev

203− , fromArrayStreamK

204

205

205 -- *** Construct from arrays

206 -- *** Construct from arrays

206 -- get chunks without copying

207 -- get chunks without copying

@@ -2090,7 +2091,27 @@ fromByteStr# addr = do

2090 return (arr {arrEnd = lenInt})

2091 return (arr {arrEnd = lenInt})

2091

2092

2092-------------------------------------------------------------------------------

2093-------------------------------------------------------------------------------

2093−-- convert stream to a single array

2094+-- convert a stream of arrays to a single array by reallocating and copying

2095+-------------------------------------------------------------------------------

2096+

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

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

2099+

2100+{-# INLINE fromArrayStreamRealloced #-}

2101+fromArrayStreamRealloced :: forall m a. (MonadIO m, Unbox a)

2102+ => Stream m (MutArray a) -> m (MutArray a)

2103+fromArrayStreamRealloced s = do

2104+ res <- D.uncons s

2105+ case res of

2106+ Just (a, strm) -> do

2107+ arr <- D.foldlM' spliceExp (pure a) strm

2108+ -- Reallocation is exponential so there may be 50% empty space in

2109+ -- worst case. One more reallocation to reclaim the space.

2110+ liftIO $ rightSize arr

2111+ Nothing -> pure nil

2112+

2113+-------------------------------------------------------------------------------

2114+-- convert a stream of arrays to a single array by buffering arrays first

2094-------------------------------------------------------------------------------

2115-------------------------------------------------------------------------------

2095

2116

2096{-# INLINE arrayStreamKLength #-}

2117{-# INLINE arrayStreamKLength #-}