Merge pull request #117 from pjrt/master

Add `iscan` with its left, right, lazy and strict versions

Author: dolio

Data/Vector.hs

@@ -143,9 +143,11 @@ module Data.Vector (
   prescanl, prescanl',
   postscanl, postscanl',
   scanl, scanl', scanl1, scanl1',
+  iscanl, iscanl',
   prescanr, prescanr',
   postscanr, postscanr',
   scanr, scanr', scanr1, scanr1',
+  iscanr, iscanr',
 
   -- * Conversions
 
@@ -1456,6 +1458,16 @@ scanl' :: (a -> b -> a) -> a -> Vector b -> Vector a
 {-# INLINE scanl' #-}
 scanl' = G.scanl'
 
+-- | /O(n)/ Scan over a vector with its index
+iscanl :: (Int -> a -> b -> a) -> a -> Vector b -> Vector a
+{-# INLINE iscanl #-}
+iscanl = G.iscanl
+
+-- | /O(n)/ Scan over a vector (strictly) with its index
+iscanl' :: (Int -> a -> b -> a) -> a -> Vector b -> Vector a
+{-# INLINE iscanl' #-}
+iscanl' = G.iscanl'
+
 -- | /O(n)/ Scan over a non-empty vector
 --
 -- > scanl f <x1,...,xn> = <y1,...,yn>
@@ -1506,6 +1518,16 @@ scanr' :: (a -> b -> b) -> b -> Vector a -> Vector b
 {-# INLINE scanr' #-}
 scanr' = G.scanr'
 
+-- | /O(n)/ Right-to-left scan over a vector with its index
+iscanr :: (Int -> a -> b -> b) -> b -> Vector a -> Vector b
+{-# INLINE iscanr #-}
+iscanr = G.iscanr
+
+-- | /O(n)/ Right-to-left scan over a vector (strictly) with its index
+iscanr' :: (Int -> a -> b -> b) -> b -> Vector a -> Vector b
+{-# INLINE iscanr' #-}
+iscanr' = G.iscanr'
+
 -- | /O(n)/ Right-to-left scan over a non-empty vector
 scanr1 :: (a -> a -> a) -> Vector a -> Vector a
 {-# INLINE scanr1 #-}

Data/Vector/Generic.hs

@@ -126,9 +126,11 @@ module Data.Vector.Generic (
   prescanl, prescanl',
   postscanl, postscanl',
   scanl, scanl', scanl1, scanl1',
+  iscanl, iscanl',
   prescanr, prescanr',
   postscanr, postscanr',
   scanr, scanr', scanr1, scanr1',
+  iscanr, iscanr',
 
   -- * Conversions
 
@@ -1721,6 +1723,23 @@ scanl' :: (Vector v a, Vector v b) => (a -> b -> a) -> a -> v b -> v a
 {-# INLINE scanl' #-}
 scanl' f z = unstream . Bundle.scanl' f z . stream
 
+-- | /O(n)/ Scan over a vector with its index
+iscanl :: (Vector v a, Vector v b) => (Int -> a -> b -> a) -> a -> v b -> v a
+{-# INLINE iscanl #-}
+iscanl f z =
+    unstream
+  . inplace (S.scanl (\a (i, b) -> f i a b) z . S.indexed) id
+  . stream
+
+-- | /O(n)/ Scan over a vector (strictly) with its index
+iscanl' :: (Vector v a, Vector v b) => (Int -> a -> b -> a) -> a -> v b -> v a
+{-# INLINE iscanl' #-}
+iscanl' f z =
+    unstream
+  . inplace (S.scanl' (\a (i, b) -> f i a b) z . S.indexed) id
+  . stream
+
+
 -- | /O(n)/ Scan over a non-empty vector
 --
 -- > scanl f <x1,...,xn> = <y1,...,yn>
@@ -1771,6 +1790,22 @@ scanr' :: (Vector v a, Vector v b) => (a -> b -> b) -> b -> v a -> v b
 {-# INLINE scanr' #-}
 scanr' f z = unstreamR . Bundle.scanl' (flip f) z . streamR
 
+-- | /O(n)/ Right-to-left scan over a vector with its index
+iscanr :: (Vector v a, Vector v b) => (Int -> a -> b -> b) -> b -> v a -> v b
+{-# INLINE iscanr #-}
+iscanr f z =
+    unstreamR
+  . inplace (S.scanl (flip $ uncurry f) z . S.indexed) id
+  . streamR
+
+-- | /O(n)/ Right-to-left scan over a vector (strictly) with its index
+iscanr' :: (Vector v a, Vector v b) => (Int -> a -> b -> b) -> b -> v a -> v b
+{-# INLINE iscanr' #-}
+iscanr' f z =
+    unstreamR
+  . inplace (S.scanl' (flip $ uncurry f) z . S.indexed) id
+  . streamR
+
 -- | /O(n)/ Right-to-left scan over a non-empty vector
 scanr1 :: Vector v a => (a -> a -> a) -> v a -> v a
 {-# INLINE scanr1 #-}

tests/Tests/Vector.hs

@@ -191,10 +191,12 @@ testPolymorphicFunctions _ = $(testProperties [
         'prop_prescanl, 'prop_prescanl',
         'prop_postscanl, 'prop_postscanl',
         'prop_scanl, 'prop_scanl', 'prop_scanl1, 'prop_scanl1',
+        'prop_iscanl, 'prop_iscanl',
 
         'prop_prescanr, 'prop_prescanr',
         'prop_postscanr, 'prop_postscanr',
-        'prop_scanr, 'prop_scanr', 'prop_scanr1, 'prop_scanr1'
+        'prop_scanr, 'prop_scanr', 'prop_scanr1, 'prop_scanr1',
+        'prop_iscanr, 'prop_iscanr'
     ])
   where
     -- Prelude
@@ -357,6 +359,10 @@ testPolymorphicFunctions _ = $(testProperties [
                  V.scanl1 `eq` scanl1
     prop_scanl1' :: P ((a -> a -> a) -> v a -> v a) = notNull2 ===>
                  V.scanl1' `eq` scanl1
+    prop_iscanl :: P ((Int -> a -> a -> a) -> a -> v a -> v a)
+                = V.iscanl `eq` iscanl
+    prop_iscanl' :: P ((Int -> a -> a -> a) -> a -> v a -> v a)
+               = V.iscanl' `eq` iscanl
 
     prop_prescanr :: P ((a -> a -> a) -> a -> v a -> v a)
                 = V.prescanr `eq` prescanr
@@ -370,6 +376,10 @@ testPolymorphicFunctions _ = $(testProperties [
                 = V.scanr `eq` scanr
     prop_scanr' :: P ((a -> a -> a) -> a -> v a -> v a)
                = V.scanr' `eq` scanr
+    prop_iscanr :: P ((Int -> a -> a -> a) -> a -> v a -> v a)
+                = V.iscanr `eq` iscanr
+    prop_iscanr' :: P ((Int -> a -> a -> a) -> a -> v a -> v a)
+               = V.iscanr' `eq` iscanr
     prop_scanr1 :: P ((a -> a -> a) -> v a -> v a) = notNull2 ===>
                  V.scanr1 `eq` scanr1
     prop_scanr1' :: P ((a -> a -> a) -> v a -> v a) = notNull2 ===>

tests/Utilities.hs

@@ -297,6 +297,12 @@ indexedLeftFold fld f z = fld (uncurry . f) z . zip [0..]
 ifoldl :: (a -> Int -> a -> a) -> a -> [a] -> a
 ifoldl = indexedLeftFold foldl
 
+iscanl :: (Int -> a -> b -> a) -> a -> [b] -> [a]
+iscanl f z = scanl (\a (i, b) -> f i a b) z . zip [0..]
+
+iscanr :: (Int -> a -> b -> b) -> b -> [a] -> [b]
+iscanr f z = scanr (uncurry f) z . zip [0..]
+
 ifoldr :: (Int -> a -> b -> b) -> b -> [a] -> b
 ifoldr f z = foldr (uncurry f) z . zip [0..]