da29139 · streamly · git.iko.soy

git.iko.soy/streamly

Use Unfold to read and concat arrays efficiently

The performance is similar to flattenArrays, the code fuses perfectly. But in the linecount benchmark Unfold case performs 50% slower even though this is not due to fusion, just less optimal code being generated.

commitda29139f455161bdd56df425a8e3486c27db76c9 authorHarendra Kumar <harendra.kumar@gmail.com> dateSun 18 Aug 2019, 05:00 (7 years ago) parent3b0f6df

6 files changed+60

compare with master

unifiedsplit

benchmark/FileIO.hs+6 src/Streamly/Internal.hs+10 src/Streamly/Memory/Array/Types.hs+27 src/Streamly/Memory/ArrayStream.hs+2 src/Streamly/Prelude/Internal.hs+11 src/Streamly/Unfold.hs+4

benchmark/FileIO.hs+6

@@ -133,6 +133,12 @@ main = do

133133 $ SS.foldLines FL.drain

134134 $ SS.decodeChar8

135135 $ FH.read inh

136+ , mkBench "linecountU" href $ do

137+ Handles inh _ <- readIORef href

138+ S.length

139+ $ SS.foldLines FL.drain

140+ $ SS.decodeChar8

141+ $ Internal.concatMapU Internal.readU (FH.readArrays inh)

136142 , mkBench "sum" href $ do

137143 Handles inh _ <- readIORef href

138144 S.sum $ FH.read inh

src/Streamly/Internal.hs+10

@@ -52,6 +52,8 @@ module Streamly.Internal

5252 , classifySessionsBy

5353 , classifyKeepAliveSessions

5454

55+ , concatMapU

56+

5557 -- * Streamly.Fold Experimental Exports

5658 , Fold (..)

5759 , lmap

@@ -67,6 +69,12 @@ module Streamly.Internal

6769 , initialize

6870 , runStep

6971

72+ -- * Streamly.Unfold Experimental Exports

73+ , Unfold (..)

74+

75+ -- * Streamly.Memory.Array

76+ , readU

77+

7078 -- * Streamly.FileSystem.Handle

7179 , writeS

7280 )

@@ -75,8 +83,10 @@ where

7583import Prelude hiding (break, span, splitAt)

7684

7785import Streamly.FileSystem.Handle.Internal (writeS)

86+import Streamly.Memory.Array.Types (readU)

7887import Streamly.Streams.Combinators (inspectMode)

7988import Streamly.Streams.Parallel (tapAsync)

89+import Streamly.Unfold.Types (Unfold(..))

8090

8191import Streamly.Prelude.Internal

8292import Streamly.Fold.Types

src/Streamly/Memory/Array/Types.hs+27

@@ -62,6 +62,7 @@ module Streamly.Memory.Array.Types

6262 , toList

6363 , writeN

6464 , read

65+ , readU

6566

6667 -- * Utilities

6768 , defaultChunkSize

@@ -100,6 +101,7 @@ import GHC.IO (IO(IO), unsafePerformIO)

100101import GHC.Ptr (Ptr(..))

101102

102103import Streamly.Fold.Types (Fold(..))

104+import Streamly.Unfold.Types (Unfold(..))

