Add lots of index-related collective operations
authorRoman Leshchinskiy <rl@cse.unsw.edu.au>
Tue, 8 Dec 2009 15:40:40 +0000 (15:40 +0000)
committerRoman Leshchinskiy <rl@cse.unsw.edu.au>
Tue, 8 Dec 2009 15:40:40 +0000 (15:40 +0000)
Data/Vector.hs
Data/Vector/Fusion/Stream.hs
Data/Vector/Fusion/Stream/Monadic.hs
Data/Vector/Generic.hs
Data/Vector/Primitive.hs
Data/Vector/Storable.hs
Data/Vector/Unboxed.hs

index 5afa0e6..98142e0 100644 (file)
@@ -34,24 +34,26 @@ module Data.Vector (
   backpermute, reverse,
 
   -- * Mapping
-  map, concatMap,
+  map, imap, concatMap,
 
   -- * Zipping and unzipping
   zipWith, zipWith3, zipWith4, zipWith5, zipWith6,
+  izipWith, izipWith3, izipWith4, izipWith5, izipWith6,
   zip, zip3, zip4, zip5, zip6,
   unzip, unzip3, unzip4, unzip5, unzip6,
 
   -- * Filtering
-  filter, takeWhile, dropWhile,
+  filter, ifilter, takeWhile, dropWhile,
 
   -- * Searching
   elem, notElem, find, findIndex,
 
   -- * Folding
   foldl, foldl1, foldl', foldl1', foldr, foldr1,
+  ifoldl, ifoldl', ifoldr,
 
   -- * Specialised folds
-  and, or, sum, product, maximum, minimum,
+  and, or, sum, product, maximum, minimum, minIndex, maxIndex,
 
   -- * Unfolding
   unfoldr,
@@ -324,6 +326,11 @@ map :: (a -> b) -> Vector a -> Vector b
 {-# INLINE map #-}
 map = G.map
 
+-- | Apply a function to every index/value pair
+imap :: (Int -> a -> b) -> Vector a -> Vector b
+{-# INLINE imap #-}
+imap = G.imap
+
 concatMap :: (a -> Vector b) -> Vector a -> Vector b
 {-# INLINE concatMap #-}
 concatMap = G.concatMap
@@ -358,6 +365,34 @@ zipWith6 :: (a -> b -> c -> d -> e -> f -> g)
 {-# INLINE zipWith6 #-}
 zipWith6 = G.zipWith6
 
+-- | Zip two vectors and their indices with the given function.
+izipWith :: (Int -> a -> b -> c) -> Vector a -> Vector b -> Vector c
+{-# INLINE izipWith #-}
+izipWith = G.izipWith
+
+-- | Zip three vectors and their indices with the given function.
+izipWith3 :: (Int -> a -> b -> c -> d)
+          -> Vector a -> Vector b -> Vector c -> Vector d
+{-# INLINE izipWith3 #-}
+izipWith3 = G.izipWith3
+
+izipWith4 :: (Int -> a -> b -> c -> d -> e)
+          -> Vector a -> Vector b -> Vector c -> Vector d -> Vector e
+{-# INLINE izipWith4 #-}
+izipWith4 = G.izipWith4
+
+izipWith5 :: (Int -> a -> b -> c -> d -> e -> f)
+          -> Vector a -> Vector b -> Vector c -> Vector d -> Vector e
+          -> Vector f
+{-# INLINE izipWith5 #-}
+izipWith5 = G.izipWith5
+
+izipWith6 :: (Int -> a -> b -> c -> d -> e -> f -> g)
+          -> Vector a -> Vector b -> Vector c -> Vector d -> Vector e
+          -> Vector f -> Vector g
+{-# INLINE izipWith6 #-}
+izipWith6 = G.izipWith6
+
 zip :: Vector a -> Vector b -> Vector (a, b)
 {-# INLINE zip #-}
 zip = G.zip
@@ -411,6 +446,12 @@ filter :: (a -> Bool) -> Vector a -> Vector a
 {-# INLINE filter #-}
 filter = G.filter
 
+-- | Drop elements that do not satisfy the predicate (applied to values and
+-- their indices)
+ifilter :: (Int -> a -> Bool) -> Vector a -> Vector a
+{-# INLINE ifilter #-}
+ifilter = G.ifilter
+
 -- | Yield the longest prefix of elements satisfying the predicate.
 takeWhile :: (a -> Bool) -> Vector a -> Vector a
 {-# INLINE takeWhile #-}
@@ -481,6 +522,22 @@ foldr1 :: (a -> a -> a) -> Vector a -> a
 {-# INLINE foldr1 #-}
 foldr1 = G.foldr1
 
+-- | Left fold (function applied to each element and its index)
+ifoldl :: (a -> Int -> b -> a) -> a -> Vector b -> a
+{-# INLINE ifoldl #-}
+ifoldl = G.ifoldl
+
+-- | Left fold with strict accumulator (function applied to each element and
+-- its index)
+ifoldl' :: (a -> Int -> b -> a) -> a -> Vector b -> a
+{-# INLINE ifoldl' #-}
+ifoldl' = G.ifoldl'
+
+-- | Right fold (function applied to each element and its index)
+ifoldr :: (Int -> a -> b -> b) -> b -> Vector a -> b
+{-# INLINE ifoldr #-}
+ifoldr = G.ifoldr
+
 -- Specialised folds
 -- -----------------
 
@@ -508,6 +565,14 @@ minimum :: Ord a => Vector a -> a
 {-# INLINE minimum #-}
 minimum = G.minimum
 
+minIndex :: Ord a => Vector a -> Int
+{-# INLINE minIndex #-}
+minIndex = G.minIndex
+
+maxIndex :: Ord a => Vector a -> Int
+{-# INLINE maxIndex #-}
+maxIndex = G.maxIndex
+
 -- Unfolding
 -- ---------
 
index 804e625..c41fbfb 100644 (file)
@@ -38,6 +38,7 @@ module Data.Vector.Fusion.Stream (
   map, concatMap, unbox,
   
   -- * Zipping
+  indexed,
   zipWith, zipWith3, zipWith4, zipWith5, zipWith6,
   zip, zip3, zip4, zip5, zip6,
 
@@ -243,6 +244,11 @@ concatMap = M.concatMap
 -- Zipping
 -- -------
 
+-- | Associate all 'Stream' elements with their indices
+indexed :: Stream a -> Stream (Int,a)
+{-# INLINE indexed #-}
+indexed = M.indexed
+
 -- | Zip two 'Stream's with the given function
 zipWith :: (a -> b -> c) -> Stream a -> Stream b -> Stream c
 {-# INLINE zipWith #-}
index 6b64856..e495368 100644 (file)
@@ -34,6 +34,7 @@ module Data.Vector.Fusion.Stream.Monadic (
   map, mapM, mapM_, trans, unbox, concatMap,
   
   -- * Zipping
+  indexed,
   zipWithM, zipWith3M, zipWith4M, zipWith5M, zipWith6M,
   zipWith, zipWith3, zipWith4, zipWith5, zipWith6,
   zip, zip3, zip4, zip5, zip6,
@@ -346,6 +347,7 @@ map :: Monad m => (a -> b) -> Stream m a -> Stream m b
 {-# INLINE map #-}
 map f = mapM (return . f)
 
+
 -- | Map a monadic function over a 'Stream'
 mapM :: Monad m => (a -> m b) -> Stream m a -> Stream m b
 {-# INLINE_STREAM mapM #-}
@@ -393,6 +395,20 @@ unbox (Stream step s n) = Stream step' s n
 -- Zipping
 -- -------
 
+-- | Pair each element in a 'Stream' with its index
+indexed :: Monad m => Stream m a -> Stream m (Int,a)
+{-# INLINE_STREAM indexed #-}
+indexed (Stream step s n) = Stream step' (s,0) n
+  where
+    {-# INLINE_INNER step' #-}
+    step' (s,i) = i `seq`
+                  do
+                    r <- step s
+                    case r of
+                      Yield x s' -> return $ Yield (i,x) (s', i+1)
+                      Skip    s' -> return $ Skip        (s', i)
+                      Done       -> return Done
+
 -- | Zip two 'Stream's with the given monadic function
 zipWithM :: Monad m => (a -> b -> m c) -> Stream m a -> Stream m b -> Stream m c
 {-# INLINE_STREAM zipWithM #-}
index 8c8b320..7ef4382 100644 (file)
@@ -35,10 +35,11 @@ module Data.Vector.Generic (
   backpermute, reverse,
 
   -- * Mapping
-  map, concatMap,
+  map, imap, concatMap,
 
   -- * Zipping and unzipping
   zipWith, zipWith3, zipWith4, zipWith5, zipWith6,
+  izipWith, izipWith3, izipWith4, izipWith5, izipWith6,
   zip, zip3, zip4, zip5, zip6,
   unzip, unzip3, unzip4, unzip5, unzip6,
 
@@ -46,16 +47,17 @@ module Data.Vector.Generic (
   eq, cmp,
 
   -- * Filtering
-  filter, takeWhile, dropWhile,
+  filter, ifilter, takeWhile, dropWhile,
 
   -- * Searching
   elem, notElem, find, findIndex,
 
   -- * Folding
   foldl, foldl1, foldl', foldl1', foldr, foldr1,
+  ifoldl, ifoldl', ifoldr,
  
   -- * Specialised folds
-  and, or, sum, product, maximum, minimum,
+  and, or, sum, product, maximum, minimum, minIndex, maxIndex,
 
   -- * Unfolding
   unfoldr,
@@ -523,6 +525,12 @@ map :: (Vector v a, Vector v b) => (a -> b) -> v a -> v b
 {-# INLINE map #-}
 map f = unstream . inplace (MStream.map f) . stream
 
+-- | Apply a function to every index/value pair
+imap :: (Vector v a, Vector v b) => (Int -> a -> b) -> v a -> v b
+{-# INLINE imap #-}
+imap f = unstream . inplace (MStream.map (uncurry f) . MStream.indexed)
+                  . stream
+
 concatMap :: (Vector v a, Vector v b) => (a -> v b) -> v a -> v b
 {-# INLINE concatMap #-}
 concatMap f = unstream . Stream.concatMap (stream . f) . stream
@@ -531,7 +539,8 @@ concatMap f = unstream . Stream.concatMap (stream . f) . stream
 -- -----------------
 
 -- | Zip two vectors with the given function.
-zipWith :: (Vector v a, Vector v b, Vector v c) => (a -> b -> c) -> v a -> v b -> v c
+zipWith :: (Vector v a, Vector v b, Vector v c)
+        => (a -> b -> c) -> v a -> v b -> v c
 {-# INLINE zipWith #-}
 zipWith f xs ys = unstream (Stream.zipWith f (stream xs) (stream ys))
 
@@ -577,6 +586,57 @@ zipWith6 f as bs cs ds es fs
                                 (stream es)
                                 (stream fs))
 
+-- | Zip two vectors and their indices with the given function.
+izipWith :: (Vector v a, Vector v b, Vector v c)
+        => (Int -> a -> b -> c) -> v a -> v b -> v c
+{-# INLINE izipWith #-}
+izipWith f xs ys = unstream
+                  (Stream.zipWith (uncurry f) (Stream.indexed (stream xs))
+                                                              (stream ys))
+
+-- | Zip three vectors and their indices with the given function.
+izipWith3 :: (Vector v a, Vector v b, Vector v c, Vector v d)
+         => (Int -> a -> b -> c -> d) -> v a -> v b -> v c -> v d
+{-# INLINE izipWith3 #-}
+izipWith3 f as bs cs
+  = unstream (Stream.zipWith3 (uncurry f) (Stream.indexed (stream as))
+                                                          (stream bs)
+                                                          (stream cs))
+
+izipWith4 :: (Vector v a, Vector v b, Vector v c, Vector v d, Vector v e)
+         => (Int -> a -> b -> c -> d -> e) -> v a -> v b -> v c -> v d -> v e
+{-# INLINE izipWith4 #-}
+izipWith4 f as bs cs ds
+  = unstream (Stream.zipWith4 (uncurry f) (Stream.indexed (stream as))
+                                                          (stream bs)
+                                                          (stream cs)
+                                                          (stream ds))
+
+izipWith5 :: (Vector v a, Vector v b, Vector v c, Vector v d, Vector v e,
+             Vector v f)
+         => (Int -> a -> b -> c -> d -> e -> f) -> v a -> v b -> v c -> v d
+                                                -> v e -> v f
+{-# INLINE izipWith5 #-}
+izipWith5 f as bs cs ds es
+  = unstream (Stream.zipWith5 (uncurry f) (Stream.indexed (stream as))
+                                                          (stream bs)
+                                                          (stream cs)
+                                                          (stream ds)
+                                                          (stream es))
+
+izipWith6 :: (Vector v a, Vector v b, Vector v c, Vector v d, Vector v e,
+             Vector v f, Vector v g)
+         => (Int -> a -> b -> c -> d -> e -> f -> g)
+         -> v a -> v b -> v c -> v d -> v e -> v f -> v g
+{-# INLINE izipWith6 #-}
+izipWith6 f as bs cs ds es fs
+  = unstream (Stream.zipWith6 (uncurry f) (Stream.indexed (stream as))
+                                                          (stream bs)
+                                                          (stream cs)
+                                                          (stream ds)
+                                                          (stream es)
+                                                          (stream fs))
+
 zip :: (Vector v a, Vector v b, Vector v (a,b)) => v a -> v b -> v (a, b)
 {-# INLINE zip #-}
 zip = zipWith (,)
@@ -658,11 +718,20 @@ cmp xs ys = compare (stream xs) (stream ys)
 -- Filtering
 -- ---------
 
--- | Drop elements which do not satisfy the predicate
+-- | Drop elements that do not satisfy the predicate
 filter :: Vector v a => (a -> Bool) -> v a -> v a
 {-# INLINE filter #-}
 filter f = unstream . inplace (MStream.filter f) . stream
 
+-- | Drop elements that do not satisfy the predicate (applied to values and
+-- their indices)
+ifilter :: Vector v a => (Int -> a -> Bool) -> v a -> v a
+{-# INLINE ifilter #-}
+ifilter f = unstream
+          . inplace (MStream.map snd . MStream.filter (uncurry f)
+                                     . MStream.indexed)
+          . stream
+
 -- | Yield the longest prefix of elements satisfying the predicate.
 takeWhile :: Vector v a => (a -> Bool) -> v a -> v a
 {-# INLINE takeWhile #-}
@@ -708,7 +777,7 @@ foldl :: Vector v b => (a -> b -> a) -> a -> v b -> a
 {-# INLINE foldl #-}
 foldl f z = Stream.foldl f z . stream
 
--- | Lefgt fold on non-empty vectors
+-- | Left fold on non-empty vectors
 foldl1 :: Vector v a => (a -> a -> a) -> v a -> a
 {-# INLINE foldl1 #-}
 foldl1 f = Stream.foldl1 f . stream
@@ -733,6 +802,22 @@ foldr1 :: Vector v a => (a -> a -> a) -> v a -> a
 {-# INLINE foldr1 #-}
 foldr1 f = Stream.foldr1 f . stream
 
+-- | Left fold (function applied to each element and its index)
+ifoldl :: Vector v b => (a -> Int -> b -> a) -> a -> v b -> a
+{-# INLINE ifoldl #-}
+ifoldl f z = Stream.foldl (uncurry . f) z . Stream.indexed . stream
+
+-- | Left fold with strict accumulator (function applied to each element and
+-- its index)
+ifoldl' :: Vector v b => (a -> Int -> b -> a) -> a -> v b -> a
+{-# INLINE ifoldl' #-}
+ifoldl' f z = Stream.foldl' (uncurry . f) z . Stream.indexed . stream
+
+-- | Right fold (function applied to each element and its index)
+ifoldr :: Vector v a => (Int -> a -> b -> b) -> b -> v a -> b
+{-# INLINE ifoldr #-}
+ifoldr f z = Stream.foldr (uncurry f) z . Stream.indexed . stream
+
 -- Specialised folds
 -- -----------------
 
@@ -760,6 +845,21 @@ minimum :: (Vector v a, Ord a) => v a -> a
 {-# INLINE minimum #-}
 minimum = Stream.foldl1' min . stream
 
+minIndex :: (Vector v a, Ord a) => v a -> Int
+{-# INLINE minIndex #-}
+minIndex = fst . Stream.foldl1' imin . Stream.indexed . stream
+  where
+    imin (i,x) (j,y) | x <= y    = (i,x)
+                     | otherwise = (j,y)
+
+maxIndex :: (Vector v a, Ord a) => v a -> Int
+{-# INLINE maxIndex #-}
+maxIndex = fst . Stream.foldl1' imax . Stream.indexed . stream
+  where
+    imax (i,x) (j,y) | x >= y    = (i,x)
+                     | otherwise = (j,y)
+
+
 -- Unfolding
 -- ---------
 
index f4939b6..64ea276 100644 (file)
@@ -32,22 +32,24 @@ module Data.Vector.Primitive (
   accum, accumulate_, (//), update_, backpermute, reverse,
 
   -- * Mapping
-  map, concatMap,
+  map, imap, concatMap,
 
   -- * Zipping and unzipping
-  zipWith, zipWith3,
+  zipWith, zipWith3, zipWith4, zipWith5, zipWith6,
+  izipWith, izipWith3, izipWith4, izipWith5, izipWith6,
 
   -- * Filtering
-  filter, takeWhile, dropWhile,
+  filter, ifilter, takeWhile, dropWhile,
 
   -- * Searching
   elem, notElem, find, findIndex,
 
   -- * Folding
   foldl, foldl1, foldl', foldl1', foldr, foldr1,
+  ifoldl, ifoldl', ifoldr,
 
   -- * Specialised folds
-  sum, product, maximum, minimum,
+  sum, product, maximum, minimum, minIndex, maxIndex,
 
   -- * Unfolding
   unfoldr,
@@ -291,6 +293,11 @@ map :: (Prim a, Prim b) => (a -> b) -> Vector a -> Vector b
 {-# INLINE map #-}
 map = G.map
 
+-- | Apply a function to every index/value pair
+imap :: (Prim a, Prim b) => (Int -> a -> b) -> Vector a -> Vector b
+{-# INLINE imap #-}
+imap = G.imap
+
 concatMap :: (Prim a, Prim b) => (a -> Vector b) -> Vector a -> Vector b
 {-# INLINE concatMap #-}
 concatMap = G.concatMap
@@ -310,6 +317,59 @@ zipWith3 :: (Prim a, Prim b, Prim c, Prim d)
 {-# INLINE zipWith3 #-}
 zipWith3 = G.zipWith3
 
+zipWith4 :: (Prim a, Prim b, Prim c, Prim d, Prim e)
+         => (a -> b -> c -> d -> e)
+         -> Vector a -> Vector b -> Vector c -> Vector d -> Vector e
+{-# INLINE zipWith4 #-}
+zipWith4 = G.zipWith4
+
+zipWith5 :: (Prim a, Prim b, Prim c, Prim d, Prim e, Prim f)
+         => (a -> b -> c -> d -> e -> f)
+         -> Vector a -> Vector b -> Vector c -> Vector d -> Vector e
+         -> Vector f
+{-# INLINE zipWith5 #-}
+zipWith5 = G.zipWith5
+
+zipWith6 :: (Prim a, Prim b, Prim c, Prim d, Prim e, Prim f, Prim g)
+         => (a -> b -> c -> d -> e -> f -> g)
+         -> Vector a -> Vector b -> Vector c -> Vector d -> Vector e
+         -> Vector f -> Vector g
+{-# INLINE zipWith6 #-}
+zipWith6 = G.zipWith6
+
+-- | Zip two vectors and their indices with the given function.
+izipWith :: (Prim a, Prim b, Prim c)
+         => (Int -> a -> b -> c) -> Vector a -> Vector b -> Vector c
+{-# INLINE izipWith #-}
+izipWith = G.izipWith
+
+-- | Zip three vectors and their indices with the given function.
+izipWith3 :: (Prim a, Prim b, Prim c, Prim d)
+          => (Int -> a -> b -> c -> d)
+          -> Vector a -> Vector b -> Vector c -> Vector d
+{-# INLINE izipWith3 #-}
+izipWith3 = G.izipWith3
+
+izipWith4 :: (Prim a, Prim b, Prim c, Prim d, Prim e)
+          => (Int -> a -> b -> c -> d -> e)
+          -> Vector a -> Vector b -> Vector c -> Vector d -> Vector e
+{-# INLINE izipWith4 #-}
+izipWith4 = G.izipWith4
+
+izipWith5 :: (Prim a, Prim b, Prim c, Prim d, Prim e, Prim f)
+          => (Int -> a -> b -> c -> d -> e -> f)
+          -> Vector a -> Vector b -> Vector c -> Vector d -> Vector e
+          -> Vector f
+{-# INLINE izipWith5 #-}
+izipWith5 = G.izipWith5
+
+izipWith6 :: (Prim a, Prim b, Prim c, Prim d, Prim e, Prim f, Prim g)
+          => (Int -> a -> b -> c -> d -> e -> f -> g)
+          -> Vector a -> Vector b -> Vector c -> Vector d -> Vector e
+          -> Vector f -> Vector g
+{-# INLINE izipWith6 #-}
+izipWith6 = G.izipWith6
+
 -- Filtering
 -- ---------
 
@@ -318,6 +378,12 @@ filter :: Prim a => (a -> Bool) -> Vector a -> Vector a
 {-# INLINE filter #-}
 filter = G.filter
 
+-- | Drop elements that do not satisfy the predicate (applied to values and
+-- their indices)
+ifilter :: Prim a => (Int -> a -> Bool) -> Vector a -> Vector a
+{-# INLINE ifilter #-}
+ifilter = G.ifilter
+
 -- | Yield the longest prefix of elements satisfying the predicate.
 takeWhile :: Prim a => (a -> Bool) -> Vector a -> Vector a
 {-# INLINE takeWhile #-}
@@ -388,6 +454,22 @@ foldr1 :: Prim a => (a -> a -> a) -> Vector a -> a
 {-# INLINE foldr1 #-}
 foldr1 = G.foldr1
 
+-- | Left fold (function applied to each element and its index)
+ifoldl :: Prim b => (a -> Int -> b -> a) -> a -> Vector b -> a
+{-# INLINE ifoldl #-}
+ifoldl = G.ifoldl
+
+-- | Left fold with strict accumulator (function applied to each element and
+-- its index)
+ifoldl' :: Prim b => (a -> Int -> b -> a) -> a -> Vector b -> a
+{-# INLINE ifoldl' #-}
+ifoldl' = G.ifoldl'
+
+-- | Right fold (function applied to each element and its index)
+ifoldr :: Prim a => (Int -> a -> b -> b) -> b -> Vector a -> b
+{-# INLINE ifoldr #-}
+ifoldr = G.ifoldr
+
 -- Specialised folds
 -- -----------------
 
@@ -407,6 +489,14 @@ minimum :: (Prim a, Ord a) => Vector a -> a
 {-# INLINE minimum #-}
 minimum = G.minimum
 
+minIndex :: (Prim a, Ord a) => Vector a -> Int
+{-# INLINE minIndex #-}
+minIndex = G.minIndex
+
+maxIndex :: (Prim a, Ord a) => Vector a -> Int
+{-# INLINE maxIndex #-}
+maxIndex = G.maxIndex
+
 -- Unfolding
 -- ---------
 
index a242e7d..66e44c6 100644 (file)
@@ -32,19 +32,21 @@ module Data.Vector.Storable (
   accum, accumulate_, (//), update_, backpermute, reverse,
 
   -- * Mapping
-  map, concatMap,
+  map, imap, concatMap,
 
   -- * Zipping and unzipping
-  zipWith, zipWith3,
+  zipWith, zipWith3, zipWith4, zipWith5, zipWith6,
+  izipWith, izipWith3, izipWith4, izipWith5, izipWith6,
 
   -- * Filtering
-  filter, takeWhile, dropWhile,
+  filter, ifilter, takeWhile, dropWhile,
 
   -- * Searching
-  elem, notElem, find, findIndex,
+  elem, notElem, find, findIndex, minIndex, maxIndex,
 
   -- * Folding
   foldl, foldl1, foldl', foldl1', foldr, foldr1,
+  ifoldl, ifoldl', ifoldr,
 
   -- * Specialised folds
   and, or, sum, product, maximum, minimum,
@@ -322,6 +324,11 @@ map :: (Storable a, Storable b) => (a -> b) -> Vector a -> Vector b
 {-# INLINE map #-}
 map = G.map
 
+-- | Apply a function to every index/value pair
+imap :: (Storable a, Storable b) => (Int -> a -> b) -> Vector a -> Vector b
+{-# INLINE imap #-}
+imap = G.imap
+
 concatMap :: (Storable a, Storable b) => (a -> Vector b) -> Vector a -> Vector b
 {-# INLINE concatMap #-}
 concatMap = G.concatMap
@@ -341,6 +348,63 @@ zipWith3 :: (Storable a, Storable b, Storable c, Storable d)
 {-# INLINE zipWith3 #-}
 zipWith3 = G.zipWith3
 
+zipWith4 :: (Storable a, Storable b, Storable c, Storable d, Storable e)
+         => (a -> b -> c -> d -> e)
+         -> Vector a -> Vector b -> Vector c -> Vector d -> Vector e
+{-# INLINE zipWith4 #-}
+zipWith4 = G.zipWith4
+
+zipWith5 :: (Storable a, Storable b, Storable c, Storable d, Storable e,
+             Storable f)
+         => (a -> b -> c -> d -> e -> f)
+         -> Vector a -> Vector b -> Vector c -> Vector d -> Vector e
+         -> Vector f
+{-# INLINE zipWith5 #-}
+zipWith5 = G.zipWith5
+
+zipWith6 :: (Storable a, Storable b, Storable c, Storable d, Storable e,
+             Storable f, Storable g)
+         => (a -> b -> c -> d -> e -> f -> g)
+         -> Vector a -> Vector b -> Vector c -> Vector d -> Vector e
+         -> Vector f -> Vector g
+{-# INLINE zipWith6 #-}
+zipWith6 = G.zipWith6
+
+-- | Zip two vectors and their indices with the given function.
+izipWith :: (Storable a, Storable b, Storable c)
+         => (Int -> a -> b -> c) -> Vector a -> Vector b -> Vector c
+{-# INLINE izipWith #-}
+izipWith = G.izipWith
+
+-- | Zip three vectors and their indices with the given function.
+izipWith3 :: (Storable a, Storable b, Storable c, Storable d)
+          => (Int -> a -> b -> c -> d)
+          -> Vector a -> Vector b -> Vector c -> Vector d
+{-# INLINE izipWith3 #-}
+izipWith3 = G.izipWith3
+
+izipWith4 :: (Storable a, Storable b, Storable c, Storable d, Storable e)
+          => (Int -> a -> b -> c -> d -> e)
+          -> Vector a -> Vector b -> Vector c -> Vector d -> Vector e
+{-# INLINE izipWith4 #-}
+izipWith4 = G.izipWith4
+
+izipWith5 :: (Storable a, Storable b, Storable c, Storable d, Storable e,
+              Storable f)
+          => (Int -> a -> b -> c -> d -> e -> f)
+          -> Vector a -> Vector b -> Vector c -> Vector d -> Vector e
+          -> Vector f
+{-# INLINE izipWith5 #-}
+izipWith5 = G.izipWith5
+
+izipWith6 :: (Storable a, Storable b, Storable c, Storable d, Storable e,
+              Storable f, Storable g)
+          => (Int -> a -> b -> c -> d -> e -> f -> g)
+          -> Vector a -> Vector b -> Vector c -> Vector d -> Vector e
+          -> Vector f -> Vector g
+{-# INLINE izipWith6 #-}
+izipWith6 = G.izipWith6
+
 -- Filtering
 -- ---------
 
@@ -349,6 +413,12 @@ filter :: Storable a => (a -> Bool) -> Vector a -> Vector a
 {-# INLINE filter #-}
 filter = G.filter
 
+-- | Drop elements that do not satisfy the predicate (applied to values and
+-- their indices)
+ifilter :: Storable a => (Int -> a -> Bool) -> Vector a -> Vector a
+{-# INLINE ifilter #-}
+ifilter = G.ifilter
+
 -- | Yield the longest prefix of elements satisfying the predicate.
 takeWhile :: Storable a => (a -> Bool) -> Vector a -> Vector a
 {-# INLINE takeWhile #-}
@@ -419,6 +489,22 @@ foldr1 :: Storable a => (a -> a -> a) -> Vector a -> a
 {-# INLINE foldr1 #-}
 foldr1 = G.foldr1
 
+-- | Left fold (function applied to each element and its index)
+ifoldl :: Storable b => (a -> Int -> b -> a) -> a -> Vector b -> a
+{-# INLINE ifoldl #-}
+ifoldl = G.ifoldl
+
+-- | Left fold with strict accumulator (function applied to each element and
+-- its index)
+ifoldl' :: Storable b => (a -> Int -> b -> a) -> a -> Vector b -> a
+{-# INLINE ifoldl' #-}
+ifoldl' = G.ifoldl'
+
+-- | Right fold (function applied to each element and its index)
+ifoldr :: Storable a => (Int -> a -> b -> b) -> b -> Vector a -> b
+{-# INLINE ifoldr #-}
+ifoldr = G.ifoldr
+
 -- Specialised folds
 -- -----------------
 
@@ -446,6 +532,14 @@ minimum :: (Storable a, Ord a) => Vector a -> a
 {-# INLINE minimum #-}
 minimum = G.minimum
 
+minIndex :: (Storable a, Ord a) => Vector a -> Int
+{-# INLINE minIndex #-}
+minIndex = G.minIndex
+
+maxIndex :: (Storable a, Ord a) => Vector a -> Int
+{-# INLINE maxIndex #-}
+maxIndex = G.maxIndex
+
 -- Unfolding
 -- ---------
 
index 125e5bf..090b554 100644 (file)
@@ -32,24 +32,26 @@ module Data.Vector.Unboxed (
   backpermute, reverse,
 
   -- * Mapping
-  map, concatMap,
+  map, imap, concatMap,
 
   -- * Zipping and unzipping
   zipWith, zipWith3, zipWith4, zipWith5, zipWith6,
+  izipWith, izipWith3, izipWith4, izipWith5, izipWith6,
   zip, zip3, zip4, zip5, zip6,
   unzip, unzip3, unzip4, unzip5, unzip6,
 
   -- * Filtering
-  filter, takeWhile, dropWhile,
+  filter, ifilter, takeWhile, dropWhile,
 
   -- * Searching
   elem, notElem, find, findIndex,
 
   -- * Folding
   foldl, foldl1, foldl', foldl1', foldr, foldr1,
+  ifoldl, ifoldl', ifoldr,
 
   -- * Specialised folds
-  and, or, sum, product, maximum, minimum,
+  and, or, sum, product, maximum, minimum, minIndex, maxIndex,
 
   -- * Unfolding
   unfoldr,
@@ -281,6 +283,11 @@ map :: (Unbox a, Unbox b) => (a -> b) -> Vector a -> Vector b
 {-# INLINE map #-}
 map = G.map
 
+-- | Apply a function to every index/value pair
+imap :: (Unbox a, Unbox b) => (Int -> a -> b) -> Vector a -> Vector b
+{-# INLINE imap #-}
+imap = G.imap
+
 concatMap :: (Unbox a, Unbox b) => (a -> Vector b) -> Vector a -> Vector b
 {-# INLINE concatMap #-}
 concatMap = G.concatMap
@@ -320,6 +327,39 @@ zipWith6 :: (Unbox a, Unbox b, Unbox c, Unbox d, Unbox e, Unbox f, Unbox g)
 {-# INLINE zipWith6 #-}
 zipWith6 = G.zipWith6
 
+-- | Zip two vectors and their indices with the given function.
+izipWith :: (Unbox a, Unbox b, Unbox c)
+         => (Int -> a -> b -> c) -> Vector a -> Vector b -> Vector c
+{-# INLINE izipWith #-}
+izipWith = G.izipWith
+
+-- | Zip three vectors and their indices with the given function.
+izipWith3 :: (Unbox a, Unbox b, Unbox c, Unbox d)
+          => (Int -> a -> b -> c -> d)
+          -> Vector a -> Vector b -> Vector c -> Vector d
+{-# INLINE izipWith3 #-}
+izipWith3 = G.izipWith3
+
+izipWith4 :: (Unbox a, Unbox b, Unbox c, Unbox d, Unbox e)
+          => (Int -> a -> b -> c -> d -> e)
+          -> Vector a -> Vector b -> Vector c -> Vector d -> Vector e
+{-# INLINE izipWith4 #-}
+izipWith4 = G.izipWith4
+
+izipWith5 :: (Unbox a, Unbox b, Unbox c, Unbox d, Unbox e, Unbox f)
+          => (Int -> a -> b -> c -> d -> e -> f)
+          -> Vector a -> Vector b -> Vector c -> Vector d -> Vector e
+          -> Vector f
+{-# INLINE izipWith5 #-}
+izipWith5 = G.izipWith5
+
+izipWith6 :: (Unbox a, Unbox b, Unbox c, Unbox d, Unbox e, Unbox f, Unbox g)
+          => (Int -> a -> b -> c -> d -> e -> f -> g)
+          -> Vector a -> Vector b -> Vector c -> Vector d -> Vector e
+          -> Vector f -> Vector g
+{-# INLINE izipWith6 #-}
+izipWith6 = G.izipWith6
+
 -- Filtering
 -- ---------
 
@@ -328,6 +368,12 @@ filter :: Unbox a => (a -> Bool) -> Vector a -> Vector a
 {-# INLINE filter #-}
 filter = G.filter
 
+-- | Drop elements that do not satisfy the predicate (applied to values and
+-- their indices)
+ifilter :: Unbox a => (Int -> a -> Bool) -> Vector a -> Vector a
+{-# INLINE ifilter #-}
+ifilter = G.ifilter
+
 -- | Yield the longest prefix of elements satisfying the predicate.
 takeWhile :: Unbox a => (a -> Bool) -> Vector a -> Vector a
 {-# INLINE takeWhile #-}
@@ -398,6 +444,22 @@ foldr1 :: Unbox a => (a -> a -> a) -> Vector a -> a
 {-# INLINE foldr1 #-}
 foldr1 = G.foldr1
 
+-- | Left fold (function applied to each element and its index)
+ifoldl :: Unbox b => (a -> Int -> b -> a) -> a -> Vector b -> a
+{-# INLINE ifoldl #-}
+ifoldl = G.ifoldl
+
+-- | Left fold with strict accumulator (function applied to each element and
+-- its index)
+ifoldl' :: Unbox b => (a -> Int -> b -> a) -> a -> Vector b -> a
+{-# INLINE ifoldl' #-}
+ifoldl' = G.ifoldl'
+
+-- | Right fold (function applied to each element and its index)
+ifoldr :: Unbox a => (Int -> a -> b -> b) -> b -> Vector a -> b
+{-# INLINE ifoldr #-}
+ifoldr = G.ifoldr
+
 -- Specialised folds
 -- -----------------
 
@@ -425,6 +487,14 @@ minimum :: (Unbox a, Ord a) => Vector a -> a
 {-# INLINE minimum #-}
 minimum = G.minimum
 
+minIndex :: (Unbox a, Ord a) => Vector a -> Int
+{-# INLINE minIndex #-}
+minIndex = G.minIndex
+
+maxIndex :: (Unbox a, Ord a) => Vector a -> Int
+{-# INLINE maxIndex #-}
+maxIndex = G.maxIndex
+
 -- Unfolding
 -- ---------