a0a2281 · streamly · git.iko.soy

git.iko.soy/streamly

Add unfoldMany to Data.Fold

commita0a228196cbc8d2c08a654e0698b2d496783eba7 authorAdithya Kumar <adithya@composewell.com> dateSat 12 Feb 2022, 09:32 (5 years ago) parent324f2e5

2 files changed+51

compare with master

unifiedsplit

benchmark/Streamly/Benchmark/Data/Fold.hs+13 core/src/Streamly/Internal/Data/Fold.hs+38

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

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

2020

2121import qualified Data.Map.Strict as Map

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

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

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

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

2526import qualified Streamly.Prelude as S

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

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

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

235236

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+

236248-------------------------------------------------------------------------------

237249-- Benchmarks

238250-------------------------------------------------------------------------------

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

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

366378 , benchIOSink value "classifyScanWith sum"

367379 $ classifyScanWith (fst . fn)

380+ , bench "unfoldMany" $ unfoldMany value

368381 ]

369382 ]

370383

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

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

236236 , chunksBetween

237237

238238 -- ** Nesting

239+ , unfoldMany

239240 , concatSequence

240241 , concatMap

241242

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

257258 )

258259where

259260

261+#include "inline.hs"

262+

260263import Data.Bifunctor (first)

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

262265import Data.Functor.Identity (Identity(..))

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

266269import Streamly.Internal.Data.Either.Strict

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

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

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

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

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

271275

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

273277import qualified Data.Map.Strict as Map

274278import qualified Data.Set as Set

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

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

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

277282import qualified Prelude

278283

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

19211926{-# INLINE toStreamRev #-}

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

19231928toStreamRev = 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