Fix recycling for various functions
authorRoman Leshchinskiy <rl@cse.unsw.edu.au>
Fri, 23 Apr 2010 02:29:21 +0000 (02:29 +0000)
committerRoman Leshchinskiy <rl@cse.unsw.edu.au>
Fri, 23 Apr 2010 02:29:21 +0000 (02:29 +0000)
Data/Vector/Generic.hs
Data/Vector/Generic/New.hs

index 0d9f6cb..e7d4f12 100644 (file)
@@ -504,8 +504,8 @@ unsafeDrop n v = unsafeSlice n (length v - n) v
 
 unsafeAccum_stream
   :: Vector v a => (a -> b -> a) -> v a -> Stream (Int,b) -> v a
-{-# INLINE_STREAM unsafeAccum_stream #-}
-unsafeAccum_stream f v s = s `seq` modify (\mv -> M.unsafeAccum f mv s) v
+{-# INLINE unsafeAccum_stream #-}
+unsafeAccum_stream f = modifyWithStream (M.unsafeAccum f)
 
 unsafeAccum :: Vector v a => (a -> b -> a) -> v a -> [(Int,b)] -> v a
 {-# INLINE unsafeAccum #-}
@@ -523,8 +523,8 @@ unsafeAccumulate_ f v is xs
   = unsafeAccum_stream f v (Stream.zipWith (,) (stream is) (stream xs))
 
 accum_stream :: Vector v a => (a -> b -> a) -> v a -> Stream (Int,b) -> v a
-{-# INLINE_STREAM accum_stream #-}
-accum_stream f v s = s `seq` modify (\mv -> M.accum f mv s) v
+{-# INLINE accum_stream #-}
+accum_stream f = modifyWithStream (M.accum f)
 
 accum :: Vector v a => (a -> b -> a) -> v a -> [(Int,b)] -> v a
 {-# INLINE accum #-}
@@ -543,8 +543,8 @@ accumulate_ f v is xs = accum_stream f v (Stream.zipWith (,) (stream is)
                                         
 
 unsafeUpdate_stream :: Vector v a => v a -> Stream (Int,a) -> v a
-{-# INLINE_STREAM unsafeUpdate_stream #-}
-unsafeUpdate_stream v s = s `seq` modify (\mv -> M.unsafeUpdate mv s) v
+{-# INLINE unsafeUpdate_stream #-}
+unsafeUpdate_stream = modifyWithStream M.unsafeUpdate
 
 unsafeUpd :: Vector v a => v a -> [(Int, a)] -> v a
 {-# INLINE unsafeUpd #-}
@@ -560,8 +560,8 @@ unsafeUpdate_ v is w
   = unsafeUpdate_stream v (Stream.zipWith (,) (stream is) (stream w))
 
 update_stream :: Vector v a => v a -> Stream (Int,a) -> v a
-{-# INLINE_STREAM update_stream #-}
-update_stream v s = s `seq` modify (\mv -> M.update mv s) v
+{-# INLINE update_stream #-}
+update_stream = modifyWithStream M.update
 
 (//) :: Vector v a => v a -> [(Int, a)] -> v a
 {-# INLINE (//) #-}
@@ -1333,6 +1333,14 @@ modify :: Vector v a => (forall s. Mutable v s a -> ST s ()) -> v a -> v a
 {-# INLINE modify #-}
 modify p = new . New.modify p . clone
 
+-- We have to make sure that this is strict in the stream but we can't seq on
+-- it while fusion is happening. Hence this ugliness.
+modifyWithStream :: Vector v a
+                 => (forall s. Mutable v s a -> Stream b -> ST s ())
+                 -> v a -> Stream b -> v a
+{-# INLINE modifyWithStream #-}
+modifyWithStream p v s = new (New.modifyWithStream p (clone v) s)
+
 -- | Copy an immutable vector into a mutable one. The two vectors must have
 -- the same length. This is not checked.
 unsafeCopy
index ac4cfe6..6a8ef6e 100644 (file)
@@ -13,7 +13,7 @@
 --
 
 module Data.Vector.Generic.New (
-  New(..), create, run, apply, modify,
+  New(..), create, run, apply, modify, modifyWithStream,
   unstream, transform, unstreamR, transformR,
   slice, init, tail, take, drop,
   unsafeSlice, unsafeInit, unsafeTail
@@ -51,6 +51,11 @@ modify :: (forall s. Mutable v s a -> ST s ()) -> New v a -> New v a
 {-# INLINE modify #-}
 modify f (New p) = New (do { v <- p; f v; return v })
 
+modifyWithStream :: (forall s. Mutable v s a -> Stream b -> ST s ())
+                 -> New v a -> Stream b -> New v a
+{-# INLINE_STREAM modifyWithStream #-}
+modifyWithStream f (New p) s = s `seq` New (do { v <- p; f v s; return v })
+
 unstream :: Vector v a => Stream a -> New v a
 {-# INLINE_STREAM unstream #-}
 unstream s = s `seq` New (MVector.unstream s)