Stream.bound and Stream.unfold
[darcs-mirrors/vector.git] / Data / Vector / Stream.hs
1 {-# LANGUAGE ExistentialQuantification, BangPatterns, CPP #-}
2
3 #include "phases.h"
4
5 module Data.Vector.Stream (
6 Step(..), Stream(..),
7
8 bound, unfold,
9 empty, singleton, replicate, (++),
10 map, filter, zipWith,
11 foldr, foldl, foldl',
12 mapM_, foldM
13 ) where
14
15 import Prelude hiding ( replicate, (++), map, filter, zipWith,
16 foldr, foldl,
17 mapM_ )
18
19 data Step s a = Yield a s
20 | Skip s
21 | Done
22
23 data Stream a = forall s. Stream (s -> Step s a) s Int
24
25 bound :: Stream a -> Int
26 {-# INLINE bound #-}
27 bound (Stream _ _ n) = n
28
29 unfold :: (s -> Maybe (a, s)) -> s -> Int -> Stream a
30 {-# INLINE_STREAM unfold #-}
31 unfold f s n = Stream step s n
32 where
33 {-# INLINE step #-}
34 step s = case f s of
35 Just (x, s') -> Yield x s'
36 Nothing -> Done
37
38 empty :: Stream a
39 {-# INLINE_STREAM empty #-}
40 empty = Stream (const Done) () 0
41
42 singleton :: a -> Stream a
43 {-# INLINE_STREAM singleton #-}
44 singleton x = Stream step True 1
45 where
46 {-# INLINE step #-}
47 step True = Yield x False
48 step False = Done
49
50 replicate :: Int -> a -> Stream a
51 {-# INLINE_STREAM replicate #-}
52 replicate n x = Stream step n (max n 0)
53 where
54 {-# INLINE step #-}
55 step i | i > 0 = Yield x (i-1)
56 | otherwise = Done
57
58 infixr ++
59 (++) :: Stream a -> Stream a -> Stream a
60 {-# INLINE_STREAM (++) #-}
61 Stream stepa sa na ++ Stream stepb sb nb = Stream step (Left sa) (na + nb)
62 where
63 step (Left sa) = case stepa sa of
64 Yield x sa' -> Yield x (Left sa')
65 Skip sa' -> Skip (Left sa')
66 Done -> Skip (Right sb)
67 step (Right sb) = case stepb sb of
68 Yield x sb' -> Yield x (Right sb')
69 Skip sb' -> Skip (Right sb')
70 Done -> Done
71
72 map :: (a -> b) -> Stream a -> Stream b
73 {-# INLINE_STREAM map #-}
74 map f (Stream step s n) = Stream step' s n
75 where
76 {-# INLINE step' #-}
77 step' s = case step s of
78 Yield x s' -> Yield (f x) s'
79 Skip s' -> Skip s'
80 Done -> Done
81
82 filter :: (a -> Bool) -> Stream a -> Stream a
83 {-# INLINE_STREAM filter #-}
84 filter f (Stream step s n) = Stream step' s n
85 where
86 {-# INLINE step' #-}
87 step' s = case step s of
88 Yield x s' | f x -> Yield x s'
89 | otherwise -> Skip s'
90 Skip s' -> Skip s'
91 Done -> Done
92
93 zipWith :: (a -> b -> c) -> Stream a -> Stream b -> Stream c
94 {-# INLINE_STREAM zipWith #-}
95 zipWith f (Stream stepa sa na) (Stream stepb sb nb)
96 = Stream step (sa, sb, Nothing) (min na nb)
97 where
98 {-# INLINE step #-}
99 step (sa, sb, Nothing) = case stepa sa of
100 Yield x sa' -> Skip (sa', sb, Just x)
101 Skip sa' -> Skip (sa', sb, Nothing)
102 Done -> Done
103
104 step (sa, sb, Just x) = case stepb sb of
105 Yield y sb' -> Yield (f x y) (sa, sb', Nothing)
106 Skip sb' -> Skip (sa, sb', Just x)
107 Done -> Done
108
109 foldl :: (a -> b -> b) -> b -> Stream a -> b
110 {-# INLINE_STREAM foldl #-}
111 foldl f z (Stream step s _) = foldl_go z s
112 where
113 foldl_go z s = case step s of
114 Yield x s' -> foldl_go (f x z) s'
115 Skip s' -> foldl_go z s'
116 Done -> z
117
118 foldl' :: (a -> b -> b) -> b -> Stream a -> b
119 {-# INLINE_STREAM foldl' #-}
120 foldl' f z (Stream step s _) = foldl_go z s
121 where
122 foldl_go !z s = case step s of
123 Yield x s' -> foldl_go (f x z) s'
124 Skip s' -> foldl_go z s'
125 Done -> z
126
127 foldr :: (a -> b -> b) -> b -> Stream a -> b
128 {-# INLINE_STREAM foldr #-}
129 foldr f z (Stream step s _) = foldr_go s
130 where
131 foldr_go s = case step s of
132 Yield x s' -> f x (foldr_go s')
133 Skip s' -> foldr_go s'
134 Done -> z
135
136 mapM_ :: Monad m => (a -> m ()) -> Stream a -> m ()
137 {-# INLINE_STREAM mapM_ #-}
138 mapM_ m (Stream step s _) = mapM_go s
139 where
140 mapM_go s = case step s of
141 Yield x s' -> do { m x; mapM_go s' }
142 Skip s' -> mapM_go s'
143 Done -> return ()
144
145 foldM :: Monad m => (a -> b -> m a) -> a -> Stream b -> m a
146 {-# INLINE_STREAM foldM #-}
147 foldM m z (Stream step s _) = foldM_go z s
148 where
149 foldM_go z s = case step s of
150 Yield x s' -> do { z' <- m z x; foldM_go z' s' }
151 Skip s' -> foldM_go z s'
152 Done -> return z
153