@@ -0,0 +1,231 @@
1+{-# LANGUAGE BangPatterns #-}
3+{-# LANGUAGE ConstraintKinds #-}
4+{-# LANGUAGE ExistentialQuantification #-}
5+{-# LANGUAGE FlexibleContexts #-}
6+{-# LANGUAGE FlexibleInstances #-}
7+{-# LANGUAGE LambdaCase #-}
8+{-# LANGUAGE MagicHash #-}
9+{-# LANGUAGE MultiParamTypeClasses #-}
10+{-# LANGUAGE RankNTypes #-}
11+{-# LANGUAGE UnboxedTuples #-}
16+-- Module : Streamly.Streams.StreamD
17+-- Copyright : (c) 2018 Harendra Kumar
20+-- Maintainer : harendra.kumar@gmail.com
21+-- Stability : experimental
24+-- Direct style re-implementation of CPS style stream in StreamK module. The
25+-- symbol or suffix 'D' in this module denotes the "Direct" style. GHC is able
26+-- to INLINE and fuse direct style better, providing better performance than
27+-- CPS implementation.
29+-- import qualified Streamly.Streams.StreamD as D
31+module Streamly.Streams.StreamD
61+import qualified Streamly.Streams.StreamK as K
62+import Prelude hiding (mapM, mapM_, repeat)
64+------------------------------------------------------------------------------
65+-- The direct style stream type
66+------------------------------------------------------------------------------
68+-- | A stream is a succession of 'Step's. A 'Yield' produces a single value and
69+-- the next state of the stream. Stop indicates there are no more values in the
71+data Step s a = Yield a s | Stop
73+instance Functor (Step s) where
75+ fmap f (Yield x s) = Yield (f x) s
78+-- | A stream consists of a step function that generates the next step given a
79+-- current state, and the current state.
80+data Stream m a = forall s. Stream (s -> m (Step s a)) s
82+-------------------------------------------------------------------------------
84+-------------------------------------------------------------------------------
86+-- Convert a direct stream to and from CPS encoded stream
87+{-# INLINE_NORMAL toStreamK #-}
88+toStreamK :: Monad m => Stream m a -> K.Stream m a
89+toStreamK (Stream step state) = go state
91+ go st = K.Stream $ \_ stp _ yld -> do
94+ Yield x s -> yld x (go s)
97+{-# INLINE_NORMAL fromStreamK #-}
98+fromStreamK :: Monad m => K.Stream m a -> Stream m a
99+fromStreamK m = Stream step m
102+ let stop = return Stop
103+ single a = return $ Yield a K.nil
104+ yieldk a r = return $ Yield a r
105+ in K.unStream m1 Nothing stop single yieldk
108+instance K.IsStream Stream where
109+ toStream = toStreamK
110+ fromStream = fromStreamK
116+------------------------------------------------------------------------------
118+------------------------------------------------------------------------------
120+-- | An empty 'Stream'.
121+{-# INLINE_NORMAL nil #-}
122+nil :: Monad m => Stream m a
123+nil = Stream (const $ return Stop) ()
125+-- | Create a singleton 'Stream' from a pure value.
126+{-# INLINE_NORMAL yield #-}
127+yield :: Monad m => a -> Stream m a
128+yield x = Stream (return . step) True
130+ {-# INLINE_LATE step #-}
131+ step True = Yield x False
134+-- | Create a singleton 'Stream' from a monadic action.
135+{-# INLINE_NORMAL yieldM #-}
136+yieldM :: Monad m => m a -> Stream m a
137+yieldM m = Stream step True
139+ {-# INLINE_LATE step #-}
140+ step True = m >>= \x -> return $ Yield x False
141+ step False = return Stop
143+------------------------------------------------------------------------------
144+-- Generation by unfold
145+------------------------------------------------------------------------------
147+{-# INLINE_NORMAL unfoldrM #-}
148+unfoldrM :: Monad m => (s -> m (Maybe (a, s))) -> s -> Stream m a
149+unfoldrM next state = Stream step state
151+ {-# INLINE_LATE step #-}
155+ Just (x, s) -> Yield x s
158+{-# INLINE_NORMAL unfoldr #-}
159+unfoldr :: Monad m => (s -> Maybe (a, s)) -> s -> Stream m a
160+unfoldr f = unfoldrM (return . f)
162+------------------------------------------------------------------------------
164+------------------------------------------------------------------------------
166+repeat :: Monad m => a -> Stream m a
167+repeat x = Stream (const $ return $ Yield x ()) ()
169+{-# INLINE_NORMAL enumFromStepN #-}
170+enumFromStepN :: (Num a, Monad m) => a -> a -> Int -> Stream m a
171+enumFromStepN from stride n =
172+ from `seq` stride `seq` n `seq` Stream step (from, n)
174+ {-# INLINE_LATE step #-}
175+ step (x, i) | i > 0 = return $ Yield x (x + stride, i - 1)
176+ | otherwise = return $ Stop
178+-- | Convert a list to a 'Stream'
179+{-# INLINE_NORMAL fromList #-}
180+fromList :: Monad m => [a] -> Stream m a
181+fromList zs = Stream step zs
183+ {-# INLINE_LATE step #-}
184+ step (x:xs) = return $ Yield x xs
185+ step [] = return Stop
187+------------------------------------------------------------------------------
189+------------------------------------------------------------------------------
191+-- | Map a monadic function over a 'Stream'
192+{-# INLINE_NORMAL mapM #-}
193+mapM :: Monad m => (a -> m b) -> Stream m a -> Stream m b
194+mapM f (Stream step state) = Stream step' state
196+ {-# INLINE_LATE step' #-}
200+ Yield x s -> f x >>= \a -> return $ Yield a s
203+-------------------------------------------------------------------------------
205+-------------------------------------------------------------------------------
207+{-# INLINE_NORMAL uncons #-}
208+uncons :: Monad m => Stream m a -> m (Maybe (a, Stream m a))
209+uncons (Stream step state) = go state
214+ Yield x s -> Just (x, (Stream step s))
217+-- | Run a streaming composition, discard the results.
218+{-# INLINE_NORMAL runStream #-}
219+runStream :: Monad m => Stream m a -> m ()
220+runStream (Stream step state) = go state
228+-- | Execute a monadic action for each element of the 'Stream'
229+{-# INLINE_NORMAL mapM_ #-}
230+mapM_ :: Monad m => (a -> m b) -> Stream m a -> m ()
231+mapM_ m = runStream . mapM m