Right-to-left streaming/unstreaming
authorRoman Leshchinskiy <rl@cse.unsw.edu.au>
Thu, 10 Dec 2009 04:41:41 +0000 (04:41 +0000)
committerRoman Leshchinskiy <rl@cse.unsw.edu.au>
Thu, 10 Dec 2009 04:41:41 +0000 (04:41 +0000)
Data/Vector/Generic.hs
Data/Vector/Generic/Mutable.hs
Data/Vector/Generic/New.hs

index 96c3ffe..e90ce8a 100644 (file)
@@ -83,7 +83,7 @@ module Data.Vector.Generic (
   toList, fromList,
 
   -- * Conversion to/from Streams
-  stream, unstream,
+  stream, unstream, streamR, unstreamR,
 
   -- * MVector-based initialisation
   new
@@ -220,6 +220,48 @@ unstream s = new (New.unstream s)
 
  #-}
 
+-- | Convert a vector to a 'Stream'
+streamR :: Vector v a => v a -> Stream a
+{-# INLINE_STREAM streamR #-}
+streamR v = v `seq` (Stream.unfoldr get n `Stream.sized` Exact n)
+  where
+    n = length v
+
+    -- NOTE: the False case comes first in Core so making it the recursive one
+    -- makes the code easier to read
+    {-# INLINE get #-}
+    get 0 = Nothing
+    get i = let i' = i-1
+            in
+            case basicUnsafeIndexM v i' of Box x -> Just (x, i')
+
+-- | Create a vector from a 'Stream'
+unstreamR :: Vector v a => Stream a -> v a
+{-# INLINE unstreamR #-}
+unstreamR s = new (New.unstreamR s)
+
+{-# RULES
+
+"streamR/unstreamR [Vector]" forall v s.
+  streamR (new' v (New.unstreamR s)) = s
+
+"New.unstreamR/streamR/new [Vector]" forall v p.
+  New.unstreamR (streamR (new' v p)) = p
+
+ #-}
+
+{-# RULES
+
+"inplace [Vector]"
+  forall (f :: forall m. Monad m => MStream m a -> MStream m a) v m.
+  New.unstreamR (inplace f (streamR (new' v m))) = New.transformR f m
+
+"uninplace [Vector]"
+  forall (f :: forall m. Monad m => MStream m a -> MStream m a) v m.
+  streamR (new' v (New.transformR f m)) = inplace f (streamR (new' v m))
+
+ #-}
+
 -- Length
 -- ------
 
index 6586d87..cd49b13 100644 (file)
@@ -26,7 +26,8 @@ module Data.Vector.Generic.Mutable (
   unsafeCopy, unsafeGrow,
 
   -- * Internal operations
-  unstream, transform, unsafeAccum, accum, unsafeUpdate, update, reverse,
+  unstream, transform, unstreamR, transformR,
+  unsafeAccum, accum, unsafeUpdate, update, reverse,
   unstablePartition, unstablePartitionStream
 ) where
 
@@ -103,6 +104,7 @@ class MVector v a where
   -- called directly, use 'unsafeGrow' instead.
   basicUnsafeGrow  :: PrimMonad m => v (PrimState m) a -> Int
                                                        -> m (v (PrimState m) a)
+
   {-# INLINE basicUnsafeNewWith #-}
   basicUnsafeNewWith n x
     = do
@@ -190,6 +192,16 @@ unsafeGrow :: (PrimMonad m, MVector v a)
 unsafeGrow v n = UNSAFE_CHECK(checkLength) "unsafeGrow" n
                $ basicUnsafeGrow v n
 
+unsafeGrowFront :: (PrimMonad m, MVector v a)
+                        => v (PrimState m) a -> Int -> m (v (PrimState m) a)
+{-# INLINE unsafeGrowFront #-}
+unsafeGrowFront v by = UNSAFE_CHECK(checkLength) "unsafeGrowFront" by
+                     $ do
+                         let n = length v
+                         v' <- basicUnsafeNew (by+n)
+                         basicUnsafeCopy (basicUnsafeSlice by n v') v
+                         return v'
+
 -- | Length of the mutable vector.
 length :: MVector v a => v s a -> Int
 {-# INLINE length #-}
@@ -255,15 +267,24 @@ grow :: (PrimMonad m, MVector v a)
 grow v by = BOUNDS_CHECK(checkLength) "grow" by
           $ unsafeGrow v by
 
+enlarge_delta v = max 1
+                $ double2Int
+                $ int2Double (length v) * gROWTH_FACTOR
+
 -- | Grow a vector logarithmically
 enlarge :: (PrimMonad m, MVector v a)
                 => v (PrimState m) a -> m (v (PrimState m) a)
 {-# INLINE enlarge #-}
-enlarge v = unsafeGrow v
-          $ max 1
-          $ double2Int
-          $ int2Double (length v) * gROWTH_FACTOR
-
+enlarge v = unsafeGrow v (enlarge_delta v)
+
+enlargeFront :: (PrimMonad m, MVector v a)
+                => v (PrimState m) a -> m (v (PrimState m) a, Int)
+{-# INLINE enlargeFront #-}
+enlargeFront v = do
+                   v' <- unsafeGrowFront v by
+                   return (v', by)
+  where
+    by = enlarge_delta v
 
 -- | Yield a part of the mutable vector without copying it.
 slice :: MVector v a => Int -> Int -> v s a -> v s a
@@ -325,6 +346,21 @@ unsafeAppend1 v i x
                        $ unsafeWrite v' i x
                      return v'
 
+unsafePrepend1 :: (PrimMonad m, MVector v a)
+        => v (PrimState m) a -> Int -> a -> m (v (PrimState m) a, Int)
+{-# INLINE_INNER unsafePrepend1 #-}
+unsafePrepend1 v i x
+  | i /= 0    = do
+                  let i' = i-1
+                  unsafeWrite v i' x
+                  return (v, i')
+  | otherwise = do
+                  (v', i) <- enlargeFront v
+                  let i' = i-1
+                  INTERNAL_CHECK(checkIndex) "unsafePrepend1" i' (length v')
+                    $ unsafeWrite v' i' x
+                  return (v', i')
+
 
 mstream :: (PrimMonad m, MVector v a) => v (PrimState m) a -> MStream m a
 {-# INLINE mstream #-}
@@ -368,10 +404,10 @@ mrstream v = v `seq` (MStream.unfoldrM get n `MStream.sized` Exact n)
       where
         j = i-1
 
-mrunstream :: (PrimMonad m, MVector v a)
+munstreamR :: (PrimMonad m, MVector v a)
            => v (PrimState m) a -> MStream m a -> m (v (PrimState m) a)
-{-# INLINE mrunstream #-}
-mrunstream v s = v `seq` do
+{-# INLINE munstreamR #-}
+munstreamR v s = v `seq` do
                            i <- MStream.foldM put n s
                            return $ unsafeSlice i (n-i) v
   where
@@ -384,10 +420,10 @@ mrunstream v s = v `seq` do
       where
         j = i-1
 
-rtransform :: (PrimMonad m, MVector v a)
+transformR :: (PrimMonad m, MVector v a)
   => (MStream m a -> MStream m a) -> v (PrimState m) a -> m (v (PrimState m) a)
-{-# INLINE_STREAM rtransform #-}
-rtransform f v = mrunstream v (f (mrstream v))
+{-# INLINE_STREAM transformR #-}
+transformR f v = munstreamR v (f (mrstream v))
 
 -- | Create a new mutable vector and fill it with elements from the 'Stream'.
 -- The vector will grow logarithmically if the 'Size' hint of the 'Stream' is
@@ -428,6 +464,45 @@ unstreamUnknown s
                     v' <- unsafeAppend1 v i x
                     return (v',i+1)
 
+-- | Create a new mutable vector and fill it with elements from the 'Stream'.
+-- The vector will grow logarithmically if the 'Size' hint of the 'Stream' is
+-- inexact.
+unstreamR :: (PrimMonad m, MVector v a) => Stream a -> m (v (PrimState m) a)
+{-# INLINE_STREAM unstreamR #-}
+unstreamR s = case upperBound (Stream.size s) of
+               Just n  -> unstreamRMax     s n
+               Nothing -> unstreamRUnknown s
+
+unstreamRMax
+  :: (PrimMonad m, MVector v a) => Stream a -> Int -> m (v (PrimState m) a)
+{-# INLINE unstreamRMax #-}
+unstreamRMax s n
+  = do
+      v <- INTERNAL_CHECK(checkLength) "unstreamRMax" n
+           $ unsafeNew n
+      let put i x = do
+                      let i' = i-1
+                      INTERNAL_CHECK(checkIndex) "unstreamRMax" i' n
+                        $ unsafeWrite v i x
+                      return i
+      i <- Stream.foldM' put n s
+      return $ INTERNAL_CHECK(checkSlice) "unstreamRMax" i (n-i) n
+             $ unsafeSlice i (n-i) v
+
+unstreamRUnknown
+  :: (PrimMonad m, MVector v a) => Stream a -> m (v (PrimState m) a)
+{-# INLINE unstreamRUnknown #-}
+unstreamRUnknown s
+  = do
+      v <- unsafeNew 0
+      (v', i) <- Stream.foldM put (v, 0) s
+      let n = length v'
+      return $ INTERNAL_CHECK(checkSlice) "unstreamRUnknown" i (n-i) n
+             $ slice i (n-i) v'
+  where
+    {-# INLINE_INNER put #-}
+    put (v,i) x = unsafePrepend1 v i x
+
 unsafeAccum :: (PrimMonad m, MVector v a)
             => (a -> b -> a) -> v (PrimState m) a -> Stream (Int, b) -> m ()
 {-# INLINE unsafeAccum #-}
index de601d5..aeba427 100644 (file)
@@ -13,7 +13,8 @@
 --
 
 module Data.Vector.Generic.New (
-  New(..), run, unstream, transform, accum, update, reverse,
+  New(..), run, unstream, transform, unstreamR, transformR,
+  accum, update, reverse,
   slice, init, tail, take, drop,
   unsafeSlice, unsafeInit, unsafeTail,
   unsafeAccum, unsafeUpdate
@@ -68,6 +69,30 @@ transform f (New p) = New (MVector.transform f =<< p)
 
  #-}
 
+
+unstreamR :: Stream a -> New a
+{-# INLINE_STREAM unstreamR #-}
+unstreamR s = s `seq` New (MVector.unstreamR s)
+
+transformR :: (forall m. Monad m => MStream m a -> MStream m a) -> New a -> New a
+{-# INLINE_STREAM transformR #-}
+transformR f (New p) = New (MVector.transformR f =<< p)
+
+{-# RULES
+
+"transformR/transformR [New]"
+  forall (f :: forall m. Monad m => MStream m a -> MStream m a)
+         (g :: forall m. Monad m => MStream m a -> MStream m a)
+         p .
+  transformR f (transformR g p) = transformR (f . g) p
+
+"transformR/unstreamR [New]"
+  forall (f :: forall m. Monad m => MStream m a -> MStream m a)
+         s.
+  transformR f (unstreamR s) = unstreamR (f s)
+
+ #-}
+
 slice :: Int -> Int -> New a -> New a
 {-# INLINE_STREAM slice #-}
 slice i n m = apply (MVector.slice i n) m