author Roman Leshchinskiy Thu, 10 Dec 2009 04:41:41 +0000 (04:41 +0000) committer Roman Leshchinskiy Thu, 10 Dec 2009 04:41:41 +0000 (04:41 +0000)

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