pure mapVectorWithIndex

2 files changed, 17 insertions, 2 deletions

diff --git a/lib/Data/Packed/Internal/Vector.hs b/lib/Data/Packed/Internal/Vector.hs
index 078cb36..8f403f4 100644
--- a/lib/Data/Packed/Internal/Vector.hs
+++ b/lib/Data/Packed/Internal/Vector.hs
@@ -17,7 +17,7 @@ module Data.Packed.Internal.Vector (
     Vector, dim,
     fromList, toList, (|>),
     join, (@>), safe, at, at', subVector, takesV,
-    mapVector, zipVectorWith, unzipVectorWith,
+    mapVector, mapVectorWithIndex, zipVectorWith, unzipVectorWith,
     mapVectorM, mapVectorM_, mapVectorWithIndexM, mapVectorWithIndexM_,
     foldVector, foldVectorG, foldLoop, foldVectorWithIndex,
     createVector, vec,
@@ -497,6 +497,21 @@ mapVectorWithIndexM_ f v = do
                                     mapVectorM' (k+1) t
 {-# INLINE mapVectorWithIndexM_ #-}
 
+
+mapVectorWithIndex :: (Storable a, Storable b) => (Int -> a -> b) -> Vector a -> Vector b
+--mapVectorWithIndex g = head . mapVectorWithIndexM (\a b -> [g a b])
+mapVectorWithIndex f v = unsafePerformIO $ do
+    w <- createVector (dim v)
+    unsafeWith v $ \p ->
+        unsafeWith w $ \q -> do
+            let go (-1) = return ()
+                go !k = do x <- peekElemOff p k
+                           pokeElemOff      q k (f k x)
+                           go (k-1)
+            go (dim v -1)
+    return w
+{-# INLINE mapVectorWithIndex #-}
+
 -------------------------------------------------------------------
 
 
diff --git a/lib/Data/Packed/Vector.hs b/lib/Data/Packed/Vector.hs
index f90f8e4..dad5d28 100644
--- a/lib/Data/Packed/Vector.hs
+++ b/lib/Data/Packed/Vector.hs
@@ -20,7 +20,7 @@ module Data.Packed.Vector (
     fromList, (|>), toList, buildVector,
     dim, (@>),
     subVector, takesV, join,
-    mapVector, zipVector, zipVectorWith, unzipVector, unzipVectorWith,
+    mapVector, mapVectorWithIndex, zipVector, zipVectorWith, unzipVector, unzipVectorWith,
     mapVectorM, mapVectorM_, mapVectorWithIndexM, mapVectorWithIndexM_,
     foldLoop, foldVector, foldVectorG, foldVectorWithIndex
 ) where