New combinators
[darcs-mirrors/vector.git] / Data / Vector / Stream.hs
index a342cd9..ad1b2c7 100644 (file)
@@ -5,13 +5,21 @@
 module Data.Vector.Stream (
   Step(..), Stream(..),
 
 module Data.Vector.Stream (
   Step(..), Stream(..),
 
-  empty, singleton, replicate,
-  map, filter, zipWith,
+  size, sized, unfold, toList, fromList,
+  empty, singleton, replicate, (++),
+  take, drop,
+  map, zipWith,
+  filter, takeWhile, dropWhile,
   foldr, foldl, foldl',
   mapM_, foldM
 ) where
 
   foldr, foldl, foldl',
   mapM_, foldM
 ) where
 
-import Prelude hiding ( replicate, map, filter, zipWith,
+import Data.Vector.Stream.Size
+
+import Prelude hiding ( replicate, (++),
+                        take, drop,
+                        map, zipWith,
+                        filter, takeWhile, dropWhile,
                         foldr, foldl,
                         mapM_ )
 
                         foldr, foldl,
                         mapM_ )
 
@@ -19,15 +27,46 @@ data Step s a = Yield a s
               | Skip    s
               | Done
 
               | Skip    s
               | Done
 
-data Stream a = forall s. Stream (s -> Step s a) s Int
+data Stream a = forall s. Stream (s -> Step s a) s Size
+
+size :: Stream a -> Size
+{-# INLINE size #-}
+size (Stream _ _ sz) = sz
+
+sized :: Stream a -> Size -> Stream a
+{-# INLINE_STREAM sized #-}
+sized (Stream step s _) sz = Stream step s sz
+
+unfold :: (s -> Maybe (a, s)) -> s -> Stream a
+{-# INLINE_STREAM unfold #-}
+unfold f s = Stream step s Unknown
+  where
+    {-# INLINE step #-}
+    step s = case f s of
+               Just (x, s') -> Yield x s'
+               Nothing      -> Done
+
+toList :: Stream a -> [a]
+{-# INLINE toList #-}
+toList s = foldr (:) [] s
+
+fromList :: [a] -> Stream a
+{-# INLINE_STREAM fromList #-}
+fromList xs = Stream step xs Unknown
+  where
+    step (x:xs) = Yield x xs
+    step []     = Done
+
+-- Construction
+-- ------------
 
 empty :: Stream a
 {-# INLINE_STREAM empty #-}
 
 empty :: Stream a
 {-# INLINE_STREAM empty #-}
-empty = Stream (const Done) () 0
+empty = Stream (const Done) () (Exact 0)
 
 singleton :: a -> Stream a
 {-# INLINE_STREAM singleton #-}
 
 singleton :: a -> Stream a
 {-# INLINE_STREAM singleton #-}
-singleton x = Stream step True 1
+singleton x = Stream step True (Exact 1)
   where
     {-# INLINE step #-}
     step True  = Yield x False
   where
     {-# INLINE step #-}
     step True  = Yield x False
@@ -35,12 +74,66 @@ singleton x = Stream step True 1
 
 replicate :: Int -> a -> Stream a
 {-# INLINE_STREAM replicate #-}
 
 replicate :: Int -> a -> Stream a
 {-# INLINE_STREAM replicate #-}
-replicate n x = Stream step n (max n 0)
+replicate n x = Stream step n (Exact (max n 0))
   where
     {-# INLINE step #-}
     step i | i > 0     = Yield x (i-1)
            | otherwise = Done
 
   where
     {-# INLINE step #-}
     step i | i > 0     = Yield x (i-1)
            | otherwise = Done
 
+infixr ++
+(++) :: Stream a -> Stream a -> Stream a
+{-# INLINE_STREAM (++) #-}
+Stream stepa sa na ++ Stream stepb sb nb = Stream step (Left sa) (na + nb)
+  where
+    step (Left  sa) = case stepa sa of
+                        Yield x sa' -> Yield x (Left  sa')
+                        Skip    sa' -> Skip    (Left  sa')
+                        Done        -> Skip    (Right sb)
+    step (Right sb) = case stepb sb of
+                        Yield x sb' -> Yield x (Right sb')
+                        Skip    sb' -> Skip    (Right sb')
+                        Done        -> Done
+
+-- Substreams
+-- ----------
+
+take :: Int -> Stream a -> Stream a
+{-# INLINE_STREAM take #-}
+take n (Stream step s sz) = Stream step' (s, 0) (smaller (Exact n) sz)
+  where
+    {-# INLINE step' #-}
+    step' (s, i) | i < n = case step s of
+                             Yield x s' -> Yield x (s', i+1)
+                             Skip    s' -> Skip    (s', i)
+                             Done       -> Done
+    step' (s, i) = Done
+
+data Drop s = Drop_Drop s Int | Drop_Keep s
+
+drop :: Int -> Stream a -> Stream a
+{-# INLINE_STREAM drop #-}
+drop n (Stream step s sz) = Stream step' (Drop_Drop s 0) (sz - Exact n)
+  where
+    {-# INLINE step' #-}
+    step' (Drop_Drop s i) | i < n = case step s of
+                                      Yield x s' -> Skip (Drop_Drop s' (i+1))
+                                      Skip    s' -> Skip (Drop_Drop s' i)
+                                      Done       -> Done
+                          | otherwise = Skip (Drop_Keep s)
+
+    step' (Drop_Keep s) = case step s of
+                            Yield x s' -> Yield x (Drop_Keep s')
+                            Skip    s' -> Skip    (Drop_Keep s')
+                            Done       -> Done
+                     
+
+-- Mapping/zipping
+-- ---------------
+
+instance Functor Stream where
+  {-# INLINE_STREAM fmap #-}
+  fmap = map
+
 map :: (a -> b) -> Stream a -> Stream b
 {-# INLINE_STREAM map #-}
 map f (Stream step s n) = Stream step' s n
 map :: (a -> b) -> Stream a -> Stream b
 {-# INLINE_STREAM map #-}
 map f (Stream step s n) = Stream step' s n
@@ -51,21 +144,10 @@ map f (Stream step s n) = Stream step' s n
                 Skip    s' -> Skip        s'
                 Done       -> Done
 
                 Skip    s' -> Skip        s'
                 Done       -> Done
 
-filter :: (a -> Bool) -> Stream a -> Stream a
-{-# INLINE_STREAM filter #-}
-filter f (Stream step s n) = Stream step' s n
-  where
-    {-# INLINE step' #-}
-    step' s = case step s of
-                Yield x s' | f x       -> Yield x s'
-                           | otherwise -> Skip    s'
-                Skip    s'             -> Skip    s'
-                Done                   -> Done
-
 zipWith :: (a -> b -> c) -> Stream a -> Stream b -> Stream c
 {-# INLINE_STREAM zipWith #-}
 zipWith f (Stream stepa sa na) (Stream stepb sb nb)
 zipWith :: (a -> b -> c) -> Stream a -> Stream b -> Stream c
 {-# INLINE_STREAM zipWith #-}
 zipWith f (Stream stepa sa na) (Stream stepb sb nb)
-  = Stream step (sa, sb, Nothing) (min na nb)
+  = Stream step (sa, sb, Nothing) (smaller na nb)
   where
     {-# INLINE step #-}
     step (sa, sb, Nothing) = case stepa sa of
   where
     {-# INLINE step #-}
     step (sa, sb, Nothing) = case stepa sa of
@@ -78,21 +160,74 @@ zipWith f (Stream stepa sa na) (Stream stepb sb nb)
                                Skip    sb' -> Skip          (sa, sb', Just x)
                                Done        -> Done
 
                                Skip    sb' -> Skip          (sa, sb', Just x)
                                Done        -> Done
 
-foldl :: (a -> b -> b) -> b -> Stream a -> b
+-- Filtering
+-- ---------
+
+filter :: (a -> Bool) -> Stream a -> Stream a
+{-# INLINE_STREAM filter #-}
+filter f (Stream step s n) = Stream step' s (toMax n)
+  where
+    {-# INLINE step' #-}
+    step' s = case step s of
+                Yield x s' | f x       -> Yield x s'
+                           | otherwise -> Skip    s'
+                Skip    s'             -> Skip    s'
+                Done                   -> Done
+
+takeWhile :: (a -> Bool) -> Stream a -> Stream a
+{-# INLINE_STREAM takeWhile #-}
+takeWhile f (Stream step s n) = Stream step' s (toMax n)
+  where
+    {-# INLINE step' #-}
+    step' s = case step s of
+                Yield x s' | f x       -> Yield x s'
+                           | otherwise -> Done
+                Skip    s'             -> Skip s'
+                Done                   -> Done
+
+
+data DropWhile s a = DropWhile_Drop s | DropWhile_Yield a s | DropWhile_Next s
+
+dropWhile :: (a -> Bool) -> Stream a -> Stream a
+{-# INLINE_STREAM dropWhile #-}
+dropWhile f (Stream step s n) = Stream step' (DropWhile_Drop s) (toMax n)
+  where
+    -- NOTE: we jump through hoops here to have only one Yield; local data
+    -- declarations would be nice!
+
+    {-# INLINE step' #-}
+    step' (DropWhile_Drop s)
+      = case step s of
+          Yield x s' | f x       -> Skip    (DropWhile_Drop    s')
+                     | otherwise -> Skip    (DropWhile_Yield x s')
+          Skip    s'             -> Skip    (DropWhile_Drop    s')
+          Done                   -> Done
+
+    step' (DropWhile_Yield x s) = Yield x (DropWhile_Next s)
+
+    step' (DropWhile_Next s) = case step s of
+                                 Yield x s' -> Skip    (DropWhile_Yield x s')
+                                 Skip    s' -> Skip    (DropWhile_Next    s')
+                                 Done       -> Done
+
+-- Folding
+-- -------
+
+foldl :: (a -> b -> a) -> a -> Stream b -> a
 {-# INLINE_STREAM foldl #-}
 foldl f z (Stream step s _) = foldl_go z s
   where
     foldl_go z s = case step s of
 {-# INLINE_STREAM foldl #-}
 foldl f z (Stream step s _) = foldl_go z s
   where
     foldl_go z s = case step s of
-                     Yield x s' -> foldl_go (f x z) s'
+                     Yield x s' -> foldl_go (f z x) s'
                      Skip    s' -> foldl_go z       s'
                      Done       -> z
 
                      Skip    s' -> foldl_go z       s'
                      Done       -> z
 
-foldl' :: (a -> b -> b) -> b -> Stream a -> b
+foldl' :: (a -> b -> a) -> a -> Stream b -> a
 {-# INLINE_STREAM foldl' #-}
 {-# INLINE_STREAM foldl' #-}
-foldl' f z (Stream step s _) = foldl_go z s
+foldl' f !z (Stream step s _) = foldl_go z s
   where
     foldl_go !z s = case step s of
   where
     foldl_go !z s = case step s of
-                      Yield x s' -> foldl_go (f x z) s'
+                      Yield x s' -> foldl_go (f z x) s'
                       Skip    s' -> foldl_go z       s'
                       Done       -> z
 
                       Skip    s' -> foldl_go z       s'
                       Done       -> z