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