Fix backpermute
authorRoman Leshchinskiy <rl@cse.unsw.edu.au>
Wed, 2 Dec 2009 06:49:34 +0000 (06:49 +0000)
committerRoman Leshchinskiy <rl@cse.unsw.edu.au>
Wed, 2 Dec 2009 06:49:34 +0000 (06:49 +0000)
Data/Vector/Fusion/Stream.hs
Data/Vector/Fusion/Stream/Monadic.hs
Data/Vector/Generic.hs

index 15e6adb..27b8bae 100644 (file)
@@ -37,7 +37,7 @@ module Data.Vector.Fusion.Stream (
   extract, init, tail, take, drop,
 
   -- * Mapping
-  map, concatMap,
+  map, concatMap, unbox,
   
   -- * Zipping
   zipWith, zipWith3,
@@ -248,6 +248,10 @@ map :: (a -> b) -> Stream a -> Stream b
 {-# INLINE map #-}
 map = M.map
 
+unbox :: Stream (Box a) -> Stream a
+{-# INLINE unbox #-}
+unbox = M.unbox
+
 concatMap :: (a -> Stream b) -> Stream a -> Stream b
 {-# INLINE concatMap #-}
 concatMap = M.concatMap
index 8f39566..85cee7c 100644 (file)
@@ -33,7 +33,7 @@ module Data.Vector.Fusion.Stream.Monadic (
   extract, init, tail, take, drop,
 
   -- * Mapping
-  map, mapM, mapM_, trans, concatMap,
+  map, mapM, mapM_, trans, unbox, concatMap,
   
   -- * Zipping
   zipWith, zipWithM, zipWith3, zipWith3M,
@@ -66,6 +66,7 @@ module Data.Vector.Fusion.Stream.Monadic (
 ) where
 
 import Data.Vector.Fusion.Stream.Size
+import Data.Vector.Fusion.Util ( Box(..) )
 
 import Control.Monad  ( liftM )
 import Prelude hiding ( length, null,
@@ -347,6 +348,18 @@ trans :: (Monad m, Monad m') => (forall a. m a -> m' a)
 {-# INLINE_STREAM trans #-}
 trans f (Stream step s n) = Stream (f . step) s n
 
+unbox :: Monad m => Stream m (Box a) -> Stream m a
+{-# INLINE_STREAM unbox #-}
+unbox (Stream step s n) = Stream step' s n
+  where
+    {-# INLINE_INNER step' #-}
+    step' s = do
+                r <- step s
+                case r of
+                  Yield (Box x) s' -> return $ Yield x s'
+                  Skip          s' -> return $ Skip    s'
+                  Done             -> return $ Done
+
 -- Zipping
 -- -------
 
index 58991a4..16b04bb 100644 (file)
@@ -406,9 +406,8 @@ backpermute :: (Vector v a, Vector v Int) => v a -> v Int -> v a
 {-# INLINE backpermute #-}
 backpermute v is = seq v
                  $ unstream
-                 $ MStream.trans (Id . unBox)
-                 $ MStream.mapM (indexM v)
-                 $ MStream.trans (Box . unId)
+                 $ Stream.unbox
+                 $ Stream.map (indexM v)
                  $ stream is
 
 reverse :: (Vector v a) => v a -> v a