From 2cb2293eeb617baa404f444944bb4613c645133a Mon Sep 17 00:00:00 2001
From: Alberto Ruiz <aruiz@um.es>
Date: Sun, 10 Jul 2011 17:31:24 +0000
Subject: pure mapVectorWithIndex

---
 lib/Data/Packed/Internal/Vector.hs | 17 ++++++++++++++++-
 lib/Data/Packed/Vector.hs          |  2 +-
 2 files changed, 17 insertions(+), 2 deletions(-)

(limited to 'lib')

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
-- 
cgit v1.2.3