Ticket #888: Stream.hs

File Stream.hs, 3.9 KB (added by simonpj, 9 years ago)

Stream example

Line 
1module StreamA
2where
3
4data Step s a = Done | Yield a s | Skip s
5data Stream a = forall s. Stream (s -> Step s a) s
6
7mkStream :: (s -> Step s a) -> s -> Stream a
8mkStream = Stream
9
10next :: Stream a -> Step (Stream a) a
11{-# INLINE next #-}
12next (Stream step s) = case step s of
13                         Done       -> Done
14                         Skip s'    -> Skip (Stream step s')
15                         Yield x s' -> Yield x (Stream step s')
16
17unstream :: Stream a -> [a]
18{-# INLINE [1] unstream #-}
19unstream s = go s
20  where
21    go s = case next s of
22             Done       -> []
23             Yield x s' -> x : go s'
24             Skip s'    -> go s'
25
26stream :: [a] -> Stream a
27{-# INLINE [1] stream #-}
28stream xs = mkStream step xs
29  where
30    step []     = Done
31    step (x:xs) = Yield x xs
32
33{-# RULES
34
35"stream/unstream" forall s.
36  stream (unstream s) = s
37
38 #-}
39
40emptyS :: Stream a
41emptyS = mkStream (const Done) ()
42
43mapS :: (a -> b) -> Stream a -> Stream b
44{-# INLINE [1] mapS #-}
45mapS f s = mkStream step s
46  where
47    step s = case next s of
48               Done       -> Done
49               Skip s'    -> Skip s'
50               Yield x s' -> Yield (f x) s'
51
52filterS :: (a -> Bool) -> Stream a -> Stream a
53{-# INLINE [1] filterS #-}
54filterS p s = mkStream step s
55  where
56    step s = case next s of
57               Done                   -> Done
58               Skip s'                -> Skip s'
59               Yield x s' | p x       -> Yield x s'
60                          | otherwise -> Skip s'
61
62appendS :: Stream a -> Stream a -> Stream a
63{-# INLINE [1] appendS #-}
64appendS s1 s2 = mkStream step (Left s1)
65  where
66    step (Left s1)  = case next s1 of
67                        Done        -> Skip (Right s2)
68                        Skip s1'    -> Skip (Left s1')
69                        Yield x s1' -> Yield x (Left s1')
70    step (Right s2) = case next s2 of
71                        Done        -> Done
72                        Skip s2'    -> Skip (Right s2')
73                        Yield x s2' -> Yield x (Right s2')
74
75concatMapS :: (a -> Stream b) -> Stream a -> Stream b
76{-# INLINE [1] concatMapS #-}
77concatMapS f s = Stream step (s, emptyS)
78  where
79    step (s, t) =
80      case next t of
81        Done       -> case next s of
82                        Done       -> Done
83                        Skip s'    -> Skip (s', emptyS)
84                        Yield x s' -> Skip (s', f x)
85        Skip t'    -> Skip (s, t')
86        Yield x t' -> Yield x (s, t')
87
88foldlS :: (a -> b -> a) -> a -> Stream b -> a
89{-# INLINE [1] foldlS #-}
90foldlS f z s = go z s
91  where
92    go z s = case next s of
93               Done       -> z
94               Skip s'    -> go z s'
95               Yield x s' -> go (f z x) s'
96
97foldrS :: (a -> b -> b) -> b -> Stream a -> b
98{-# INLINE [1] foldrS #-}
99foldrS f z s = go s
100  where
101    go s = case next s of
102             Done       -> z
103             Skip s'    -> go s'
104             Yield x s' -> f x (go s')
105
106enumS :: Int -> Int -> Stream Int
107{-# INLINE [1] enumS #-}
108enumS m n = mkStream step m
109  where
110    step m | m > n     = Done
111           | otherwise = Yield m (m+1)
112               
113mapL :: (a -> b) -> [a] -> [b]
114{-# INLINE mapL #-}
115mapL f = unstream . mapS f . stream
116
117filterL :: (a -> Bool) -> [a] -> [a]
118{-# INLINE filterL #-}
119filterL f = unstream . filterS f . stream
120
121appendL :: [a] -> [a] -> [a]
122{-# INLINE appendL #-}
123appendL xs ys = unstream (appendS (stream xs) (stream ys))
124
125foldlL :: (a -> b -> a) -> a -> [b] -> a
126{-# INLINE foldlL #-}
127foldlL f z = foldlS f z . stream
128
129foldrL :: (a -> b -> b) -> b -> [a] -> b
130{-# INLINE foldrL #-}
131foldrL f z = foldrS f z . stream
132
133enumL :: Int -> Int -> [Int]
134{-# INLINE enumL #-}
135enumL m = unstream . enumS m
136
137concatMapL :: (a -> [b]) -> [a] -> [b]
138{-# INLINE concatMapL #-}
139concatMapL f = unstream . concatMapS (stream . f) . stream
140