add withIndex traversal

2 files changed, 46 insertions, 4 deletions

diff --git a/lib/Data/Packed/Internal/Vector.hs b/lib/Data/Packed/Internal/Vector.hs
index be2fcbb..a47c376 100644
--- a/lib/Data/Packed/Internal/Vector.hs
+++ b/lib/Data/Packed/Internal/Vector.hs
@@ -18,8 +18,8 @@ module Data.Packed.Internal.Vector (
     fromList, toList, (|>),
     join, (@>), safe, at, at', subVector, takesV,
     mapVector, zipVectorWith, unzipVectorWith,
-    mapVectorM, mapVectorM_,
-    foldVector, foldVectorG, foldLoop,
+    mapVectorM, mapVectorM_, mapVectorWithIndexM, mapVectorWithIndexM_,
+    foldVector, foldVectorG, foldLoop, foldVectorWithIndex,
     createVector, vec,
     asComplex, asReal, float2DoubleV, double2FloatV,
     fwriteVector, freadVector, fprintfVector, fscanfVector,
@@ -364,6 +364,16 @@ foldVector f x v = unsafePerformIO $
         go (dim v -1) x
 {-# INLINE foldVector #-}
 
+-- the zero-indexed index is passed to the folding function
+foldVectorWithIndex :: Storable a => (Int -> a -> b -> b) -> b -> Vector a -> b
+foldVectorWithIndex f x v = unsafePerformIO $
+    unsafeWith v $ \p -> do
+        let go (-1) s = return s
+            go !k !s = do y <- peekElemOff p k
+                          go (k-1::Int) (f k y s)
+        go (dim v -1) x
+{-# INLINE foldVectorWithIndex #-}
+
 foldLoop f s0 d = go (d - 1) s0
      where
        go 0 s = f (0::Int) s
@@ -408,6 +418,38 @@ mapVectorM_ f v = do
                                     mapVectorM' f' v' (k+1) t
 {-# INLINE mapVectorM_ #-}
 
+-- | monadic map over Vectors with the zero-indexed index passed to the mapping function
+mapVectorWithIndexM :: (Storable a, Storable b, Monad m) => (Int -> a -> m b) -> Vector a -> m (Vector b)
+mapVectorWithIndexM f v = do
+    w <- return $! unsafePerformIO $! createVector (dim v)
+    mapVectorM' f v w 0 (dim v -1)
+    return w
+    where mapVectorM' f' v' w' !k !t
+              | k == t               = do
+                                       x <- return $! inlinePerformIO $! unsafeWith v' $! \p -> peekElemOff p k 
+                                       y <- f' k x
+                                       return $! inlinePerformIO $! unsafeWith w' $! \q -> pokeElemOff q k y
+              | otherwise            = do
+                                       x <- return $! inlinePerformIO $! unsafeWith v' $! \p -> peekElemOff p k 
+                                       y <- f' k x
+                                       _ <- return $! inlinePerformIO $! unsafeWith w' $! \q -> pokeElemOff q k y
+                                       mapVectorM' f' v' w' (k+1) t
+{-# INLINE mapVectorWithIndexM #-}
+
+-- | monadic map over Vectors with the zero-indexed index passed to the mapping function
+mapVectorWithIndexM_ :: (Storable a, Monad m) => (Int -> a -> m ()) -> Vector a -> m ()
+mapVectorWithIndexM_ f v = do
+    mapVectorM' f v 0 (dim v -1)
+    where mapVectorM' f' v' !k !t
+              | k == t            = do
+                                    x <- return $! inlinePerformIO $! unsafeWith v' $! \p -> peekElemOff p k
+                                    f' k x
+              | otherwise         = do
+                                    x <- return $! inlinePerformIO $! unsafeWith v' $! \p -> peekElemOff p k 
+                                    _ <- f' k x
+                                    mapVectorM' f' v' (k+1) t
+{-# INLINE mapVectorWithIndexM_ #-}
+
 -------------------------------------------------------------------
 
 
diff --git a/lib/Data/Packed/Vector.hs b/lib/Data/Packed/Vector.hs
index a526caa..ad690f9 100644
--- a/lib/Data/Packed/Vector.hs
+++ b/lib/Data/Packed/Vector.hs
@@ -27,9 +27,9 @@ module Data.Packed.Vector (
 --    vectorMax, vectorMin,
     vectorMaxIndex, vectorMinIndex,
     mapVector, zipVector, zipVectorWith, unzipVector, unzipVectorWith,
-    mapVectorM, mapVectorM_,
+    mapVectorM, mapVectorM_, mapVectorWithIndexM, mapVectorWithIndexM_,
     fscanfVector, fprintfVector, freadVector, fwriteVector,
-    foldLoop, foldVector, foldVectorG
+    foldLoop, foldVector, foldVectorG, foldVectorWithIndex
 ) where
 
 import Data.Packed.Internal.Vector