8d806c7 · streamly · git.iko.soy

git.iko.soy/streamly

on master help sign in

Add foldEither and foldConcat

8d806c7Harendra Kumar4 years ago

core/src/Streamly/Internal/Data/Array/Unboxed/Stream.hs−3

compare with master

unifiedsplit

@@ -315,9 +315,6 @@ splitOnSuffix byte s =

315-- Elimination - Running folds

316-------------------------------------------------------------------------------

317

318−-- XXX This should be written using CPS (as foldK) if we want it to scale wrt

319−-- to the number of times it can be called on the same stream.

320−--

321{-# INLINE_NORMAL foldBreakD #-}

318{-# INLINE_NORMAL foldBreakD #-}

322foldBreakD :: forall m a b. (MonadIO m, Unbox a) =>

319foldBreakD :: forall m a b. (MonadIO m, Unbox a) =>

323 Fold m a b -> D.Stream m (Array a) -> m (b, D.Stream m (Array a))

320 Fold m a b -> D.Stream m (Array a) -> m (b, D.Stream m (Array a))

core/src/Streamly/Internal/Data/Stream/Bottom.hs+120−7

@@ -23,6 +23,10 @@ module Streamly.Internal.Data.Stream.Bottom

23 , fold

24 , foldContinue

25 , foldBreak

26+ , foldBreak2

27+ , foldEither

28+ , foldEither2

29+ , foldConcat

27 -- * Scans

31 -- * Scans

28 , smapM

32 , smapM

@@ -61,10 +65,13 @@ where

61#include "inline.hs"

65#include "inline.hs"

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

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

68+import GHC.Types (SPEC(..))

64import Streamly.Internal.Data.Fold.Type (Fold (..))

69import Streamly.Internal.Data.Fold.Type (Fold (..))

65import Streamly.Internal.Data.Time.Units (AbsTime, RelTime64, addToAbsTime64)

70import Streamly.Internal.Data.Time.Units (AbsTime, RelTime64, addToAbsTime64)

66import Streamly.Internal.Data.Unboxed (Unbox)

71import Streamly.Internal.Data.Unboxed (Unbox)

72+import Streamly.Internal.Data.Producer.Type (Producer(..))

67import Streamly.Internal.System.IO (defaultChunkSize)

73import Streamly.Internal.System.IO (defaultChunkSize)

74+import Streamly.Internal.Data.SVar.Type (defState)

69import qualified Streamly.Internal.Data.Array.Unboxed.Type as A

76import qualified Streamly.Internal.Data.Array.Unboxed.Type as A

70import qualified Streamly.Internal.Data.Fold as Fold

77import qualified Streamly.Internal.Data.Fold as Fold

@@ -214,21 +221,127 @@ fold fl strm = D.fold fl $ D.fromStreamK $ toStreamK strm

214--

221--

215{-# INLINE foldBreak #-}

222{-# INLINE foldBreak #-}

216foldBreak :: Monad m => Fold m a b -> Stream m a -> m (b, Stream m a)

223foldBreak :: Monad m => Fold m a b -> Stream m a -> m (b, Stream m a)

217−{-

218−-- XXX This shows quadratic performance when used recursively perhaps because

219−-- of StreamK to StreamD conversions not getting eliminated sue to recursion.

220−foldBreak fl (Stream strm) = fmap f $ D.foldBreak fl $ D.fromStreamK strm

224+foldBreak fl strm = fmap f $ K.foldBreak fl (toStreamK strm)

221

225

222 where

226 where

223

227

224− f (b, str) = (b, Stream (D.toStreamK str))

228+ f (b, str) = (b, fromStreamK str)

225−-}

229+

230+-- XXX The quadratic slowdown in recursive use is because recursive function

231+-- cannot be inlined and StreamD/StreamK conversions pile up and cannot be

232+-- eliminated by rewrite rules.

233+

234+-- | Like 'foldBreak' but fuses.

235+--

236+-- /Note:/ Unlike 'foldBreak', recursive application on the resulting stream

237+-- would lead to quadratic slowdown. If you need recursion with fusion (within

238+-- one iteration of recursion) use StreamD.foldBreak directly.

239+--

240+-- /Internal/

241+{-# INLINE foldBreak2 #-}

242+foldBreak2 :: Monad m => Fold m a b -> Stream m a -> m (b, Stream m a)

243+foldBreak2 fl strm = fmap f $ D.foldBreak fl $ toStreamD strm

244+

245+ where

246+

247+ f (b, str) = (b, fromStreamD str)

248+

249+-- | Fold resulting in either breaking the stream or continuation of the fold

250+-- Instead of supplying the input stream in one go we can run the fold multiple

251+-- times each time supplying the next segment of the input stream. If the fold

252+-- has not yet finished it returns a fold that can be run again otherwise it

253+-- returns the fold result and the residual stream.

254+--

255+-- /Internal/

256+{-# INLINE foldEither #-}

257+foldEither :: Monad m =>

258+ Fold m a b -> Stream m a -> m (Either (Fold m a b) (b, Stream m a))

226−foldBreak fl strm = fmap f $ K.foldBreak fl (toStreamK strm)

259+foldEither fl strm = fmap (fmap f) $ K.foldEither fl $ toStreamK strm

227

260

228 where

261 where

229

262

230 f (b, str) = (b, fromStreamK str)

263 f (b, str) = (b, fromStreamK str)

231

264

265+-- | Like 'foldEither' but fuses. However, recursive application on resulting

266+-- stream would lead to quadratic slowdown.

267+--

268+-- /Internal/

269+{-# INLINE foldEither2 #-}

270+foldEither2 :: Monad m =>

271+ Fold m a b -> Stream m a -> m (Either (Fold m a b) (b, Stream m a))

272+foldEither2 fl strm = fmap (fmap f) $ D.foldEither fl $ toStreamD strm

273+

274+ where

275+

276+ f (b, str) = (b, fromStreamD str)

277+

278+-- XXX Array folds can be implemented using this.

279+-- foldContainers? Specialized to foldArrays.

280+

281+-- | Generate streams from individual elements of a stream and fold the

282+-- concatenation of those streams using the supplied fold. Return the result of

283+-- the fold and residual stream.

284+--

285+-- For example, this can be used to efficiently fold an Array Word8 stream

286+-- using Word8 folds.

287+--

288+-- The outer stream forces CPS to allow scalable appends and the inner stream

289+-- forces direct style for stream fusion.

290+--

291+-- /Internal/

292+{-# INLINE foldConcat #-}

293+foldConcat :: Monad m =>

294+ Producer m a b -> Fold m b c -> Stream m a -> m (c, Stream m a)

295+foldConcat

296+ (Producer pstep pinject pextract)

297+ (Fold fstep begin done)

298+ stream = do

299+

300+ res <- begin

301+ case res of

302+ Fold.Partial fs -> go fs streamK

303+ Fold.Done fb -> return (fb, fromStreamK streamK)

304+

305+ where

306+

307+ streamK = toStreamK stream

308+

309+ go !acc m1 = do

310+ let stop = do

311+ r <- done acc

312+ return (r, fromStreamK K.nil)

313+ single a = do

314+ st <- pinject a

315+ res <- go1 SPEC acc st

316+ case res of

317+ Left fs -> do

318+ r <- done fs

319+ return (r, fromStreamK K.nil)

320+ Right (b, s) -> do

321+ x <- pextract s

322+ return (b, fromStreamK (K.fromPure x))

323+ yieldk a r = do

324+ st <- pinject a

325+ res <- go1 SPEC acc st

326+ case res of

327+ Left fs -> go fs r

328+ Right (b, s) -> do

329+ x <- pextract s

330+ return (b, fromStreamK (x `K.cons` r))

331+ in K.foldStream defState yieldk single stop m1

332+

333+ {-# INLINE go1 #-}

334+ go1 !_ !fs st = do

335+ r <- pstep st

336+ case r of

337+ D.Yield x s -> do

338+ res <- fstep fs x

339+ case res of

340+ Fold.Done b -> return $ Right (b, s)

341+ Fold.Partial fs1 -> go1 SPEC fs1 s

342+ D.Skip s -> go1 SPEC fs s

343+ D.Stop -> return $ Left fs

344+

232------------------------------------------------------------------------------

345------------------------------------------------------------------------------

233-- Transformation

346-- Transformation

234------------------------------------------------------------------------------

347------------------------------------------------------------------------------

core/src/Streamly/Internal/Data/Stream/StreamD/Type.hs+26−19

@@ -44,7 +44,7 @@ module Streamly.Internal.Data.Stream.StreamD.Type

44 , fold

45 , foldBreak

46 , foldContinue

47− , foldStream

47+ , foldEither

49 -- * Right Folds

50 , foldrT

@@ -282,13 +282,14 @@ fold fld strm = do

282 (b, _) <- foldBreak fld strm

283 return b

284

285−{-# INLINE_NORMAL foldBreak #-}

285+{-# INLINE_NORMAL foldEither #-}

286+foldEither :: Monad m =>

286−foldBreak :: Monad m => Fold m a b -> Stream m a -> m (b, Stream m a)

287+ Fold m a b -> Stream m a -> m (Either (Fold m a b) (b, Stream m a))

287−foldBreak (Fold fstep begin done) (UnStream step state) = do

288+foldEither (Fold fstep begin done) (UnStream step state) = do

288 res <- begin

289 res <- begin

289 case res of

290 case res of

290 FL.Partial fs -> go SPEC fs state

291 FL.Partial fs -> go SPEC fs state

291− FL.Done fb -> return $! (fb, Stream step state)

292+ FL.Done fb -> return $! Right (fb, Stream step state)

292

293

293 where

294 where

294

295

@@ -299,12 +300,28 @@ foldBreak (Fold fstep begin done) (UnStream step state) = do

299 Yield x s -> do

300 Yield x s -> do

300 res <- fstep fs x

301 res <- fstep fs x

301 case res of

302 case res of

302− FL.Done b -> return $! (b, Stream step s)

303+ FL.Done b -> return $! Right (b, Stream step s)

303 FL.Partial fs1 -> go SPEC fs1 s

304 FL.Partial fs1 -> go SPEC fs1 s

304 Skip s -> go SPEC fs s

305 Skip s -> go SPEC fs s

305− Stop -> do

306+ Stop -> return $! Left (Fold fstep (return $ FL.Partial fs) done)

306− b <- done fs

307+

308+{-# INLINE_NORMAL foldBreak #-}

309+foldBreak :: Monad m => Fold m a b -> Stream m a -> m (b, Stream m a)

310+foldBreak fld strm = do

311+ r <- foldEither fld strm

312+ case r of

313+ Right res -> return res

314+ Left (Fold _ initial extract) -> do

315+ res <- initial

316+ case res of

317+ FL.Done _ -> error "foldBreak: unreachable state"

318+ FL.Partial s -> do

319+ b <- extract s

320+ return (b, nil)

321+

322+ where

323+

307− return $! (b, Stream (\ _ _ -> return Stop) ())

324+ nil = Stream (\_ _ -> return Stop) ()

308

325

309-- | If the fold finishes before the stream, we can detect that the fold is

326-- | If the fold finishes before the stream, we can detect that the fold is

310-- done by checking if the initial action returns Done. But the remaining

327-- done by checking if the initial action returns Done. But the remaining

@@ -334,16 +351,6 @@ foldContinue (Fold fstep finitial fextract) (Stream sstep state) =

334 Skip s -> go SPEC fs s

351 Skip s -> go SPEC fs s

335 Stop -> return $ FL.Partial fs

352 Stop -> return $ FL.Partial fs

336

353

337−-- | Returns when either the fold or the stream finishes. If the Fold finishes

338−-- first we can check that using Fold.done. If the fold is not done then stream

339−-- would be nil.

340−--

341−-- /Unimplemented/

342−{-# INLINE_NORMAL foldStream #-}

343−foldStream :: -- Monad m =>

344− Fold m a b -> Stream m a -> m (Fold m a b, Stream m a)

345−foldStream = undefined

346−

347------------------------------------------------------------------------------

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

348-- Right Folds

355-- Right Folds

349------------------------------------------------------------------------------

356------------------------------------------------------------------------------

core/src/Streamly/Internal/Data/Stream/StreamK.hs+26−8

@@ -83,6 +83,7 @@ module Streamly.Internal.Data.Stream.StreamK

83 , foldlMx'

84 , fold

85 , foldBreak

86+ , foldEither

86 , parseBreak

87 , parseBreak

88 -- ** Specialized Folds

89 -- ** Specialized Folds

@@ -180,6 +181,7 @@ where

180import Control.Monad.Catch (MonadThrow, throwM)

181import Control.Monad.Catch (MonadThrow, throwM)

181import Control.Monad.Trans.Class (MonadTrans(lift))

182import Control.Monad.Trans.Class (MonadTrans(lift))

182import Control.Monad (void, join)

183import Control.Monad (void, join)

184+import Streamly.Internal.Data.Fold.Type (Fold(..))

183import Streamly.Internal.Data.SVar.Type (adaptState, defState)

185import Streamly.Internal.Data.SVar.Type (adaptState, defState)

184

186

185import qualified Streamly.Internal.Data.Fold.Type as FL

187import qualified Streamly.Internal.Data.Fold.Type as FL

@@ -333,30 +335,46 @@ fold (FL.Fold step begin done) m = do

333 FL.Done b1 -> return b1

335 FL.Done b1 -> return b1

334 in foldStream defState yieldk single stop m1

336 in foldStream defState yieldk single stop m1

335

337

336−{-# INLINE foldBreak #-}

338+{-# INLINE foldEither #-}

339+foldEither :: Monad m =>

337−foldBreak :: Monad m => FL.Fold m a b -> Stream m a -> m (b, Stream m a)

340+ Fold m a b -> Stream m a -> m (Either (Fold m a b) (b, Stream m a))

338−foldBreak (FL.Fold step begin done) m = do

341+foldEither (FL.Fold step begin done) m = do

339 res <- begin

342 res <- begin

340 case res of

343 case res of

341 FL.Partial fs -> go fs m

344 FL.Partial fs -> go fs m

342− FL.Done fb -> return (fb, m)

345+ FL.Done fb -> return $ Right (fb, m)

343

346

344 where

347 where

345

348

346 go !acc m1 =

349 go !acc m1 =

347− let stop = (, nil) <$> done acc

350+ let stop = return $ Left (Fold step (return $ FL.Partial acc) done)

348 single a =

351 single a =

349 step acc a

352 step acc a

350 >>= \case

353 >>= \case

351− FL.Partial s -> (, nil) <$> done s

354+ FL.Partial s ->

355+ return $ Left (Fold step (return $ FL.Partial s) done)

352− FL.Done b1 -> return (b1, nil)

356+ FL.Done b1 -> return $ Right (b1, nil)

353 yieldk a r =

357 yieldk a r =

354 step acc a

358 step acc a

355 >>= \case

359 >>= \case

356 FL.Partial s -> go s r

360 FL.Partial s -> go s r

357− FL.Done b1 -> return (b1, r)

361+ FL.Done b1 -> return $ Right (b1, r)

358 in foldStream defState yieldk single stop m1

362 in foldStream defState yieldk single stop m1

359

363

364+{-# INLINE foldBreak #-}

365+foldBreak :: Monad m => Fold m a b -> Stream m a -> m (b, Stream m a)

366+foldBreak fld strm = do

367+ r <- foldEither fld strm

368+ case r of

369+ Right res -> return res

370+ Left (Fold _ initial extract) -> do

371+ res <- initial

372+ case res of

373+ FL.Done _ -> error "foldBreak: unreachable state"

374+ FL.Partial s -> do

375+ b <- extract s

376+ return (b, nil)

377+

360-- | Like 'foldl'' but with a monadic step function.

378-- | Like 'foldl'' but with a monadic step function.

361{-# INLINE foldlM' #-}

379{-# INLINE foldlM' #-}

362foldlM' :: Monad m => (b -> a -> m b) -> m b -> Stream m a -> m b

380foldlM' :: Monad m => (b -> a -> m b) -> m b -> Stream m a -> m b