a0a2281 · streamly · git.iko.soy

git.iko.soy/streamly

on master help sign in

Add unfoldMany to Data.Fold

a0a2281Adithya Kumar5 years ago

benchmark/Streamly/Benchmark/Data/Fold.hs+13

compare with master

unifiedsplit

@@ -20,6 +20,7 @@ import Streamly.Internal.Data.Fold (Fold(..))

21import qualified Data.Map.Strict as Map

22import qualified Streamly.Internal.Data.Fold as FL

23+import qualified Streamly.Internal.Data.Unfold as Unfold

23import qualified Streamly.Internal.Data.Pipe as Pipe

24import qualified Streamly.Internal.Data.Pipe as Pipe

24import qualified Streamly.Internal.Data.Stream.IsStream as IP

25import qualified Streamly.Internal.Data.Stream.IsStream as IP

25import qualified Streamly.Prelude as S

26import qualified Streamly.Prelude as S

@@ -233,6 +234,17 @@ unzipWithMinM = do

233 let f = \a -> return (a + 1, a)

234 let f = \a -> return (a + 1, a)

234 IP.fold (FL.unzipWithMinM f FL.sum FL.length)

235 IP.fold (FL.unzipWithMinM f FL.sum FL.length)

235

236

237+-------------------------------------------------------------------------------

238+-- Nested

239+-------------------------------------------------------------------------------

240+

241+{-# INLINE unfoldMany #-}

242+unfoldMany :: Int -> Benchmarkable

243+unfoldMany val =

244+ nfIO

245+ $ IP.fold (FL.unfoldMany (Unfold.replicateM val) FL.drain)

246+ $ IP.fromPure (randomRIO (1, 1 :: Int))

247+

236-------------------------------------------------------------------------------

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

237-- Benchmarks

249-- Benchmarks

238-------------------------------------------------------------------------------

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

@@ -365,6 +377,7 @@ o_1_space_serial_composition value =

365 , benchIOSink value "classifyWith sum" $ classifyWith (fst . fn)

377 , benchIOSink value "classifyWith sum" $ classifyWith (fst . fn)

366 , benchIOSink value "classifyScanWith sum"

378 , benchIOSink value "classifyScanWith sum"

367 $ classifyScanWith (fst . fn)

379 $ classifyScanWith (fst . fn)

380+ , bench "unfoldMany" $ unfoldMany value

368 ]

381 ]

369 ]

382 ]

370

383

core/src/Streamly/Internal/Data/Fold.hs+38

@@ -236,6 +236,7 @@ module Streamly.Internal.Data.Fold

236 , chunksBetween

237

238 -- ** Nesting

239+ , unfoldMany

239 , concatSequence

240 , concatSequence

240 , concatMap

241 , concatMap

241

242

@@ -257,6 +258,8 @@ module Streamly.Internal.Data.Fold

257 )

258 )

258where

259where

259

260

261+#include "inline.hs"

262+

260import Data.Bifunctor (first)

263import Data.Bifunctor (first)

261import Data.Either (isLeft, isRight, fromLeft, fromRight)

264import Data.Either (isLeft, isRight, fromLeft, fromRight)

262import Data.Functor.Identity (Identity(..))

265import Data.Functor.Identity (Identity(..))

@@ -266,6 +269,7 @@ import Data.Maybe (isJust, fromJust)

266import Streamly.Internal.Data.Either.Strict

269import Streamly.Internal.Data.Either.Strict

267 (Either'(..), fromLeft', fromRight', isLeft', isRight')

270 (Either'(..), fromLeft', fromRight', isLeft', isRight')

268import Streamly.Internal.Data.Pipe.Type (Pipe (..), PipeState(..))

271import Streamly.Internal.Data.Pipe.Type (Pipe (..), PipeState(..))

272+import Streamly.Internal.Data.Unfold.Type (Unfold(..))

269import Streamly.Internal.Data.Tuple.Strict (Tuple'(..), Tuple3'(..))

273import Streamly.Internal.Data.Tuple.Strict (Tuple'(..), Tuple3'(..))

270import Streamly.Internal.Data.Stream.Serial (SerialT(..))

274import Streamly.Internal.Data.Stream.Serial (SerialT(..))

271

275

@@ -273,6 +277,7 @@ import qualified Data.IntSet as IntSet

273import qualified Data.Map.Strict as Map

277import qualified Data.Map.Strict as Map

274import qualified Data.Set as Set

278import qualified Data.Set as Set

275import qualified Streamly.Internal.Data.Pipe.Type as Pipe

279import qualified Streamly.Internal.Data.Pipe.Type as Pipe

280+import qualified Streamly.Internal.Data.Stream.StreamD.Type as StreamD

276-- import qualified Streamly.Internal.Data.Stream.IsStream.Enumeration as Stream

281-- import qualified Streamly.Internal.Data.Stream.IsStream.Enumeration as Stream

277import qualified Prelude

282import qualified Prelude

278

283

@@ -1921,3 +1926,36 @@ toStream = fmap SerialT toStreamK

1921{-# INLINE toStreamRev #-}

1926{-# INLINE toStreamRev #-}

1922toStreamRev :: Monad m => Fold m a (SerialT n a)

1927toStreamRev :: Monad m => Fold m a (SerialT n a)

1923toStreamRev = fmap SerialT toStreamKRev

1928toStreamRev = fmap SerialT toStreamKRev

1929+

1930+-- XXX This does not fuse. It contains a recursive step function. We will need

1931+-- a Skip input constructor in the fold type to make it fuse.

1932+--

1933+-- | Unfold and flatten the input stream of a fold.

1934+--

1935+-- @

1936+-- Stream.fold (unfoldMany u f) = Stream.fold f . Stream.unfoldMany u

1937+-- @

1938+--

1939+-- /Pre-release/

1940+{-# INLINE unfoldMany #-}

1941+unfoldMany :: Monad m => Unfold m a b -> Fold m b c -> Fold m a c

1942+unfoldMany (Unfold ustep inject) (Fold fstep initial extract) =

1943+ Fold consume initial extract

1944+

1945+ where

1946+

1947+ {-# INLINE produce #-}

1948+ produce fs us = do

1949+ ures <- ustep us

1950+ case ures of

1951+ StreamD.Yield b us1 -> do

1952+ fres <- fstep fs b

1953+ case fres of

1954+ Partial fs1 -> produce fs1 us1

1955+ -- XXX What to do with the remaining stream?

1956+ Done c -> return $ Done c

1957+ StreamD.Skip us1 -> produce fs us1

1958+ StreamD.Stop -> return $ Partial fs

1959+

1960+ {-# INLINE_LATE consume #-}

1961+ consume s a = inject a >>= produce s