de5c0ca · streamly · git.iko.soy

git.iko.soy/streamly

use direct style scan, push mapMaybe from prelude to StreamK

commitde5c0ca22e9c4ff16ebb7cf495a89f2a60623e0f authorHarendra Kumar <harendra.kumar@gmail.com> dateMon 25 Jun 2018, 03:15 (8 years ago) parentc7b46ad

6 files changed+179−71

compare with master

unifiedsplit

Changelog.md+1 benchmark/Linear.hs+7 benchmark/LinearOps.hs+26−5 src/Streamly/Prelude.hs+59−63 src/Streamly/Streams/StreamD.hs+33−2 src/Streamly/Streams/StreamK.hs+53−1

Changelog.md+1

@@ -11,6 +11,7 @@

1111* Add `fromList` and `fromListM` to generate streams from lists, faster than

1212 `fromFoldable` and `fromFoldableM`

1313* Add `map` as a synonym of fmap

14+* Add `scanlM'`, the monadic version of scanl'

1415* Significant performance improvements

1516

1617## 0.3.0

benchmark/Linear.hs+7

@@ -15,10 +15,13 @@ import Gauge

1515

1616-- We need a monadic bind here to make sure that the function f does not get

1717-- completely optimized out by the compiler in some cases.

18+--

19+-- | Takes a fold method, and uses it with a default source.

1820{-# INLINE benchIO #-}

1921benchIO :: (IsStream t, NFData b) => String -> (t IO Int -> IO b) -> Benchmark

2022benchIO name f = bench name $ nfIO $ randomRIO (1,1000) >>= f . Ops.source

2123

24+-- | Takes a source, and uses it with a default drain/fold method.

2225{-# INLINE benchSrcIO #-}

2326benchSrcIO

2427 :: (t IO Int -> SerialT IO Int)

@@ -64,6 +67,10 @@ main = do

6467 , benchIO "map" Ops.map

6568 , benchIO "fmap" Ops.fmap

6669 , benchIO "mapM" (Ops.mapM serially)

70+ , benchIO "mapMaybe" Ops.mapMaybe

71+ , benchIO "mapMaybeM" Ops.mapMaybeM

72+ , bench "sequence" $ nfIO $ randomRIO (1,1000) >>= \n ->

73+ (Ops.sequence serially) (Ops.sourceUnfoldrMAction n)

6774 , benchIO "concat" Ops.concat

6875 ]

6976 , bgroup "filtering"

benchmark/LinearOps.hs+26−5

@@ -11,7 +11,7 @@ module LinearOps where

1111

1212import Prelude

1313 (Monad, Int, (+), ($), (.), return, fmap, even, (>), (<=),

14− subtract, undefined, Maybe(..))

14+ subtract, undefined, Maybe(..), odd)

1515

1616import qualified Streamly as S

1717import qualified Streamly.Prelude as S

@@ -61,10 +61,17 @@ last :: Monad m => Stream m Int -> m (Maybe Int)