103105import Streamly.Strict (Tuple'(..))

104106import Streamly.SVar (adaptState)

105107

@@ -388,6 +390,31 @@ byteCapacity Array{..} =

388390 len = aBound `minusPtr` p

389391 in assert (len >= 0) len

390392

393+{-# INLINE_NORMAL readU #-}

394+readU :: forall m a. (Monad m, Storable a) => Unfold m (Array a) a

395+readU = Unfold step inject

396+ where

397+

398+ -- Note we have repurposed the Array type here

399+ inject Array{..} =

400+ return $ Array aStart (unsafeForeignPtrToPtr aStart) aEnd

401+

402+ {-# INLINE_LATE step #-}

403+ step (Array _ p end) | p == end = return D.Stop

404+ step (Array start p end) = do

405+ -- unsafeInlineIO allows us to run this in Identity monad for pure

406+ -- toList/foldr case which makes them much faster due to not

407+ -- accumulating the list and fusing better with the pure consumers.

408+ --

409+ -- This should be safe as the array contents are guaranteed to be

410+ -- evaluated/written to before we peek at them.

411+ let !x = unsafeInlineIO $ do

412+ r <- peek p

413+ touchForeignPtr start

414+ return r

415+ return $ D.Yield x

416+ (Array start (p `plusPtr` (sizeOf (undefined :: a))) end)

417+

391418{-# INLINE_NORMAL toStreamD #-}

392419toStreamD :: forall m a. (Monad m, Storable a) => Array a -> D.Stream m a

393420toStreamD Array{..} =

src/Streamly/Memory/ArrayStream.hs+2

@@ -68,6 +68,8 @@ import qualified Streamly.Streams.Prelude as P

6868{-# INLINE concat #-}

6969concat :: (IsStream t, MonadIO m, Storable a) => t m (Array a) -> t m a

7070concat m = D.fromStreamD $ A.flattenArrays (D.toStreamD m)

71+-- concat m = D.fromStreamD $ D.concatMap A.toStreamD (D.toStreamD m)

72+-- concat m = D.fromStreamD $ D.concatMapU A.readU (D.toStreamD m)

7173

7274-- XXX should we have a reverseArrays API to reverse the stream of arrays

7375-- instead?

src/Streamly/Prelude/Internal.hs+11

@@ -231,6 +231,7 @@ module Streamly.Prelude.Internal

231231

232232 -- ** Nested Streams

233233 , concatMapM

234+ , concatMapU

234235 , concatMap

235236 , concatMapWith

236237 -- , interposeBy

@@ -379,6 +380,7 @@ import qualified System.IO as IO

379380

380381import Streamly.Enumeration (Enumerable(..), enumerate, enumerateTo)

381382import Streamly.Fold.Types (Fold (..))

383+import Streamly.Unfold.Types (Unfold)

382384import Streamly.Memory.Array.Types (Array)

383385-- import Streamly.Memory.Ring (Ring)

384386import Streamly.SVar (MonadAsync, defState)

@@ -2099,6 +2101,15 @@ concatMap f m = fromStreamD $ D.concatMap (toStreamD . f) (toStreamD m)

20992101concatMapM :: (IsStream t, Monad m) => (a -> m (t m b)) -> t m a -> t m b

21002102concatMapM f m = fromStreamD $ D.concatMapM (fmap toStreamD . f) (toStreamD m)

21012103

2104+-- | Like 'concatMap' but uses an 'Unfold' for stream generation. Unlike

2105+-- 'concatMap' this can fuse the 'Unfold' code with the inner loop and

2106+-- therefore provide many times better performance.

2107+--

2108+-- @since 0.7.0

2109+{-# INLINE concatMapU #-}

2110+concatMapU ::(IsStream t, Monad m) => Unfold m a b -> t m a -> t m b

2111+concatMapU u m = fromStreamD $ D.concatMapU u (toStreamD m)

2112+

21022113{-

21032114-- | A generalization of insertBy to inserting sequences.

21042115interposeBy :: (IsStream t, Monad m)

src/Streamly/Unfold.hs+4

@@ -2,6 +2,7 @@

22{-# LANGUAGE CPP #-}

33{-# LANGUAGE ExistentialQuantification #-}

44{-# LANGUAGE FlexibleContexts #-}

5+{-# LANGUAGE PatternSynonyms #-}

56{-# LANGUAGE RankNTypes #-}

67{-# LANGUAGE RecordWildCards #-}

78{-# LANGUAGE ScopedTypeVariables #-}

@@ -34,6 +35,9 @@ module Streamly.Unfold

3435where

3536

3637import Streamly.Streams.StreamD.Type (Stream(..), Step(..))

38+#if __GLASGOW_HASKELL__ < 800

39+import Streamly.Streams.StreamD.Type (pattern Stream)

40+#endif

3741import Streamly.Unfold.Types (Unfold(..))

3842

3943-------------------------------------------------------------------------------