fix mapVectorM(_) and add example

3 files changed, 63 insertions, 18 deletions

diff --git a/examples/vector-map.hs b/examples/vector-map.hs
new file mode 100644
index 0000000..d778358
--- /dev/null
+++ b/examples/vector-map.hs
@@ -0,0 +1,42 @@
+-- use of vectorMapM
+--
+
+-------------------------------------------
+
+import Data.Packed.Vector
+import Numeric.LinearAlgebra.Interface
+
+import Control.Monad.State
+import Control.Monad.Trans
+
+-------------------------------------------
+
+-- an instance of MonadIO, a monad transformer
+type VectorMonadT = StateT Int IO
+
+v :: Vector Int
+v = fromList $ take 10 [0..]
+
+test1 :: Vector Int -> IO (Vector Int)
+test1 = do
+        mapVectorM (\x -> do
+                          putStr $ (show) x ++ " "
+                          return (x + 1))
+                            
+-- we can have an arbitrary monad AND do IO
+addInitialM :: Vector Int -> VectorMonadT ()
+addInitialM = mapVectorM_ (\x -> do
+                                 i <- get
+                                 liftIO $ putStr $ (show $ x + i) ++ " "
+                                 put $ x + i
+                          )
+
+-------------------------------------------
+main = do
+       v' <- test1 v
+       putStrLn ""
+       putStrLn $ show v'
+       evalStateT (addInitialM v) 1
+       putStrLn ""
+       return ()
+
diff --git a/hmatrix.cabal b/hmatrix.cabal
index 04c3581..0bcc3aa 100644
--- a/hmatrix.cabal
+++ b/hmatrix.cabal
@@ -45,6 +45,7 @@ extra-source-files: examples/tests.hs
                     examples/devel/ej2/functions.c
                     examples/Real.hs
                     examples/vector.hs
+                    examples/vector-map.hs
 
 extra-source-files: lib/Numeric/LinearAlgebra/LAPACK/lapack-aux.h,
                     lib/Numeric/LinearAlgebra/LAPACK/clapack.h
diff --git a/lib/Data/Packed/Internal/Vector.hs b/lib/Data/Packed/Internal/Vector.hs
index 7dd1289..ec2bf3c 100644
--- a/lib/Data/Packed/Internal/Vector.hs
+++ b/lib/Data/Packed/Internal/Vector.hs
@@ -365,30 +365,32 @@ foldVectorG f s0 v = foldLoop g s0 (dim v)
 mapVectorM :: (Storable a, Storable b, MonadIO m) => (a -> m b) -> Vector a -> m (Vector b)
 mapVectorM f v = do
     w <- liftIO $ createVector (dim v)
-    mapVectorM' f v w (dim v -1)
+    mapVectorM' f v w 0 (dim v -1)
     return w
-    where mapVectorM' f' v' w' 0  = do
-                                    x <- liftIO $ unsafeWith v' $ \p -> peekElemOff p 0 
-                                    y <- f' x
-                                    liftIO $ unsafeWith w' $ \q -> pokeElemOff q 0 y
-          mapVectorM' f' v' w' !k = do
-                                    x <- liftIO $ unsafeWith v' $ \p -> peekElemOff p k 
-                                    y <- f' x
-                                    liftIO $ unsafeWith w' $ \q -> pokeElemOff q k y
-                                    mapVectorM' f' v' w' (k-1)
+    where mapVectorM' f' v' w' !k !t
+              | k == t               = do
+                                       x <- liftIO $ unsafeWith v' $ \p -> peekElemOff p k 
+                                       y <- f' x
+                                       liftIO $ unsafeWith w' $ \q -> pokeElemOff q k y
+              | otherwise            = do
+                                       x <- liftIO $ unsafeWith v' $ \p -> peekElemOff p k 
+                                       y <- f' x
+                                       liftIO $ unsafeWith w' $ \q -> pokeElemOff q k y
+                                       mapVectorM' f' v' w' (k+1) t
 {-# INLINE mapVectorM #-}
 
 -- | monadic map over Vectors
 mapVectorM_ :: (Storable a, MonadIO m) => (a -> m ()) -> Vector a -> m ()
 mapVectorM_ f v = do
-    mapVectorM' f v (dim v -1)
-    where mapVectorM' f' v' 0  = do
-                                 x <- liftIO $ unsafeWith v' $ \p -> peekElemOff p 0
-                                 f' x
-          mapVectorM' f' v' !k = do
-                                 x <- liftIO $ unsafeWith v' $ \p -> peekElemOff p k 
-                                 _ <- f' x
-                                 mapVectorM' f' v' (k-1)
+    mapVectorM' f v 0 (dim v -1)
+    where mapVectorM' f' v' !k !t
+              | k == t            = do
+                                    x <- liftIO $ unsafeWith v' $ \p -> peekElemOff p k
+                                    f' x
+              | otherwise         = do
+                                    x <- liftIO $ unsafeWith v' $ \p -> peekElemOff p k 
+                                    _ <- f' x
+                                    mapVectorM' f' v' (k+1) t
 {-# INLINE mapVectorM_ #-}
 
 -------------------------------------------------------------------