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