9ad7fca · streamly · git.iko.soy

git.iko.soy/streamly

on master help sign in

Add inner monad transformations

9ad7fcaHarendra Kumar7 years ago

benchmark/Linear.hs+4

compare with master

unifiedsplit

@@ -296,6 +296,10 @@ main =

296 , benchIOSink "tee" (Ops.transformTeeMapM serially 4)

297 , benchIOSink "zip" (Ops.transformZipMapM serially 4)

298 ]

299+ , bgroup "transformer"

300+ [ benchIOSrc serially "evalState" Ops.evalStateT

301+ , benchIOSrc serially "withState" Ops.withState

302+ ]

299 , bgroup "transformation"

303 , bgroup "transformation"

300 [ benchIOSink "scanl" (Ops.scan 1)

304 [ benchIOSink "scanl" (Ops.scan 1)

301 , benchIOSink "scanl1'" (Ops.scanl1' 1)

305 , benchIOSink "scanl1'" (Ops.scanl1' 1)

src/Streamly/Benchmark/Prelude.hs+23

@@ -25,6 +25,7 @@ module Streamly.Benchmark.Prelude where

25import Control.DeepSeq (NFData)

26import Control.Monad (when)

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

28+import Control.Monad.State.Strict (StateT, get, put)

28import Data.Functor.Identity (Identity, runIdentity)

29import Data.Functor.Identity (Identity, runIdentity)

29import Data.Maybe (fromJust)

30import Data.Maybe (fromJust)

30import GHC.Generics (Generic)

31import GHC.Generics (Generic)

@@ -167,6 +168,19 @@ sourceUnfoldrM n = S.unfoldrM step n

167 then return Nothing

168 then return Nothing

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

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

169

170

171+{-# INLINE sourceUnfoldrState #-}

172+sourceUnfoldrState :: (S.IsStream t, S.MonadAsync m)

173+ => Int -> t (StateT Int m) Int

174+sourceUnfoldrState n = S.unfoldrM step n

175+ where

176+ step cnt =

177+ if cnt > n + value

178+ then return Nothing

179+ else do

180+ s <- get

181+ put (s + 1)

182+ return (Just (s, cnt + 1))

183+

170{-# INLINE sourceUnfoldrMN #-}

184{-# INLINE sourceUnfoldrMN #-}

171sourceUnfoldrMN :: (S.IsStream t, S.MonadAsync m) => Int -> Int -> t m Int

185sourceUnfoldrMN :: (S.IsStream t, S.MonadAsync m) => Int -> Int -> t m Int

172sourceUnfoldrMN upto start = S.unfoldrM step start

186sourceUnfoldrMN upto start = S.unfoldrM step start

@@ -208,6 +222,15 @@ runStream = S.drain

208{-# INLINE toList #-}

222{-# INLINE toList #-}

209toList :: Monad m => Stream m Int -> m [Int]

223toList :: Monad m => Stream m Int -> m [Int]

210

224

225+{-# INLINE evalStateT #-}

226+evalStateT :: S.MonadAsync m => Int -> Stream m Int

227+evalStateT n = Internal.evalStateT 0 (sourceUnfoldrState n)

228+

229+{-# INLINE withState #-}

230+withState :: S.MonadAsync m => Int -> Stream m Int

231+withState n =

232+ Internal.evalStateT (0 :: Int) (Internal.liftInner (sourceUnfoldrM n))

233+

211{-# INLINE head #-}

234{-# INLINE head #-}

212{-# INLINE last #-}

235{-# INLINE last #-}

213{-# INLINE maximum #-}

236{-# INLINE maximum #-}

src/Streamly/Internal/Prelude.hs+94−11

@@ -133,8 +133,6 @@ module Streamly.Internal.Prelude

133

134 -- * Transformation

135 , transform

136− , hoist

137− , generally

138

136

139 -- ** Mapping

137 -- ** Mapping

140 , Serial.map

138 , Serial.map

@@ -376,6 +374,17 @@ module Streamly.Internal.Prelude

376 , finally

374 , finally

377 , handle

375 , handle

378

376

377+ -- * Generalize Inner Monad

378+ , hoist

379+ , generally

380+

381+ -- * Transform Inner Monad

382+ , liftInner

383+ , runReaderT

384+ , evalStateT

385+ , usingStateT

386+ , runStateT

387+

379 -- * Diagnostics

388 -- * Diagnostics

380 , inspectMode

389 , inspectMode

381

390

@@ -399,6 +408,8 @@ import Control.Exception (Exception)

399import Control.Monad (void)

408import Control.Monad (void)

400import Control.Monad.Catch (MonadCatch)

409import Control.Monad.Catch (MonadCatch)

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

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

411+import Control.Monad.Reader (ReaderT)

412+import Control.Monad.State.Strict (StateT)

402import Control.Monad.Trans (MonadTrans(..))

413import Control.Monad.Trans (MonadTrans(..))

403import Data.Functor.Identity (Identity (..))

414import Data.Functor.Identity (Identity (..))

404import Data.Heap (Entry(..))

415import Data.Heap (Entry(..))

@@ -1471,15 +1482,6 @@ toPureRev = foldl' (flip K.cons) K.nil

1471transform :: (IsStream t, Monad m) => Pipe m a b -> t m a -> t m b

1482transform :: (IsStream t, Monad m) => Pipe m a b -> t m a -> t m b

1472transform pipe xs = fromStreamD $ D.transform pipe (toStreamD xs)

1483transform pipe xs = fromStreamD $ D.transform pipe (toStreamD xs)

1473

1484

1474−{-# INLINE hoist #-}

1475−hoist :: (Monad m, Monad n)

1476− => (forall x. m x -> n x) -> SerialT m a -> SerialT n a

1477−hoist f xs = fromStreamS $ S.hoist f (toStreamS xs)

1478−

1479−{-# INLINE generally #-}

1480−generally :: (IsStream t, Monad m) => t Identity a -> t m a

1481−generally xs = fromStreamS $ S.hoist (return . runIdentity) (toStreamS xs)

1482−

1483------------------------------------------------------------------------------

1485------------------------------------------------------------------------------

1484-- Transformation by Folding (Scans)

1486-- Transformation by Folding (Scans)

1485------------------------------------------------------------------------------

1487------------------------------------------------------------------------------

@@ -3696,3 +3698,84 @@ handle :: (IsStream t, MonadCatch m, Exception e)

3696 => (e -> t m a) -> t m a -> t m a

3698 => (e -> t m a) -> t m a -> t m a

3697handle handler xs =

3699handle handler xs =

3698 D.fromStreamD $ D.handle (\e -> D.toStreamD $ handler e) $ D.toStreamD xs

3700 D.fromStreamD $ D.handle (\e -> D.toStreamD $ handler e) $ D.toStreamD xs

3701+

3702+------------------------------------------------------------------------------

3703+-- Generalize the underlying monad

3704+------------------------------------------------------------------------------

3705+

3706+-- | Transform the inner monad of a stream using a natural transformation.

3707+--

3708+-- / Internal/

3709+--

3710+{-# INLINE hoist #-}

3711+hoist :: (Monad m, Monad n)

3712+ => (forall x. m x -> n x) -> SerialT m a -> SerialT n a

3713+hoist f xs = fromStreamS $ S.hoist f (toStreamS xs)

3714+

3715+-- | Generalize the inner monad of the stream from 'Identity' to any monad.

3716+--

3717+-- / Internal/

3718+--

3719+{-# INLINE generally #-}

3720+generally :: (IsStream t, Monad m) => t Identity a -> t m a

3721+generally xs = fromStreamS $ S.hoist (return . runIdentity) (toStreamS xs)

3722+

3723+------------------------------------------------------------------------------

3724+-- Add and remove a monad transformer

3725+------------------------------------------------------------------------------

3726+

3727+-- | Lift the inner monad of a stream using a monad transformer.

3728+--

3729+-- / Internal/

3730+--

3731+{-# INLINE liftInner #-}

3732+liftInner :: (Monad m, IsStream t, MonadTrans tr, Monad (tr m))

3733+ => t m a -> t (tr m) a

3734+liftInner xs = fromStreamD $ D.liftInner (toStreamD xs)

3735+

3736+-- | Evaluate the inner monad of a stream as 'ReaderT'.

3737+--

3738+-- / Internal/

3739+--

3740+{-# INLINE runReaderT #-}

3741+runReaderT :: (IsStream t, Monad m) => s -> t (ReaderT s m) a -> t m a

3742+runReaderT s xs = fromStreamD $ D.runReaderT s (toStreamD xs)

3743+

3744+-- | Evaluate the inner monad of a stream as 'StateT'.

3745+--

3746+-- This is supported only for 'SerialT' as concurrent state updation may not be

3747+-- safe.

3748+--

3749+-- / Internal/

3750+--

3751+{-# INLINE evalStateT #-}

3752+evalStateT :: Monad m => s -> SerialT (StateT s m) a -> SerialT m a

3753+evalStateT s xs = fromStreamD $ D.evalStateT s (toStreamD xs)

3754+

3755+-- | Run a stateful (StateT) stream transformation using a given state.

3756+--

3757+-- This is supported only for 'SerialT' as concurrent state updation may not be

3758+-- safe.

3759+--

3760+-- / Internal/

3761+--

3762+{-# INLINE usingStateT #-}

3763+usingStateT

3764+ :: Monad m

3765+ => s

3766+ -> (SerialT (StateT s m) a -> SerialT (StateT s m) a)

3767+ -> SerialT m a

3768+ -> SerialT m a

3769+usingStateT s f xs = evalStateT s $ f $ liftInner xs

3770+

3771+-- | Evaluate the inner monad of a stream as 'StateT' and emit the resulting

3772+-- state and value pair after each step.

3773+--

3774+-- This is supported only for 'SerialT' as concurrent state updation may not be

3775+-- safe.

3776+--

3777+-- / Internal/

3778+--

3779+{-# INLINE runStateT #-}

3780+runStateT :: Monad m => s -> SerialT (StateT s m) a -> SerialT m (s, a)

3781+runStateT s xs = fromStreamD $ D.runStateT s (toStreamD xs)

src/Streamly/Streams/StreamD.hs+55

@@ -188,9 +188,15 @@ module Streamly.Streams.StreamD

188 , toListRev

189 , toStreamK

190 , toStreamD

191+

191 , hoist

192 , hoist

192 , generally

193 , generally

193

194

195+ , liftInner

196+ , runReaderT

197+ , evalStateT

198+ , runStateT

199+

194 -- * Transformation

200 -- * Transformation

195 , transform

201 , transform

196

202

@@ -290,6 +296,8 @@ import Control.Exception (Exception, SomeException)

290import Control.Monad (void)

296import Control.Monad (void)

291import Control.Monad.Catch (MonadCatch)

297import Control.Monad.Catch (MonadCatch)

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

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

299+import Control.Monad.Reader (ReaderT)

300+import Control.Monad.State.Strict (StateT)

293import Control.Monad.Trans (MonadTrans(lift))

301import Control.Monad.Trans (MonadTrans(lift))

294import Data.Bits (shiftR, shiftL, (.|.), (.&.))

302import Data.Bits (shiftR, shiftL, (.|.), (.&.))

295import Data.Functor.Identity (Identity(..))

303import Data.Functor.Identity (Identity(..))

@@ -312,6 +320,8 @@ import Prelude

312 reverse)

320 reverse)

313

321

314import qualified Control.Monad.Catch as MC

322import qualified Control.Monad.Catch as MC

323+import qualified Control.Monad.Reader as Reader

324+import qualified Control.Monad.State.Strict as State

315

325

316import Streamly.Internal.Memory.Array.Types (Array(..))

326import Streamly.Internal.Memory.Array.Types (Array(..))

317import Streamly.Internal.Data.Fold.Types (Fold(..))

327import Streamly.Internal.Data.Fold.Types (Fold(..))

@@ -650,6 +660,51 @@ hoist f (Stream step state) = (Stream step' state)

650generally :: Monad m => Stream Identity a -> Stream m a

660generally :: Monad m => Stream Identity a -> Stream m a

651generally = hoist (return . runIdentity)

661generally = hoist (return . runIdentity)

652

662

663+{-# INLINE_NORMAL liftInner #-}

664+liftInner :: (Monad m, MonadTrans t, Monad (t m))

665+ => Stream m a -> Stream (t m) a

666+liftInner (Stream step state) = Stream step' state

667+ where

668+ step' gst st = do

669+ r <- lift $ step (adaptState gst) st

670+ return $ case r of

671+ Yield x s -> Yield x s

672+ Skip s -> Skip s

673+ Stop -> Stop

674+

675+{-# INLINE_NORMAL runReaderT #-}

676+runReaderT :: Monad m => s -> Stream (ReaderT s m) a -> Stream m a

677+runReaderT sval (Stream step state) = Stream step' state

678+ where

679+ step' gst st = do

680+ r <- Reader.runReaderT (step (adaptState gst) st) sval

681+ return $ case r of

682+ Yield x s -> Yield x s

683+ Skip s -> Skip s

684+ Stop -> Stop

685+

686+{-# INLINE_NORMAL evalStateT #-}

687+evalStateT :: Monad m => s -> Stream (StateT s m) a -> Stream m a

688+evalStateT sval (Stream step state) = Stream step' (state, sval)

689+ where

690+ step' gst (st, sv) = do

691+ (r, sv') <- State.runStateT (step (adaptState gst) st) sv

692+ return $ case r of

693+ Yield x s -> Yield x (s, sv')

694+ Skip s -> Skip (s, sv')

695+ Stop -> Stop

696+

697+{-# INLINE_NORMAL runStateT #-}

698+runStateT :: Monad m => s -> Stream (StateT s m) a -> Stream m (s, a)

699+runStateT sval (Stream step state) = Stream step' (state, sval)

700+ where

701+ step' gst (st, sv) = do

702+ (r, sv') <- State.runStateT (step (adaptState gst) st) sv

703+ return $ case r of

704+ Yield x s -> Yield (sv', x) (s, sv')

705+ Skip s -> Skip (s, sv')

706+ Stop -> Stop

707+

653------------------------------------------------------------------------------

708------------------------------------------------------------------------------

654-- Elimination by Folds

709-- Elimination by Folds

655------------------------------------------------------------------------------

710------------------------------------------------------------------------------