git.iko.soy/streamly
on 21d779c help sign in
streamly/src/Streamly/Prelude.hs
430 lines 15 kB history raw
1{-# LANGUAGE BangPatterns #-}
2{-# LANGUAGE FlexibleContexts #-}
3{-# LANGUAGE FlexibleInstances #-}
4{-# LANGUAGE GeneralizedNewtypeDeriving#-}
5{-# LANGUAGE MultiParamTypeClasses #-}
6{-# LANGUAGE RankNTypes #-}
7{-# LANGUAGE StandaloneDeriving #-}
8{-# LANGUAGE UndecidableInstances #-} -- XXX
10-- |
11-- Module : Streamly.Prelude
12-- Copyright : (c) 2017 Harendra Kumar
13--
14-- License : BSD3
15-- Maintainer : harendra.kumar@gmail.com
16-- Stability : experimental
17-- Portability : GHC
18--
19--
20module Streamly.Prelude
21 (
22 -- * Construction
23 cons
24 , nil
25 , unfoldr
26 , unfoldrM
27 , each
28 , fromHandle
30 -- * Elimination
31 , foldr
32 , foldrM
33 , foldl
34 , foldlM
35 , uncons
37 -- * Elimination Special Folds
38 , toList
39 , toHandle
40 , all
41 , any
42 , sum
43 , product
44 , head
45 , last
46 , length
47 , elem
48 , notElem
49 , maximum
50 , minimum
52 -- * Filtering
53 , filter
54 , take
55 , takeWhile
56 , drop
57 , dropWhile
59 -- * Transformation
60 , mapM
61 , mapM_
62 , sequence
64 -- * Zipping
65 , zipWith
66 , zipWithM
67 , zipAsyncWith
68 , zipAsyncWithM
69 )
70where
72import Control.Monad (liftM)
73import Control.Monad.IO.Class (MonadIO(..))
74import Data.Semigroup (Semigroup(..))
75import Prelude hiding (filter, drop, dropWhile, take,
76 takeWhile, zipWith, foldr, foldl,
77 mapM, mapM_, sequence, all, any,
78 sum, product, elem, notElem,
79 maximum, minimum, head, last,
80 length)
81import qualified Prelude as Prelude
82import qualified System.IO as IO
84import Streamly.Core
85import Streamly.Streams
87------------------------------------------------------------------------------
88-- Construction
89------------------------------------------------------------------------------
91-- | Build a Stream by unfolding pure steps starting from a seed.
92unfoldr :: Streaming t => (b -> Maybe (a, b)) -> b -> t m a
93unfoldr step = fromStream . go
94 where
95 go s = Stream $ \_ stp yld -> do
96 case step s of
97 Nothing -> stp
98 Just (a, b) -> yld a (Just (go b))
100-- | Build a Stream by unfolding monadic steps starting from a seed.
101unfoldrM :: (Streaming t, Monad m) => (b -> m (Maybe (a, b))) -> b -> t m a
102unfoldrM step = fromStream . go
103 where
104 go s = Stream $ \_ stp yld -> do
105 mayb <- step s
106 case mayb of
107 Nothing -> stp
108 Just (a, b) -> yld a (Just (go b))
110-- XXX need eachInterleaved, eachAsync, eachParallel
111-- | Same as @foldWith (<>)@ but more efficient.
112{-# INLINE each #-}
113each :: (Foldable f, Streaming t) => f a -> t m a
114each xs = Prelude.foldr cons nil xs
116-- | Read lines from an IO Handle into a stream of Strings.
117fromHandle :: (MonadIO m, Streaming t) => IO.Handle -> t m String
118fromHandle h = fromStream $ go
119 where
120 go = Stream $ \_ stp yld -> do
121 eof <- liftIO $ IO.hIsEOF h
122 if eof
123 then stp
124 else do
125 str <- liftIO $ IO.hGetLine h
126 yld str (Just go)
128------------------------------------------------------------------------------
129-- Elimination
130------------------------------------------------------------------------------
132-- Parallel variants of folds?
134-- | Right fold.
135foldr :: (Monad m, Streaming t) => (a -> b -> b) -> b -> t m a -> m b
136foldr step acc m = go (toStream m)
137 where
138 go m1 =
139 let stop = return acc
140 yield a Nothing = return (step a acc)
141 yield a (Just x) = go x >>= \b -> return (step a b)
142 in (runStream m1) Nothing stop yield
144-- | Right fold with a monadic step function. See 'toList' for an example use.
145{-# INLINE foldrM #-}
146foldrM :: Streaming t => (a -> m b -> m b) -> m b -> t m a -> m b
147foldrM step acc m = go (toStream m)
148 where
149 go m1 =
150 let stop = acc
151 yield a Nothing = step a acc
152 yield a (Just x) = step a (go x)
153 in (runStream m1) Nothing stop yield
155-- | Strict left fold. This is typed to work with the foldl package. To use
156-- directly pass 'id' as the third argument.
157foldl :: (Monad m, Streaming t)
158 => (x -> a -> x) -> x -> (x -> b) -> t m a -> m b
159foldl step begin done m = go begin (toStream m)
160 where
161 go !acc m1 =
162 let stop = return (done acc)
163 yield a Nothing = return (done (step acc a))
164 yield a (Just x) = go (step acc a) x
165 in (runStream m1) Nothing stop yield
167-- | Strict left fold, with monadic step function. This is typed to work
168-- with the foldl package. To use directly pass 'id' as the third argument.
169foldlM :: (Monad m, Streaming t)
170 => (x -> a -> m x) -> m x -> (x -> m b) -> t m a -> m b
171foldlM step begin done m = go begin (toStream m)
172 where
173 go !acc m1 =
174 let stop = acc >>= done
175 yield a Nothing = acc >>= \b -> step b a >>= done
176 yield a (Just x) = acc >>= \b -> go (step b a) x
177 in (runStream m1) Nothing stop yield
179-- | Decompose a stream into its head and tail. If the stream is empty, returns
180-- 'Nothing'. If the stream is non-empty, returns 'Just (a, ma)', where 'a' is
181-- the head of the stream and 'ma' its tail.
182uncons :: (Streaming t, Monad m) => t m a -> m (Maybe (a, t m a))
183uncons m =
184 let stop = return Nothing
185 yield a Nothing = return (Just (a, nil))
186 yield a (Just x) = return (Just (a, (fromStream x)))
187 in (runStream (toStream m)) Nothing stop yield
189-- | Write a stream of Strings to an IO Handle.
190toHandle :: (Streaming t, MonadIO m) => IO.Handle -> t m String -> m ()
191toHandle h m = go (toStream m)
192 where
193 go m1 =
194 let stop = return ()
195 yield a Nothing = liftIO (IO.hPutStrLn h a)
196 yield a (Just x) = liftIO (IO.hPutStrLn h a) >> go x
197 in (runStream m1) Nothing stop yield
199------------------------------------------------------------------------------
200-- Special folds
201------------------------------------------------------------------------------
203-- | Convert a stream into a list in the underlying monad.
204{-# INLINABLE toList #-}
205toList :: (Monad m, Streaming t) => t m a -> m [a]
206toList = foldrM (\a xs -> liftM (a :) xs) (return [])
208-- | Take first 'n' elements from the stream and discard the rest.
209take :: Streaming t => Int -> t m a -> t m a
210take n m = fromStream $ go n (toStream m)
211 where
212 go n1 m1 = Stream $ \ctx stp yld -> do
213 let yield a Nothing = yld a Nothing
214 yield a (Just x) = yld a (Just (go (n1 - 1) x))
215 if (n1 <= 0)
216 then stp
217 else (runStream m1) ctx stp yield
219-- XXX This is not as efficient as it could be. We need a short circuiting at
220-- a lower level. Compare with simple-conduit, filtering there cuts down time
221-- due to short circuting whereas the time spent remains the same here.
223-- | Include only those elements that pass a predicate.
224{-# INLINE filter #-}
225filter :: (Streaming t, Monad (t m)) => (a -> Bool) -> t m a -> t m a
226filter p m = m >>= \x -> if p x then return x else nil
228-- | End the stream as soon as the predicate fails on an element.
229takeWhile :: Streaming t => (a -> Bool) -> t m a -> t m a
230takeWhile p m = fromStream $ go (toStream m)
231 where
232 go m1 = Stream $ \ctx stp yld -> do
233 let yield a Nothing | p a = yld a Nothing
234 | otherwise = stp
235 yield a (Just x) | p a = yld a (Just (go x))
236 | otherwise = stp
237 in (runStream m1) ctx stp yield
239-- | Discard first 'n' elements from the stream and take the rest.
240drop :: Streaming t => Int -> t m a -> t m a
241drop n m = fromStream $ go n (toStream m)
242 where
243 go n1 m1 = Stream $ \ctx stp yld -> do
244 let yield _ Nothing = stp
245 yield _ (Just x) = (runStream $ go (n1 - 1) x) ctx stp yld
246 if (n1 <= 0)
247 then (runStream m1) ctx stp yld
248 else (runStream m1) ctx stp yield
250-- | Drop elements in the stream as long as the predicate succeeds and then
251-- take the rest of the stream.
252dropWhile :: Streaming t => (a -> Bool) -> t m a -> t m a
253dropWhile p m = fromStream $ go (toStream m)
254 where
255 go m1 = Stream $ \ctx stp yld -> do
256 let yield a Nothing | p a = stp
257 | otherwise = yld a Nothing
258 yield a (Just x) | p a = (runStream (go x)) ctx stp yield
259 | otherwise = yld a (Just x)
260 in (runStream m1) ctx stp yield
262-- | Determine whether all elements of a stream satisfy a predicate.
263all :: (Streaming t, Monad m) => (a -> Bool) -> t m a -> m Bool
264all p m = go (toStream m)
265 where
266 go m1 =
267 let yield a Nothing | p a = return True
268 | otherwise = return False
269 yield a (Just x) | p a = go x
270 | otherwise = return False
271 in (runStream m1) Nothing (return True) yield
273-- | Determine whether any of the elements of a stream satisfy a predicate.
274any :: (Streaming t, Monad m) => (a -> Bool) -> t m a -> m Bool
275any p m = go (toStream m)
276 where
277 go m1 =
278 let yield a Nothing | p a = return True
279 | otherwise = return False
280 yield a (Just x) | p a = return True
281 | otherwise = go x
282 in (runStream m1) Nothing (return False) yield
284-- | Determine the sum of all elements of a stream of numbers
285sum :: (Streaming t, Monad m, Num a) => t m a -> m a
286sum = foldl (+) 0 id
288-- | Determine the product of all elements of a stream of numbers
289product :: (Streaming t, Monad m, Num a) => t m a -> m a
290product = foldl (*) 0 id
292-- | Extract the first element of the stream, if any.
293head :: (Streaming t, Monad m) => t m a -> m (Maybe a)
294head m =
295 let stop = return Nothing
296 yield a _ = return (Just a)
297 in (runStream (toStream m)) Nothing stop yield
299-- | Extract the last element of the stream, if any.
300last :: (Streaming t, Monad m) => t m a -> m (Maybe a)
301last m = go (toStream m)
302 where
303 go m1 =
304 let stop = return Nothing
305 yield a Nothing = return (Just a)
306 yield _ (Just x) = go x
307 in (runStream m1) Nothing stop yield
309-- | Determine whether an element is present in the stream.
310elem :: (Streaming t, Monad m, Eq a) => a -> t m a -> m Bool
311elem e m = go (toStream m)
312 where
313 go m1 =
314 let stop = return False
315 yield a Nothing = return (a == e)
316 yield a (Just x) = if (a == e) then return True else go x
317 in (runStream m1) Nothing stop yield
319-- | Determine whether an element is not present in the stream.
320notElem :: (Streaming t, Monad m, Eq a) => a -> t m a -> m Bool
321notElem e m = go (toStream m)
322 where
323 go m1 =
324 let stop = return True
325 yield a Nothing = return (a /= e)
326 yield a (Just x) = if (a == e) then return False else go x
327 in (runStream m1) Nothing stop yield
329-- | Determine the length of the stream.
330length :: (Streaming t, Monad m) => t m a -> m Int
331length = foldl (\n _ -> n + 1) 0 id
333-- | Determine the minimum element in a stream.
334minimum :: (Streaming t, Monad m, Ord a) => t m a -> m (Maybe a)
335minimum m = go Nothing (toStream m)
336 where
337 go r m1 =
338 let stop = return r
339 yield a Nothing = return $ min_ a r
340 yield a (Just x) = go (min_ a r) x
341 in (runStream m1) Nothing stop yield
343 min_ a r = case r of
344 Nothing -> Just a
345 Just e -> Just $ min a e
347-- | Determine the maximum element in a stream.
348maximum :: (Streaming t, Monad m, Ord a) => t m a -> m (Maybe a)
349maximum m = go Nothing (toStream m)
350 where
351 go r m1 =
352 let stop = return r
353 yield a Nothing = return $ max_ a r
354 yield a (Just x) = go (max_ a r) x
355 in (runStream m1) Nothing stop yield
357 max_ a r = case r of
358 Nothing -> Just a
359 Just e -> Just $ max a e
361------------------------------------------------------------------------------
362-- Transformation
363------------------------------------------------------------------------------
365-- XXX Parallel variants of these? mapMWith et al. sequenceWith.
367-- | Replace each element of the stream with the result of a monadic action
368-- applied on the element.
369mapM :: (Streaming t, Monad m) => (a -> m b) -> t m a -> t m b
370mapM f m = fromStream $ go (toStream m)
371 where
372 go m1 = Stream $ \_ stp yld -> do
373 let stop = stp
374 yield a Nothing = f a >>= \b -> yld b Nothing
375 yield a (Just x) = f a >>= \b -> yld b (Just (go x))
376 in (runStream m1) Nothing stop yield
378-- | Apply a monadic action to each element of the stream and discard the
379-- output of the action.
380mapM_ :: (Streaming t, Monad m) => (a -> m b) -> t m a -> m ()
381mapM_ f m = go (toStream m)
382 where
383 go m1 =
384 let stop = return ()
385 yield a Nothing = f a >> return ()
386 yield a (Just x) = f a >> go x
387 in (runStream m1) Nothing stop yield
389-- | Reduce a stream of monadic actions to a stream of the output of those
390-- actions.
391sequence :: (Streaming t, Monad m) => t m (m a) -> t m a
392sequence m = fromStream $ go (toStream m)
393 where
394 go m1 = Stream $ \_ stp yld -> do
395 let stop = stp
396 yield a Nothing = a >>= \b -> yld b Nothing
397 yield a (Just x) = a >>= \b -> yld b (Just (go x))
398 in (runStream m1) Nothing stop yield
400------------------------------------------------------------------------------
401-- Serially Zipping Streams
402------------------------------------------------------------------------------
404-- | Zip two streams serially using a monadic zipping function.
405zipWithM :: Streaming t => (a -> b -> t m c) -> t m a -> t m b -> t m c
406zipWithM f m1 m2 = fromStream $ go (toStream m1) (toStream m2)
407 where
408 go mx my = Stream $ \_ stp yld -> do
409 let merge a ra =
410 let yield2 b Nothing = (runStream (g a b)) Nothing stp yld
411 yield2 b (Just rb) =
412 (runStream ((g a b) <> (go ra rb))) Nothing stp yld
413 in (runStream my) Nothing stp yield2
414 let yield1 a Nothing = merge a snil
415 yield1 a (Just ra) = merge a ra
416 (runStream mx) Nothing stp yield1
417 g a b = toStream $ f a b
419------------------------------------------------------------------------------
420-- Parallely Zipping Streams
421------------------------------------------------------------------------------
423-- | Zip two streams asyncly (i.e. both the elements being zipped are generated
424-- concurrently) using a monadic zipping function.
425zipAsyncWithM :: (Streaming t, MonadAsync m)
426 => (a -> b -> t m c) -> t m a -> t m b -> t m c
427zipAsyncWithM f m1 m2 = fromStream $ Stream $ \_ stp yld -> do
428 ma <- async m1
429 mb <- async m2
430 (runStream (toStream (zipWithM f ma mb))) Nothing stp yld