6161{-# INLINE toNull #-}

6262toNull :: Monad m => (t m Int -> S.SerialT m Int) -> t m Int -> m ()

6363{-# INLINE mapM_ #-}

64−mapM_ :: Monad m => Stream m a -> m ()

64+mapM_ :: Monad m => Stream m Int -> m ()

6565{-# INLINE mapM #-}

6666mapM :: (S.IsStream t, S.MonadAsync m)

6767 => (t m Int -> S.SerialT m Int) -> t m Int -> m ()

68+{-# INLINE mapMaybe #-}

69+mapMaybe :: Monad m => Stream m Int -> m ()

70+{-# INLINE mapMaybeM #-}

71+mapMaybeM :: S.MonadAsync m => Stream m Int -> m ()

72+{-# INLINE sequence #-}

73+sequence :: (S.IsStream t, S.MonadAsync m)

74+ => (t m Int -> S.SerialT m Int) -> t m (m Int) -> m ()

6875{-# INLINE zipAsync #-}

6976zipAsync :: S.MonadAsync m => Stream m Int -> m ()

7077

@@ -123,14 +130,23 @@ sourceUnfoldrM n = S.unfoldrM step n

123130 then return Nothing

124131 else return (Just (cnt, cnt + 1))

125132

126−{-# INLINE runStream #-}

127−runStream :: Monad m => Stream m a -> m ()

128−runStream = S.runStream

133+{-# INLINE sourceUnfoldrMAction #-}

134+sourceUnfoldrMAction :: (S.IsStream t, S.MonadAsync m) => Int -> t m (m Int)

135+sourceUnfoldrMAction n = S.serially $ S.unfoldrM step n

136+ where

137+ step cnt =

138+ if cnt > n + value

139+ then return Nothing

140+ else return (Just (return cnt, cnt + 1))

129141

130142-------------------------------------------------------------------------------

131143-- Elimination

132144-------------------------------------------------------------------------------

133145

146+{-# INLINE runStream #-}

147+runStream :: Monad m => Stream m a -> m ()

148+runStream = S.runStream

149+

134150toNull t = runStream . t

135151mapM_ = S.mapM_ (\_ -> return ())

136152toList = S.toList

@@ -151,6 +167,11 @@ scan = transform . S.scanl' (+) 0

151167fmap = transform . Prelude.fmap (+1)

152168map = transform . S.map (+1)

153169mapM t = transform . t . S.mapM return

170+mapMaybe = transform . S.mapMaybe

171+ (\x -> if Prelude.odd x then Nothing else Just ())

172+mapMaybeM = transform . S.mapMaybeM

173+ (\x -> if Prelude.odd x then (return Nothing) else return $ Just ())

174+sequence t = transform . t . S.sequence

154175filterEven = transform . S.filter even

155176filterAllOut = transform . S.filter (> maxValue)

156177filterAllIn = transform . S.filter (<= maxValue)

src/Streamly/Prelude.hs+59−63

@@ -103,6 +103,7 @@ module Streamly.Prelude

103103

104104 -- * Scans

105105 , scanl'

106+ , scanlM'

106107 , scanx

107108

108109 -- * Mapping

@@ -160,6 +161,22 @@ import qualified Streamly.Streams.StreamK as K

160161import qualified Streamly.Streams.StreamD as D

161162import qualified Streamly.Streams.Serial as Serial

162163

164+------------------------------------------------------------------------------

165+-- Deconstruction

166+------------------------------------------------------------------------------

167+

168+-- | Decompose a stream into its head and tail. If the stream is empty, returns

169+-- 'Nothing'. If the stream is non-empty, returns @Just (a, ma)@, where @a@ is

170+-- the head of the stream and @ma@ its tail.

171+--

172+-- @since 0.1.0

173+uncons :: (IsStream t, Monad m) => SerialT m a -> m (Maybe (a, t m a))

174+uncons m =

175+ let stop = return Nothing

176+ single a = return (Just (a, K.nil))

177+ yieldk a r = return (Just (a, fromStream r))

178+ in (K.unStream (toStream m)) Nothing stop single yieldk

179+

163180------------------------------------------------------------------------------

164181-- Construction

165182------------------------------------------------------------------------------

@@ -379,40 +396,6 @@ foldr f = foldrM (\a b -> return (f a b))

379396foldrM :: Monad m => (a -> b -> m b) -> b -> SerialT m a -> m b

380397foldrM step acc m = D.foldrM step acc $ D.fromStreamK (toStream m)

381398

382−-- | Strict left scan with an extraction function. Like 'scanl'', but applies a

383−-- user supplied extraction function (the third argument) at each step. This is

384−-- designed to work with the @foldl@ library. The suffix @x@ is a mnemonic for

385−-- extraction.

386−--

387−-- @since 0.2.0

388−{-# INLINE scanx #-}

389−scanx :: IsStream t => (x -> a -> x) -> x -> (x -> b) -> t m a -> t m b

390−scanx step begin done m =

391− K.cons (done begin) $ fromStream $ go (toStream m) begin

392− where

393− go m1 !acc = K.Stream $ \_ stp sng yld ->

394− let single a = sng (done $ step acc a)

395− yieldk a r =

396− let s = step acc a

397− in yld (done s) (go r s)

398− in K.unStream m1 Nothing stp single yieldk

399−

400−-- |

401−-- @since 0.1.1

402−{-# DEPRECATED scan "Please use scanx instead." #-}

403−scan :: IsStream t => (x -> a -> x) -> x -> (x -> b) -> t m a -> t m b

404−scan = scanx

405−

406−-- | Strict left scan. Like 'foldl'', but returns the folded value at each

407−-- step, generating a stream of all intermediate fold results. The first

408−-- element of the stream is the user supplied initial value, and the last

409−-- element of the stream is the same as the result of 'foldl''.

410−--

411−-- @since 0.2.0

412−{-# INLINE scanl' #-}

413−scanl' :: IsStream t => (b -> a -> b) -> b -> t m a -> t m b

414−scanl' step begin m = scanx step begin id m

415−

416399-- | Strict left fold with an extraction function. Like the standard strict

417400-- left fold, but applies a user supplied extraction function (the third

418401-- argument) to the folded value at the end. This is designed to work with the

@@ -455,18 +438,6 @@ foldlM = foldxM

455438foldlM' :: Monad m => (b -> a -> m b) -> b -> SerialT m a -> m b

456439foldlM' step begin m = D.foldlM' step begin $ D.fromStreamK (toStream m)

457440

458−-- | Decompose a stream into its head and tail. If the stream is empty, returns

459−-- 'Nothing'. If the stream is non-empty, returns @Just (a, ma)@, where @a@ is

460−-- the head of the stream and @ma@ its tail.

461−--

462−-- @since 0.1.0

463−uncons :: (IsStream t, Monad m) => SerialT m a -> m (Maybe (a, t m a))

464−uncons m =

465− let stop = return Nothing

466− single a = return (Just (a, K.nil))

467− yieldk a r = return (Just (a, fromStream r))

468− in (K.unStream (toStream m)) Nothing stop single yieldk

469−

470441-- | Write a stream of Strings to an IO Handle.

471442--

472443-- @since 0.1.0

@@ -711,6 +682,44 @@ maximum m = go Nothing (toStream m)

711682 Nothing -> Just a

712683 Just e -> Just $ max a e

713684

685+------------------------------------------------------------------------------

686+-- Scans

687+------------------------------------------------------------------------------

688+

689+-- | Strict left scan with an extraction function. Like 'scanl'', but applies a

690+-- user supplied extraction function (the third argument) at each step. This is

691+-- designed to work with the @foldl@ library. The suffix @x@ is a mnemonic for

692+-- extraction.

693+--

694+-- @since 0.2.0

695+{-# INLINE scanx #-}

696+scanx :: IsStream t => (x -> a -> x) -> x -> (x -> b) -> t m a -> t m b

697+scanx = K.scanx

698+

699+-- |

700+-- @since 0.1.1

701+{-# DEPRECATED scan "Please use scanx instead." #-}

702+scan :: IsStream t => (x -> a -> x) -> x -> (x -> b) -> t m a -> t m b

703+scan = scanx

704+

705+-- | Like 'scanl'' but with a monadic step function.

706+--

707+-- @since 0.4.0

708+{-# INLINE scanlM' #-}

709+scanlM' :: (IsStream t, Monad m) => (b -> a -> m b) -> b -> t m a -> t m b

710+scanlM' step begin m = fromStream $ D.toStreamK $

711+ D.scanlM' step begin $ D.fromStreamK (toStream m)

712+

713+-- | Strict left scan. Like 'foldl'', but returns the folded value at each

714+-- step, generating a stream of all intermediate fold results. The first

715+-- element of the stream is the user supplied initial value, and the last

716+-- element of the stream is the same as the result of 'foldl''.

717+--

718+-- @since 0.2.0

719+{-# INLINE scanl' #-}

720+scanl' :: (IsStream t, Monad m) => (b -> a -> b) -> b -> t m a -> t m b

721+scanl' step = scanlM' (\a b -> return (step a b))

722+

714723------------------------------------------------------------------------------

715724-- Transformation

716725------------------------------------------------------------------------------

@@ -743,17 +752,8 @@ mapMSerial = Serial.mapM

743752--

744753-- @since 0.3.0

745754{-# INLINE mapMaybe #-}

746−mapMaybe :: (IsStream t) => (a -> Maybe b) -> t m a -> t m b

747−mapMaybe f m = go (toStream m)

748− where

749− go m1 = fromStream $ K.Stream $ \_ stp sng yld ->

750− let single a = case f a of

751− Just b -> sng b

752− Nothing -> stp

753− yieldk a r = case f a of

754− Just b -> yld b (toStream $ go r)

755− Nothing -> (K.unStream r) Nothing stp single yieldk

756− in (K.unStream m1) Nothing stp single yieldk

755+mapMaybe :: IsStream t => (a -> Maybe b) -> t m a -> t m b

756+mapMaybe = K.mapMaybe

757757

758758-- | Like 'mapMaybe' but maps a monadic function.

759759--

@@ -788,10 +788,6 @@ mapM_ f m = D.mapM_ f $ D.fromStreamK (toStream m)

788788-- /Concurrent (do not use with 'parallely' on infinite streams)/

789789--

790790-- @since 0.1.0

791+{-# INLINE sequence #-}

791792sequence :: (IsStream t, MonadAsync m) => t m (m a) -> t m a

792−sequence m = go (toStream m)

793− where

794− go m1 = fromStream $ K.Stream $ \svr stp sng yld ->

795− let single ma = ma >>= sng

796− yieldk ma r = K.unStream (toStream $ ma |: go r) svr stp sng yld

797− in (K.unStream m1) Nothing stp single yieldk

793+sequence = K.sequence

src/Streamly/Streams/StreamD.hs+33−2

@@ -37,9 +37,8 @@ module Streamly.Streams.StreamD

3737 , Stream (..)

3838

3939 -- * Construction

40− -- | Direct style stream does not support construction by adding more

41− -- elements to an existing stream.

4240 , nil

41+ , cons

4342

4443 -- * Deconstruction

4544 , uncons

@@ -70,6 +69,9 @@ module Streamly.Streams.StreamD

7069 , toList

7170 , last

7271

72+ -- * Scans

73+ , scanlM'

74+

7375 -- * Mapping

7476 , map

7577 , mapM

@@ -145,6 +147,17 @@ fromStreamK m = Stream step m

145147nil :: Monad m => Stream m a

146148nil = Stream (const $ return Stop) ()

147149

150+-- | Note that the cons operation is not scalable as it has O(n^2) complexity.

151+cons :: Monad m => a -> Stream m a -> Stream m a

152+cons x (Stream step state) = Stream step1 Nothing

153+ where

154+ step1 Nothing = return $ Yield x (Just state)

155+ step1 (Just st) = do

156+ r <- step st

157+ case r of

158+ Yield a s -> return $ Yield a (Just s)

159+ Stop -> return Stop

160+

148161-- | Create a singleton 'Stream' from a pure value.

149162{-# INLINE_NORMAL yield #-}

150163yield :: Monad m => a -> Stream m a

@@ -320,3 +333,21 @@ last (Stream step state) = go0 SPEC state

320333 case r of

321334 Yield y s -> go1 SPEC y s

322335 Stop -> return (Just x)

336+

337+{-# INLINE_NORMAL postscanlM' #-}

338+postscanlM' :: Monad m => (b -> a -> m b) -> b -> Stream m a -> Stream m b

339+postscanlM' fstep begin (Stream step state) =

340+ begin `seq` Stream step' (state, begin)

341+ where

342+ {-# INLINE_LATE step' #-}

343+ step' (st, acc) = acc `seq` do

344+ r <- step st

345+ case r of

346+ Yield x s -> do

347+ y <- fstep acc x

348+ y `seq` return (Yield y (s, y))

349+ Stop -> return Stop

350+

351+{-# INLINE scanlM' #-}

352+scanlM' :: Monad m => (b -> a -> m b) -> b -> Stream m a -> Stream m b

353+scanlM' fstep begin s = begin `seq` (begin `cons` postscanlM' fstep begin s)

src/Streamly/Streams/StreamK.hs+53−1

@@ -74,9 +74,15 @@ module Streamly.Streams.StreamK

7474 , toList

7575 , last

7676

77+ -- * Scans

78+ , scanx

79+ , scanl'

80+

7781 -- * Mapping

7882 , map

7983 , mapM

84+ , mapMaybe

85+ , sequence

8086

8187 -- * Semigroup Style Composition

8288 , serial

@@ -96,7 +102,7 @@ import Control.Monad (void)

96102import Control.Monad.Reader.Class (MonadReader(..))

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

98104import Data.Semigroup (Semigroup(..))

99−import Prelude hiding (foldl, foldr, last, map, mapM, mapM_, repeat)

105+import Prelude hiding (foldl, foldr, last, map, mapM, mapM_, repeat, sequence)

100106import qualified Prelude

101107

102108import Streamly.SVar

@@ -450,6 +456,26 @@ toList = foldr (:) []

450456last :: (IsStream t, Monad m) => t m a -> m (Maybe a)

451457last = foldx (\_ y -> Just y) Nothing id

452458

459+-------------------------------------------------------------------------------

460+-- Scans

461+-------------------------------------------------------------------------------

462+

463+{-# INLINE scanx #-}

464+scanx :: IsStream t => (x -> a -> x) -> x -> (x -> b) -> t m a -> t m b

465+scanx step begin done m =

466+ cons (done begin) $ fromStream $ go (toStream m) begin

467+ where

468+ go m1 !acc = Stream $ \_ stp sng yld ->

469+ let single a = sng (done $ step acc a)

470+ yieldk a r =

471+ let s = step acc a

472+ in yld (done s) (go r s)

473+ in unStream m1 Nothing stp single yieldk

474+

475+{-# INLINE scanl' #-}

476+scanl' :: IsStream t => (b -> a -> b) -> b -> t m a -> t m b

477+scanl' step begin m = scanx step begin id m

478+

453479-------------------------------------------------------------------------------

454480-- Generation

455481-------------------------------------------------------------------------------

@@ -532,6 +558,8 @@ map f m = fromStream $ Stream $ \_ stp sng yld ->

532558instance Monad m => Functor (Stream m) where

533559 fmap = map

534560

561+-- Be careful when modifying this, this uses a consM (|:) deliberately to allow

562+-- other stream types to overload it.

535563{-# INLINE mapM #-}

536564mapM :: (IsStream t, MonadAsync m) => (a -> m b) -> t m a -> t m b

537565mapM f m = go (toStream m)

@@ -541,6 +569,30 @@ mapM f m = go (toStream m)

541569 yieldk a r = unStream (toStream (f a |: (go r))) svr stp sng yld

542570 in (unStream m1) Nothing stp single yieldk

543571

572+{-# INLINE mapMaybe #-}

573+mapMaybe :: IsStream t => (a -> Maybe b) -> t m a -> t m b

574+mapMaybe f m = go (toStream m)

575+ where

576+ go m1 = fromStream $ Stream $ \_ stp sng yld ->

577+ let single a = case f a of

578+ Just b -> sng b

579+ Nothing -> stp

580+ yieldk a r = case f a of

581+ Just b -> yld b (toStream $ go r)

582+ Nothing -> (unStream r) Nothing stp single yieldk

583+ in unStream m1 Nothing stp single yieldk

584+

585+-- Be careful when modifying this, this uses a consM (|:) deliberately to allow

586+-- other stream types to overload it.

587+{-# INLINE sequence #-}

588+sequence :: (IsStream t, MonadAsync m) => t m (m a) -> t m a

589+sequence m = go (toStream m)

590+ where

591+ go m1 = fromStream $ Stream $ \svr stp sng yld ->

592+ let single ma = ma >>= sng

593+ yieldk ma r = unStream (toStream $ ma |: go r) svr stp sng yld

594+ in (unStream m1) Nothing stp single yieldk

595+

544596-------------------------------------------------------------------------------

545597-- Bind utility

546598-------------------------------------------------------------------------------