1 files changed, 15 insertions, 13 deletions

diff --git a/lib/Data/Packed/Matrix.hs b/lib/Data/Packed/Matrix.hs
index c1a9b24..fe8c159 100644
--- a/lib/Data/Packed/Matrix.hs
+++ b/lib/Data/Packed/Matrix.hs
@@ -46,6 +46,7 @@ import Data.Array
 
 import Data.List(transpose,intersperse)
 import Foreign.Storable(Storable)
+import Control.Monad(liftM)
 
 -------------------------------------------------------------------
 
@@ -351,7 +352,11 @@ toBlocksEvery r c m = toBlocks rs cs m where
 
 -------------------------------------------------------------------
 
-mk c g = \k v -> g (divMod k c) v
+-- Given a column number and a function taking matrix indexes, returns
+-- a function which takes vector indexes (that can be used on the
+-- flattened matrix).
+mk :: Int -> ((Int, Int) -> t) -> (Int -> t)
+mk c g = \k -> g (divMod k c)
 
 {- |
 
@@ -364,9 +369,8 @@ m[1,1] = 5
 m[1,2] = 6@
 -}
 mapMatrixWithIndexM_
-  :: (Element a, Num a,
-      Functor f, Monad f) =>
-      ((Int, Int) -> a -> f ()) -> Matrix a -> f ()
+  :: (Element a, Num a, Monad m) =>
+      ((Int, Int) -> a -> m ()) -> Matrix a -> m ()
 mapMatrixWithIndexM_ g m = mapVectorWithIndexM_ (mk c g) . flatten $ m
   where
     c = cols m
@@ -380,11 +384,9 @@ Just (3><3)
  ,  20.0,  21.0, 122.0 ]@
 -}
 mapMatrixWithIndexM
-  :: (Foreign.Storable.Storable t,
-      Element a, Num a,
-      Functor f, Monad f) =>
-      ((Int, Int) -> a -> f t) -> Matrix a -> f (Matrix t)
-mapMatrixWithIndexM g m = fmap (reshape c) . mapVectorWithIndexM (mk c g) . flatten $ m
+  :: (Element a, Storable b, Monad m) =>
+      ((Int, Int) -> a -> m b) -> Matrix a -> m (Matrix b)
+mapMatrixWithIndexM g m = liftM (reshape c) . mapVectorWithIndexM (mk c g) . flatten $ m 
     where
       c = cols m
 
@@ -395,10 +397,10 @@ mapMatrixWithIndexM g m = fmap (reshape c) . mapVectorWithIndexM (mk c g) . flat
  ,  10.0, 111.0,  12.0
  ,  20.0,  21.0, 122.0 ]@
  -}
-mapMatrixWithIndex :: (Foreign.Storable.Storable t,
-      Element a, Num a) =>
-      ((Int, Int) -> a -> t) -> Matrix a -> Matrix t
-mapMatrixWithIndex g m = reshape c $ mapVectorWithIndex (mk c g) $ flatten m
+mapMatrixWithIndex
+  :: (Element a, Storable b) =>
+      ((Int, Int) -> a -> b) -> Matrix a -> Matrix b
+mapMatrixWithIndex g m = reshape c . mapVectorWithIndex (mk c g) . flatten $ m
     where
       c = cols m