Fusion rules for in-place map
authorRoman Leshchinskiy <rl@cse.unsw.edu.au>
Sat, 12 Jul 2008 09:52:02 +0000 (09:52 +0000)
committerRoman Leshchinskiy <rl@cse.unsw.edu.au>
Sat, 12 Jul 2008 09:52:02 +0000 (09:52 +0000)
Data/Vector/IVector.hs
Data/Vector/MVector.hs
Data/Vector/MVector/Mut.hs

index 50eebbf..4fe9910 100644 (file)
@@ -296,6 +296,14 @@ map :: (IVector v a, IVector v b) => (a -> b) -> v a -> v b
 {-# INLINE map #-}
 map f = unstream . Stream.map f . stream
 
 {-# INLINE map #-}
 map f = unstream . Stream.map f . stream
 
+{-# RULES
+
+"in-place map [IVector]" forall f m.
+  Mut.unstream (Stream.map f (stream (new m)))
+    = Mut.map f m
+
+  #-}
+
 -- | Zip two vectors with the given function.
 zipWith :: (IVector v a, IVector v b, IVector v c) => (a -> b -> c) -> v a -> v b -> v c
 {-# INLINE zipWith #-}
 -- | Zip two vectors with the given function.
 zipWith :: (IVector v a, IVector v b, IVector v c) => (a -> b -> c) -> v a -> v b -> v c
 {-# INLINE zipWith #-}
index 38f0e8d..e4ed39a 100644 (file)
@@ -16,7 +16,7 @@
 module Data.Vector.MVector (
   MVectorPure(..), MVector(..),
 
 module Data.Vector.MVector (
   MVectorPure(..), MVector(..),
 
-  slice, new, newWith, read, write, copy, grow, unstream, update
+  slice, new, newWith, read, write, copy, grow, unstream, map, update
 ) where
 
 import qualified Data.Vector.Stream      as Stream
 ) where
 
 import qualified Data.Vector.Stream      as Stream
@@ -30,7 +30,7 @@ import GHC.Float (
     double2Int, int2Double
   )
 
     double2Int, int2Double
   )
 
-import Prelude hiding ( length, read )
+import Prelude hiding ( length, map, read )
 
 gROWTH_FACTOR :: Double
 gROWTH_FACTOR = 1.5
 
 gROWTH_FACTOR :: Double
 gROWTH_FACTOR = 1.5
@@ -201,6 +201,16 @@ unstreamUnknown s
                                  . double2Int
                                  $ int2Double (length v) * gROWTH_FACTOR
 
                                  . double2Int
                                  $ int2Double (length v) * gROWTH_FACTOR
 
+map :: MVector v m a => (a -> a) -> v a -> m ()
+{-# INLINE map #-}
+map f v = map_loop 0
+  where
+    n = length v
+
+    map_loop i | i <= n    = do
+                               x <- read v i
+                               write v i (f x)
+               | otherwise = return ()
 
 update :: MVector v m a => v a -> Stream (Int, a) -> m ()
 {-# INLINE update #-}
 
 update :: MVector v m a => v a -> Stream (Int, a) -> m ()
 {-# INLINE update #-}
index 0918e3c..2b453a4 100644 (file)
@@ -3,7 +3,7 @@
 #include "phases.h"
 
 module Data.Vector.MVector.Mut (
 #include "phases.h"
 
 module Data.Vector.MVector.Mut (
-  Mut(..), run, unstream, update
+  Mut(..), run, unstream, map, update
 ) where
 
 import qualified Data.Vector.MVector as MVector
 ) where
 
 import qualified Data.Vector.MVector as MVector
@@ -11,6 +11,8 @@ import           Data.Vector.MVector ( MVector )
 
 import           Data.Vector.Stream ( Stream )
 
 
 import           Data.Vector.Stream ( Stream )
 
+import Prelude hiding ( map )
+
 data Mut a = Mut (forall m mv. MVector mv m a => m (mv a))
 
 run :: MVector mv m a => Mut a -> m (mv a)
 data Mut a = Mut (forall m mv. MVector mv m a => m (mv a))
 
 run :: MVector mv m a => Mut a -> m (mv a)
@@ -21,6 +23,10 @@ unstream :: Stream a -> Mut a
 {-# INLINE_STREAM unstream #-}
 unstream s = Mut (MVector.unstream s)
 
 {-# INLINE_STREAM unstream #-}
 unstream s = Mut (MVector.unstream s)
 
+map :: (a -> a) -> Mut a -> Mut a
+{-# INLINE_STREAM map #-}
+map f (Mut p) = Mut (do { v <- p; MVector.map f v; return v })
+
 update :: Mut a -> Stream (Int, a) -> Mut a
 {-# INLINE_STREAM update #-}
 update (Mut p) s = Mut (do { v <- p; MVector.update v s; return v })
 update :: Mut a -> Stream (Int, a) -> Mut a
 {-# INLINE_STREAM update #-}
 update (Mut p) s = Mut (do { v <- p; MVector.update v s; return v })