@@ -0,0 +1,231 @@
1+{-# LANGUAGE BangPatterns #-}
2+{-# LANGUAGE CPP #-}
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 #-}
12+
13+#include "inline.h"
14+
15+-- |
16+-- Module : Streamly.Streams.StreamD
17+-- Copyright : (c) 2018 Harendra Kumar
18+--
19+-- License : BSD3
20+-- Maintainer : harendra.kumar@gmail.com
21+-- Stability : experimental
22+-- Portability : GHC
23+--
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.
28+--
29+-- import qualified Streamly.Streams.StreamD as D
30+--
31+module Streamly.Streams.StreamD
32+ (
33+ -- * The stream type
34+ Step (..)
35+ , Stream (..)
36+
37+ -- * Elimination
38+ , runStream
39+ , uncons
40+
41+ -- * Generation
42+ , nil
43+ , yield
44+ , yieldM
45+ , unfoldr
46+ , unfoldrM
47+ , repeat
48+ , enumFromStepN
49+ , fromList
50+
51+ -- * Transformation
52+ , mapM
53+ , mapM_
54+
55+ -- * Conversion
56+ , toStreamK
57+ , fromStreamK
58+ )
59+where
60+
61+import qualified Streamly.Streams.StreamK as K
62+import Prelude hiding (mapM, mapM_, repeat)
63+
64+------------------------------------------------------------------------------
65+-- The direct style stream type
66+------------------------------------------------------------------------------
67+
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
70+-- stream.
71+data Step s a = Yield a s | Stop
72+
73+instance Functor (Step s) where
74+ {-# INLINE fmap #-}
75+ fmap f (Yield x s) = Yield (f x) s
76+ fmap _ Stop = Stop
77+
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
81+
82+-------------------------------------------------------------------------------
83+-- IsStream Stream
84+-------------------------------------------------------------------------------
85+
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
90+ where
91+ go st = K.Stream $ \_ stp _ yld -> do
92+ r <- step st
93+ case r of
94+ Yield x s -> yld x (go s)
95+ Stop -> stp
96+
97+{-# INLINE_NORMAL fromStreamK #-}
98+fromStreamK :: Monad m => K.Stream m a -> Stream m a
99+fromStreamK m = Stream step m
100+ where
101+ step m1 =
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
106+
107+{-
108+instance K.IsStream Stream where
109+ toStream = toStreamK
110+ fromStream = fromStreamK
111+
112+ consM = undefined
113+ (|:) = undefined
114+-}
115+
116+------------------------------------------------------------------------------
117+-- Construction
118+------------------------------------------------------------------------------
119+
120+-- | An empty 'Stream'.
121+{-# INLINE_NORMAL nil #-}
122+nil :: Monad m => Stream m a
123+nil = Stream (const $ return Stop) ()
124+
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
129+ where
130+ {-# INLINE_LATE step #-}
131+ step True = Yield x False
132+ step False = Stop
133+
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
138+ where
139+ {-# INLINE_LATE step #-}
140+ step True = m >>= \x -> return $ Yield x False
141+ step False = return Stop
142+
143+------------------------------------------------------------------------------
144+-- Generation by unfold
145+------------------------------------------------------------------------------
146+
147+{-# INLINE_NORMAL unfoldrM #-}
148+unfoldrM :: Monad m => (s -> m (Maybe (a, s))) -> s -> Stream m a
149+unfoldrM next state = Stream step state
150+ where
151+ {-# INLINE_LATE step #-}
152+ step st = do
153+ r <- next st
154+ return $ case r of
155+ Just (x, s) -> Yield x s
156+ Nothing -> Stop
157+
158+{-# INLINE_NORMAL unfoldr #-}
159+unfoldr :: Monad m => (s -> Maybe (a, s)) -> s -> Stream m a
160+unfoldr f = unfoldrM (return . f)
161+
162+------------------------------------------------------------------------------
163+-- Generation
164+------------------------------------------------------------------------------
165+
166+repeat :: Monad m => a -> Stream m a
167+repeat x = Stream (const $ return $ Yield x ()) ()
168+
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)
173+ where
174+ {-# INLINE_LATE step #-}
175+ step (x, i) | i > 0 = return $ Yield x (x + stride, i - 1)
176+ | otherwise = return $ Stop
177+
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
182+ where
183+ {-# INLINE_LATE step #-}
184+ step (x:xs) = return $ Yield x xs
185+ step [] = return Stop
186+
187+------------------------------------------------------------------------------
188+-- Transformation
189+------------------------------------------------------------------------------
190+
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
195+ where
196+ {-# INLINE_LATE step' #-}
197+ step' st = do
198+ r <- step st
199+ case r of
200+ Yield x s -> f x >>= \a -> return $ Yield a s
201+ Stop -> return Stop
202+
203+-------------------------------------------------------------------------------
204+-- Elimination
205+-------------------------------------------------------------------------------
206+
207+{-# INLINE_NORMAL uncons #-}
208+uncons :: Monad m => Stream m a -> m (Maybe (a, Stream m a))
209+uncons (Stream step state) = go state
210+ where
211+ go st = do
212+ r <- step st
213+ return $ case r of
214+ Yield x s -> Just (x, (Stream step s))
215+ Stop -> Nothing
216+
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
221+ where
222+ go st = do
223+ r <- step st
224+ case r of
225+ Yield _ s -> go s
226+ Stop -> return ()
227+
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