add withIndex traversal

3 files changed, 79 insertions, 4 deletions

diff --git a/examples/vector-map.hs b/examples/vector-map.hs
index f116946..7796cc0 100644
--- a/examples/vector-map.hs
+++ b/examples/vector-map.hs
@@ -30,12 +30,45 @@ addInitialM = mapVectorM_ (\x -> do
                                  put $ x + i
                           )
 
+-- sum the values of the even indiced elements
+sumEvens :: Vector Int -> Int
+sumEvens = foldVectorWithIndex (\x a b -> if x `mod` 2 == 0 then a + b else b) 0 
+
+-- sum and print running total of evens
+sumEvensAndPrint :: Vector Int -> VectorMonadT ()
+sumEvensAndPrint = mapVectorWithIndexM_ (\ i x -> do
+                                                  when (i `mod` 2 == 0) (do
+                                                                         v <- get
+                                                                         put $ v + x
+                                                                         v' <- get
+                                                                         liftIO $ putStr $ (show v') ++ " " 
+                                                                         return ())
+                                                  return ()
+                                        )
+
+indexPlusSum :: Vector Int -> VectorMonadT ()
+indexPlusSum v' = do
+                  v <- mapVectorWithIndexM (\i x -> do
+                                                    s <- get
+                                                    let inc = x+s
+                                                    liftIO $ putStr $ show (i,inc) ++ " "
+                                                    put inc
+                                                    return inc) v'
+                  liftIO $ do
+                           putStrLn ""
+                           putStrLn $ show v
+
 -------------------------------------------
+
 main = do
        v' <- test1 v
        putStrLn ""
        putStrLn $ show v'
        evalStateT (addInitialM v) 0
        putStrLn ""
+       putStrLn $ show (sumEvens v)
+       evalStateT (sumEvensAndPrint v) 0
+       putStrLn ""
+       evalStateT (indexPlusSum v) 0
        return ()
 
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