git.iko.soy/streamly
on 54acd41 help sign in
streamly/src/Streamly/Streams/StreamDK.hs
165 lines 4.4 kB history raw
1{-# LANGUAGE BangPatterns #-}
2{-# LANGUAGE CPP #-}
3{-# LANGUAGE ExistentialQuantification #-}
4{-# LANGUAGE FlexibleContexts #-}
5{-# LANGUAGE PatternSynonyms #-}
6{-# LANGUAGE RankNTypes #-}
7{-# LANGUAGE RecordWildCards #-}
8-- {-# LANGUAGE ScopedTypeVariables #-}
10#include "inline.hs"
12-- |
13-- Module : Streamly.Streams.StreamDK
14-- Copyright : (c) 2019 Composewell Technologies
15-- License : BSD3
16-- Maintainer : harendra.kumar@gmail.com
17-- Stability : experimental
18-- Portability : GHC
19--
21module Streamly.Streams.StreamDK
22 (
23 -- * Stream Type
25 Stream
26 , Step (..)
28 -- * Construction
29 , nil
30 , cons
31 , consM
32 , unfoldr
33 , unfoldrM
34 , replicateM
36 -- * Folding
37 , uncons
38 , foldrS
40 -- * Specific Folds
41 , drain
42 )
43where
45import Streamly.Streams.StreamDK.Type (Stream(..), Step(..))
47-------------------------------------------------------------------------------
48-- Construction
49-------------------------------------------------------------------------------
51nil :: Monad m => Stream m a
52nil = Stream $ return Stop
54{-# INLINE_NORMAL cons #-}
55cons :: Monad m => a -> Stream m a -> Stream m a
56cons x xs = Stream $ return $ Yield x xs
58consM :: Monad m => m a -> Stream m a -> Stream m a
59consM eff xs = Stream $ eff >>= \x -> return $ Yield x xs
61unfoldrM :: Monad m => (s -> m (Maybe (a, s))) -> s -> Stream m a
62unfoldrM next state = Stream (step' state)
63 where
64 step' st = do
65 r <- next st
66 return $ case r of
67 Just (x, s) -> Yield x (Stream (step' s))
68 Nothing -> Stop
69{-
70unfoldrM next s0 = buildM $ \yld stp ->
71 let go s = do
72 r <- next s
73 case r of
74 Just (a, b) -> yld a (go b)
75 Nothing -> stp
76 in go s0
77-}
79{-# INLINE unfoldr #-}
80unfoldr :: Monad m => (b -> Maybe (a, b)) -> b -> Stream m a
81unfoldr next s0 = build $ \yld stp ->
82 let go s =
83 case next s of
84 Just (a, b) -> yld a (go b)
85 Nothing -> stp
86 in go s0
88replicateM :: Monad m => Int -> a -> Stream m a
89replicateM n x = Stream (step n)
90 where
91 step i = return $
92 if i <= 0
93 then Stop
94 else Yield x (Stream (step (i - 1)))
96-------------------------------------------------------------------------------
97-- Folding
98-------------------------------------------------------------------------------
100uncons :: Monad m => Stream m a -> m (Maybe (a, Stream m a))
101uncons (Stream step) = do
102 r <- step
103 return $ case r of
104 Yield x xs -> Just (x, xs)
105 Stop -> Nothing
107-- | Lazy right associative fold to a stream.
108{-# INLINE_NORMAL foldrS #-}
109foldrS :: Monad m
110 => (a -> Stream m b -> Stream m b)
111 -> Stream m b
112 -> Stream m a
113 -> Stream m b
114foldrS f streamb = go
115 where
116 go (Stream stepa) = Stream $ do
117 r <- stepa
118 case r of
119 Yield x xs -> let Stream step = f x (go xs) in step
120 Stop -> let Stream step = streamb in step
122{-# INLINE_LATE foldrM #-}
123foldrM :: Monad m => (a -> m b -> m b) -> m b -> Stream m a -> m b
124foldrM fstep acc ys = go ys
125 where
126 go (Stream step) = do
127 r <- step
128 case r of
129 Yield x xs -> fstep x (go xs)
130 Stop -> acc
132{-# INLINE_NORMAL build #-}
133build :: Monad m
134 => forall a. (forall b. (a -> b -> b) -> b -> b) -> Stream m a
135build g = g cons nil
137{-# RULES
138"foldrM/build" forall k z (g :: forall b. (a -> b -> b) -> b -> b).
139 foldrM k z (build g) = g k z #-}
141{-
142-- To fuse foldrM with unfoldrM we need the type m1 to be polymorphic such that
143-- it is either Monad m or Stream m. So that we can use cons/nil as well as
144-- monadic construction function as its arguments.
145--
146{-# INLINE_NORMAL buildM #-}
147buildM :: Monad m
148 => forall a. (forall b. (a -> m1 b -> m1 b) -> m1 b -> m1 b) -> Stream m a
149buildM g = g cons nil
150-}
152-------------------------------------------------------------------------------
153-- Specific folds
154-------------------------------------------------------------------------------
156{-# INLINE drain #-}
157drain :: Monad m => Stream m a -> m ()
158drain = foldrM (\_ xs -> xs) (return ())
159{-
160drain (Stream step) = do
161 r <- step
162 case r of
163 Yield _ next -> drain next
164 Stop -> return ()
165 -}