Add prescans
authorRoman Leshchinskiy <rl@cse.unsw.edu.au>
Mon, 15 Sep 2008 09:50:32 +0000 (09:50 +0000)
committerRoman Leshchinskiy <rl@cse.unsw.edu.au>
Mon, 15 Sep 2008 09:50:32 +0000 (09:50 +0000)
Data/Vector/Fusion/Stream.hs
Data/Vector/Fusion/Stream/Monadic.hs
Data/Vector/IVector.hs

index 7c47364..2d5ef0e 100644 (file)
@@ -48,6 +48,9 @@ module Data.Vector.Fusion.Stream (
   -- * Unfolding
   unfold,
 
+  -- * Scans
+  prescanl, prescanl',
+
   -- * Conversion to/from lists
   toList, fromList,
 
@@ -98,11 +101,6 @@ sized :: Stream a -> Size -> Stream a
 {-# INLINE sized #-}
 sized = M.sized
 
--- | Unfold
-unfold :: (s -> Maybe (a, s)) -> s -> Stream a
-{-# INLINE unfold #-}
-unfold = M.unfold
-
 -- | Convert a 'Stream' to a list
 toList :: Stream a -> [a]
 {-# INLINE toList #-}
@@ -299,6 +297,27 @@ foldr1 :: (a -> a -> a) -> Stream a -> a
 {-# INLINE foldr1 #-}
 foldr1 f = unId . M.foldr1 f
 
+-- Unfolding
+-- ---------
+
+-- | Unfold
+unfold :: (s -> Maybe (a, s)) -> s -> Stream a
+{-# INLINE unfold #-}
+unfold = M.unfold
+
+-- Scans
+-- -----
+
+-- | Prefix scan
+prescanl :: (a -> b -> a) -> a -> Stream b -> Stream a
+{-# INLINE prescanl #-}
+prescanl = M.prescanl
+
+-- | Prefix scan with strict accumulator
+prescanl' :: (a -> b -> a) -> a -> Stream b -> Stream a
+{-# INLINE prescanl' #-}
+prescanl' = M.prescanl'
+
 -- Comparisons
 -- -----------
 
index 1c740b5..027bfe6 100644 (file)
@@ -37,6 +37,9 @@ module Data.Vector.Fusion.Stream.Monadic (
   -- * Unfolding
   unfold, unfoldM,
 
+  -- * Scans
+  prescanl, prescanlM, prescanl', prescanlM',
+
   toList, fromList
 ) where
 
@@ -593,6 +596,50 @@ unfoldM f s = Stream step s Unknown
                  Nothing      -> Done
              ) (f s)
 
+-- Scans
+-- -----
+
+-- | Prefix scan
+prescanl :: Monad m => (a -> b -> a) -> a -> Stream m b -> Stream m a
+{-# INLINE prescanl #-}
+prescanl f = prescanlM (\a b -> return (f a b))
+
+-- | Prefix scan
+prescanlM :: Monad m => (a -> b -> m a) -> a -> Stream m b -> Stream m a
+{-# INLINE_STREAM prescanlM #-}
+prescanlM f z (Stream step s sz) = Stream step' (s,z) sz
+  where
+    {-# INLINE step' #-}
+    step' (s,x) = do
+                    r <- step s
+                    case r of
+                      Yield y s' -> do
+                                      z <- f x y
+                                      return $ Yield x (s', z)
+                      Skip    s' -> return $ Skip (s', x)
+                      Done       -> return Done
+
+
+-- | Prefix scan with strict accumulator
+prescanl' :: Monad m => (a -> b -> a) -> a -> Stream m b -> Stream m a
+{-# INLINE prescanl' #-}
+prescanl' f = prescanlM' (\a b -> return (f a b))
+
+-- | Prefix scan with strict accumulator
+prescanlM' :: Monad m => (a -> b -> m a) -> a -> Stream m b -> Stream m a
+{-# INLINE_STREAM prescanlM' #-}
+prescanlM' f z (Stream step s sz) = Stream step' (s,z) sz
+  where
+    {-# INLINE step' #-}
+    step' (s,x) = x `seq`
+                  do
+                    r <- step s
+                    case r of
+                      Yield y s' -> do
+                                      z <- f x y
+                                      return $ Yield x (s', z)
+                      Skip    s' -> return $ Skip (s', x)
+                      Done       -> return Done
 
 -- Conversions
 -- -----------
index 513365a..c0e567a 100644 (file)
@@ -48,6 +48,9 @@ module Data.Vector.IVector (
   -- * Folding
   foldl, foldl1, foldl', foldl1', foldr, foldr1,
 
+  -- * Scans
+  prescanl, prescanl',
+
   -- * Conversion to/from lists
   toList, fromList,
 
@@ -474,6 +477,40 @@ foldr1 :: IVector v a => (a -> a -> a) -> v a -> a
 {-# INLINE foldr1 #-}
 foldr1 f = Stream.foldr1 f . stream
 
+-- Scans
+-- -----
+
+-- | Prefix scan
+prescanl :: (IVector v a, IVector v b) => (a -> b -> a) -> a -> v b -> v a
+{-# INLINE prescanl #-}
+prescanl f z = unstream . Stream.prescanl f z . stream
+
+inplace_prescanl :: IVector v a => (a -> a -> a) -> a -> v a -> v a
+{-# INLINE inplace_prescanl #-}
+inplace_prescanl f z = unstream . inplace (MStream.prescanl f z) . stream
+
+{-# RULES
+
+"prescanl -> inplace_prescanl [IVector]" prescanl = inplace_prescanl
+
+ #-}
+
+-- | Prefix scan with strict accumulator
+prescanl' :: (IVector v a, IVector v b) => (a -> b -> a) -> a -> v b -> v a
+{-# INLINE prescanl' #-}
+prescanl' f z = unstream . Stream.prescanl' f z . stream
+
+inplace_prescanl' :: IVector v a => (a -> a -> a) -> a -> v a -> v a
+{-# INLINE inplace_prescanl' #-}
+inplace_prescanl' f z = unstream . inplace (MStream.prescanl' f z) . stream
+
+{-# RULES
+
+"prescanl' -> inplace_prescanl' [IVector]" prescanl' = inplace_prescanl'
+
+ #-}
+
+
 -- | Convert a vector to a list
 toList :: IVector v a => v a -> [a]
 {-# INLINE toList #-}