1 {-# LANGUAGE ExistentialQuantification, BangPatterns #-}
3 module Data.Vector.Stream (
4 Step(..), Stream(..),
6 empty, singleton, replicate,
7 map, filter, zipWith,
8 foldr, foldl, foldl',
9 mapM_, foldM
10 ) where
12 import Prelude hiding ( replicate, map, filter, zipWith,
13 foldr, foldl,
14 mapM_ )
16 data Step s a = Yield a s
17 | Skip s
18 | Done
20 data Stream a = forall s. Stream (s -> Step s a) s Int
22 empty :: Stream a
23 {-# INLINE_STREAM empty #-}
24 empty = Stream (const Done) () 0
26 singleton :: a -> Stream a
27 {-# INLINE_STREAM singleton #-}
28 singleton x = Stream step True 1
29 where
30 {-# INLINE step #-}
31 step True = Yield x False
32 step False = Done
34 replicate :: Int -> a -> Stream a
35 {-# INLINE_STREAM replicate #-}
36 replicate n x = Stream step n (max n 0)
37 where
38 {-# INLINE step #-}
39 step i | i > 0 = Yield x (i-1)
40 | otherwise = Done
42 map :: (a -> b) -> Stream a -> Stream b
43 {-# INLINE_STREAM map #-}
44 map f (Stream step s n) = Stream step' s n
45 where
46 {-# INLINE step' #-}
47 step' s = case step s of
48 Yield x s' -> Yield (f x) s'
49 Skip s' -> Skip s'
50 Done -> Done
52 filter :: (a -> Bool) -> Stream a -> Stream a
53 {-# INLINE_STREAM filter #-}
54 filter f (Stream step s n) = Stream step' s n
55 where
56 {-# INLINE step' #-}
57 step' s = case step s of
58 Yield x s' | f x -> Yield x s'
59 | otherwise -> Skip s'
60 Skip s' -> Skip s'
61 Done -> Done
63 zipWith :: (a -> b -> c) -> Stream a -> Stream b -> Stream c
64 {-# INLINE_STREAM zipWith #-}
65 zipWith f (Stream stepa sa na) (Stream stepb sb nb)
66 = Stream step (sa, sb, Nothing) (min na nb)
67 where
68 {-# INLINE step #-}
69 step (sa, sb, Nothing) = case stepa sa of
70 Yield x sa' -> Skip (sa', sb, Just x)
71 Skip sa' -> Skip (sa', sb, Nothing)
72 Done -> Done
74 step (sa, sb, Just x) = case stepb sb of
75 Yield y sb' -> Yield (f x y) (sa, sb', Nothing)
76 Skip sb' -> Skip (sa, sb', Just x)
77 Done -> Done
79 foldl :: (a -> b -> b) -> b -> Stream a -> b
80 {-# INLINE_STREAM foldl #-}
81 foldl f z (Stream step s _) = foldl_go z s
82 where
83 foldl_go z s = case step s of
84 Yield x s' -> foldl_go (f x z) s'
85 Skip s' -> foldl_go z s'
86 Done -> z
88 foldl' :: (a -> b -> b) -> b -> Stream a -> b
89 {-# INLINE_STREAM foldl' #-}
90 foldl' f z (Stream step s _) = foldl_go z s
91 where
92 foldl_go !z s = case step s of
93 Yield x s' -> foldl_go (f x z) s'
94 Skip s' -> foldl_go z s'
95 Done -> z
97 foldr :: (a -> b -> b) -> b -> Stream a -> b
98 {-# INLINE_STREAM foldr #-}
99 foldr f z (Stream step s _) = foldr_go s
100 where
101 foldr_go s = case step s of
102 Yield x s' -> f x (foldr_go s')
103 Skip s' -> foldr_go s'
104 Done -> z
106 mapM_ :: Monad m => (a -> m ()) -> Stream a -> m ()
107 {-# INLINE_STREAM mapM_ #-}
108 mapM_ m (Stream step s _) = mapM_go s
109 where
110 mapM_go s = case step s of
111 Yield x s' -> do { m x; mapM_go s' }
112 Skip s' -> mapM_go s'
113 Done -> return ()
115 foldM :: Monad m => (a -> b -> m a) -> a -> Stream b -> m a
116 {-# INLINE_STREAM foldM #-}
117 foldM m z (Stream step s _) = foldM_go z s
118 where
119 foldM_go z s = case step s of
120 Yield x s' -> do { z' <- m z x; foldM_go z' s' }
121 Skip s' -> foldM_go z s'
122 Done -> return z