More flexible size hints
authorRoman Leshchinskiy <rl@cse.unsw.edu.au>
Mon, 7 Jul 2008 03:26:20 +0000 (03:26 +0000)
committerRoman Leshchinskiy <rl@cse.unsw.edu.au>
Mon, 7 Jul 2008 03:26:20 +0000 (03:26 +0000)
Data/Vector.hs
Data/Vector/Mutable.hs
Data/Vector/Prim.hs
Data/Vector/Stream.hs
Data/Vector/Stream/Size.hs [new file with mode: 0644]
Data/Vector/Unboxed.hs
Data/Vector/Unboxed/Mutable.hs
Data/Vector/Unboxed/Prim.hs
vector.cabal

index 9b0bfb9..c730264 100644 (file)
@@ -8,6 +8,7 @@ where
 import qualified Data.Vector.Prim    as Prim
 import qualified Data.Vector.Mutable as Mut
 
 import qualified Data.Vector.Prim    as Prim
 import qualified Data.Vector.Mutable as Mut
 
+import           Data.Vector.Stream.Size ( Size(..) )
 import qualified Data.Vector.Stream  as Stream
 import           Data.Vector.Stream ( Stream )
 
 import qualified Data.Vector.Stream  as Stream
 import           Data.Vector.Stream ( Stream )
 
@@ -44,7 +45,7 @@ new' n x init = runST (
 
 stream :: Vector a -> Stream a
 {-# INLINE_STREAM stream #-}
 
 stream :: Vector a -> Stream a
 {-# INLINE_STREAM stream #-}
-stream (Vector i n arr) = Stream.unfold get i n
+stream (Vector i n arr) = Stream.unfold get i `Stream.sized` Exact n
   where
     n' = n+i
 
   where
     n' = n+i
 
@@ -54,10 +55,11 @@ stream (Vector i n arr) = Stream.unfold get i n
 
 unstream :: Stream a -> Vector a
 {-# INLINE_STREAM unstream #-}
 
 unstream :: Stream a -> Vector a
 {-# INLINE_STREAM unstream #-}
-unstream s = new (Stream.bound s) (\mv ->
-  do
-    n <- Mut.fill mv s
-    return $ Mut.slice mv 0 n
+unstream s = runST (do
+    mv <- Mut.unstream s
+    let (mprim, i, n) = Mut.dataOf mv
+    prim <- Prim.unsafeFreeze mprim
+    return $ Vector i n prim
   )
 
 {-# RULES
   )
 
 {-# RULES
index c958d4b..244d996 100644 (file)
@@ -3,12 +3,13 @@
 module Data.Vector.Mutable (
   Vector,
 
 module Data.Vector.Mutable (
   Vector,
 
-  new, new', length, slice, read, write, fill,
+  new, new', length, slice, read, write, unstream, fill,
   dataOf
 ) where
 
 import qualified Data.Vector.Prim as Prim
 
   dataOf
 ) where
 
 import qualified Data.Vector.Prim as Prim
 
+import           Data.Vector.Stream.Size ( upperBound )
 import qualified Data.Vector.Stream as Stream
 import           Data.Vector.Stream ( Stream )
 
 import qualified Data.Vector.Stream as Stream
 import           Data.Vector.Stream ( Stream )
 
@@ -56,6 +57,41 @@ write (Vector i n v) j x
   = assert (j < n)
   $ Prim.write v (i+j) x
 
   = assert (j < n)
   $ Prim.write v (i+j) x
 
+unstream :: Stream a -> ST s (Vector s a)
+{-# INLINE unstream #-}
+unstream s = case upperBound (Stream.size s) of
+               Just n  -> unstream_known   s n
+               Nothing -> unstream_unknown s
+
+gROWTH_FACTOR :: Double
+gROWTH_FACTOR = 1.6
+
+unstream_known :: Stream a -> Int -> ST s (Vector s a)
+{-# INLINE unstream_known #-}
+unstream_known s n
+  = do
+      v  <- new n
+      n' <- fill v s
+      return $ slice v 0 n'
+
+unstream_unknown :: Stream a -> ST s (Vector s a)
+{-# INLINE unstream_unknown #-}
+unstream_unknown s
+  = do
+      v <- Prim.new 0
+      (w, n, _) <- Stream.foldM put (v, 0, 0) s
+      return $ Vector 0 n w
+  where
+    {-# INLINE put #-}
+    put (v, i, n) x = do
+                        (v', n') <- enlarge v i n
+                        Prim.write v' i x
+                        return (v', i+1, n')
+
+    {-# INLINE enlarge #-}
+    enlarge v i n | i < n     = return (v, n)
+                  | otherwise = Prim.grow v n gROWTH_FACTOR
+
 fill :: Vector s a -> Stream a -> ST s Int
 {-# INLINE fill #-}
 fill !v s = Stream.foldM put 0 s
 fill :: Vector s a -> Stream a -> ST s Int
 {-# INLINE fill #-}
 fill !v s = Stream.foldM put 0 s
index df5c68f..f9eb004 100644 (file)
@@ -2,13 +2,16 @@
 
 module Data.Vector.Prim (
   Vector, MutableVector,
 
 module Data.Vector.Prim (
   Vector, MutableVector,
-  new, new', unsafeFreeze, at, at', read, write
+  new, new', unsafeFreeze, at, at', read, write, copy, grow
 ) where
 
 import GHC.Prim (
     Array#, MutableArray#,
     newArray#, readArray#, writeArray#, indexArray#, unsafeFreezeArray#
   )
 ) where
 
 import GHC.Prim (
     Array#, MutableArray#,
     newArray#, readArray#, writeArray#, indexArray#, unsafeFreezeArray#
   )
+import GHC.Float (
+    double2Int, int2Double
+  )
 import GHC.ST (
     ST(..)
   )
 import GHC.ST (
     ST(..)
   )
@@ -54,3 +57,25 @@ write :: MutableVector s a -> Int -> a -> ST s ()
 write (MutableVector arr#) (I# n#) x = ST $ \s# ->
   case writeArray# arr# n# x s# of s2# -> (# s2#, () #)
 
 write (MutableVector arr#) (I# n#) x = ST $ \s# ->
   case writeArray# arr# n# x s# of s2# -> (# s2#, () #)
 
+copy :: MutableVector s a -> Int -> MutableVector s a -> Int -> Int -> ST s ()
+{-# INLINE copy #-}
+copy mv i mw j n = do_copy i j n
+  where
+    do_copy i j 0 = return ()
+    do_copy i j n = do
+                      x <- read mw j
+                      write mv i x
+                      do_copy (i+1) (j+1) (n-1)
+
+grow :: MutableVector s a -> Int -> Double -> ST s (MutableVector s a, Int)
+{-# INLINE grow #-}
+grow v n r
+  = do
+      w <- new m
+      copy w 0 v 0 n
+      return (w, m)
+  where
+    n' = double2Int (int2Double n * r)
+    m | n' <= n   = n+1
+      | otherwise = n'
+
index c563c33..0ef303c 100644 (file)
@@ -5,13 +5,15 @@
 module Data.Vector.Stream (
   Step(..), Stream(..),
 
 module Data.Vector.Stream (
   Step(..), Stream(..),
 
-  bound, unfold,
+  size, sized, unfold,
   empty, singleton, replicate, (++),
   map, filter, zipWith,
   foldr, foldl, foldl',
   mapM_, foldM
 ) where
 
   empty, singleton, replicate, (++),
   map, filter, zipWith,
   foldr, foldl, foldl',
   mapM_, foldM
 ) where
 
+import Data.Vector.Stream.Size
+
 import Prelude hiding ( replicate, (++), map, filter, zipWith,
                         foldr, foldl,
                         mapM_ )
 import Prelude hiding ( replicate, (++), map, filter, zipWith,
                         foldr, foldl,
                         mapM_ )
@@ -20,15 +22,19 @@ data Step s a = Yield a s
               | Skip    s
               | Done
 
               | Skip    s
               | Done
 
-data Stream a = forall s. Stream (s -> Step s a) s Int
+data Stream a = forall s. Stream (s -> Step s a) s Size
+
+size :: Stream a -> Size
+{-# INLINE size #-}
+size (Stream _ _ sz) = sz
 
 
-bound :: Stream a -> Int
-{-# INLINE bound #-}
-bound (Stream _ _ n) = n
+sized :: Stream a -> Size -> Stream a
+{-# INLINE_STREAM sized #-}
+sized (Stream step s _) sz = Stream step s sz
 
 
-unfold :: (s -> Maybe (a, s)) -> s -> Int -> Stream a
+unfold :: (s -> Maybe (a, s)) -> s -> Stream a
 {-# INLINE_STREAM unfold #-}
 {-# INLINE_STREAM unfold #-}
-unfold f s n = Stream step s n
+unfold f s = Stream step s Unknown
   where
     {-# INLINE step #-}
     step s = case f s of
   where
     {-# INLINE step #-}
     step s = case f s of
@@ -37,11 +43,11 @@ unfold f s n = Stream step s n
 
 empty :: Stream a
 {-# INLINE_STREAM empty #-}
 
 empty :: Stream a
 {-# INLINE_STREAM empty #-}
-empty = Stream (const Done) () 0
+empty = Stream (const Done) () (Exact 0)
 
 singleton :: a -> Stream a
 {-# INLINE_STREAM singleton #-}
 
 singleton :: a -> Stream a
 {-# INLINE_STREAM singleton #-}
-singleton x = Stream step True 1
+singleton x = Stream step True (Exact 1)
   where
     {-# INLINE step #-}
     step True  = Yield x False
   where
     {-# INLINE step #-}
     step True  = Yield x False
@@ -49,7 +55,7 @@ singleton x = Stream step True 1
 
 replicate :: Int -> a -> Stream a
 {-# INLINE_STREAM replicate #-}
 
 replicate :: Int -> a -> Stream a
 {-# INLINE_STREAM replicate #-}
-replicate n x = Stream step n (max n 0)
+replicate n x = Stream step n (Exact (max n 0))
   where
     {-# INLINE step #-}
     step i | i > 0     = Yield x (i-1)
   where
     {-# INLINE step #-}
     step i | i > 0     = Yield x (i-1)
@@ -81,7 +87,7 @@ map f (Stream step s n) = Stream step' s n
 
 filter :: (a -> Bool) -> Stream a -> Stream a
 {-# INLINE_STREAM filter #-}
 
 filter :: (a -> Bool) -> Stream a -> Stream a
 {-# INLINE_STREAM filter #-}
-filter f (Stream step s n) = Stream step' s n
+filter f (Stream step s n) = Stream step' s (toMax n)
   where
     {-# INLINE step' #-}
     step' s = case step s of
   where
     {-# INLINE step' #-}
     step' s = case step s of
@@ -93,7 +99,7 @@ filter f (Stream step s n) = Stream step' s n
 zipWith :: (a -> b -> c) -> Stream a -> Stream b -> Stream c
 {-# INLINE_STREAM zipWith #-}
 zipWith f (Stream stepa sa na) (Stream stepb sb nb)
 zipWith :: (a -> b -> c) -> Stream a -> Stream b -> Stream c
 {-# INLINE_STREAM zipWith #-}
 zipWith f (Stream stepa sa na) (Stream stepb sb nb)
-  = Stream step (sa, sb, Nothing) (min na nb)
+  = Stream step (sa, sb, Nothing) (smaller na nb)
   where
     {-# INLINE step #-}
     step (sa, sb, Nothing) = case stepa sa of
   where
     {-# INLINE step #-}
     step (sa, sb, Nothing) = case stepa sa of
diff --git a/Data/Vector/Stream/Size.hs b/Data/Vector/Stream/Size.hs
new file mode 100644 (file)
index 0000000..6389256
--- /dev/null
@@ -0,0 +1,66 @@
+module Data.Vector.Stream.Size (
+  Size(..), smaller, larger, toMax, upperBound
+) where
+
+data Size = Exact Int
+          | Max   Int
+          | Unknown
+        deriving( Eq, Show )
+
+instance Num Size where
+  Exact m + Exact n = Exact (m+n)
+  Exact m + Max   n = Max   (m+n)
+
+  Max   m + Exact n = Max   (m+n)
+  Max   m + Max   n = Max   (m+n)
+
+  _       + _       = Unknown
+
+
+  Exact m - Exact n = Exact (m-n)
+  Exact m - Max   n = Max   m
+
+  Max   m - Exact n = Max   (m-n)
+  Max   m - Max   n = Max   m
+  Max   m - Unknown = Max   m
+
+  _       - _       = Unknown
+
+
+  fromInteger n     = Exact (fromInteger n)
+
+
+smaller :: Size -> Size -> Size
+smaller (Exact m) (Exact n) = Exact (m `min` n)
+smaller (Exact m) (Max   n) = Max   (m `min` n)
+smaller (Exact m) Unknown   = Max   m
+smaller (Max   m) (Exact n) = Max   (m `min` n)
+smaller (Max   m) (Max   n) = Max   (m `min` n)
+smaller (Max   m) Unknown   = Max   m
+smaller Unknown   (Exact n) = Max   n
+smaller Unknown   (Max   n) = Max   n
+smaller Unknown   Unknown   = Unknown
+
+larger :: Size -> Size -> Size
+larger (Exact m) (Exact n)             = Exact (m `max` n)
+larger (Exact m) (Max   n) | m >= n    = Exact m
+                           | otherwise = Max   n
+larger (Max   m) (Exact n) | n >= m    = Exact n
+                           | otherwise = Max   m
+larger (Max   m) (Max   n)             = Max   (m `max` n)
+larger _         _                     = Unknown
+
+toMax :: Size -> Size
+toMax (Exact n) = Max n
+toMax (Max   n) = Max n
+toMax Unknown   = Unknown
+
+lowerBound :: Size -> Int
+lowerBound (Exact n) = n
+lowerBound _         = 0
+
+upperBound :: Size -> Maybe Int
+upperBound (Exact n) = Just n
+upperBound (Max   n) = Just n
+upperBound Unknown   = Nothing
+
index 7daef6a..6deb862 100644 (file)
@@ -9,6 +9,7 @@ import qualified Data.Vector.Unboxed.Prim as Prim
 import qualified Data.Vector.Unboxed.Mutable as Mut
 import           Data.Vector.Unboxed.Unbox ( Unbox )
 
 import qualified Data.Vector.Unboxed.Mutable as Mut
 import           Data.Vector.Unboxed.Unbox ( Unbox )
 
+import           Data.Vector.Stream.Size ( Size(..) )
 import qualified Data.Vector.Stream as Stream
 import           Data.Vector.Stream ( Stream )
 
 import qualified Data.Vector.Stream as Stream
 import           Data.Vector.Stream ( Stream )
 
@@ -35,7 +36,7 @@ new n init = runST (
 
 stream :: Unbox a => Vector a -> Stream a
 {-# INLINE_STREAM stream #-}
 
 stream :: Unbox a => Vector a -> Stream a
 {-# INLINE_STREAM stream #-}
-stream (Vector i n arr) = Stream.unfold get i n
+stream (Vector i n arr) = Stream.unfold get i `Stream.sized` Exact n
   where
     n' = n+i
 
   where
     n' = n+i
 
@@ -45,10 +46,11 @@ stream (Vector i n arr) = Stream.unfold get i n
 
 unstream :: Unbox a => Stream a -> Vector a
 {-# INLINE_STREAM unstream #-}
 
 unstream :: Unbox a => Stream a -> Vector a
 {-# INLINE_STREAM unstream #-}
-unstream s = new (Stream.bound s) (\mv ->
-  do
-    n <- Mut.fill mv s
-    return $ Mut.slice mv 0 n
+unstream s = runST (do
+    mv <- Mut.unstream s
+    let (mprim, i, n) = Mut.dataOf mv
+    prim <- Prim.unsafeFreeze mprim
+    return $ Vector i n prim
   )
 
 {-# RULES
   )
 
 {-# RULES
index 97b10a4..32fb5fd 100644 (file)
@@ -3,13 +3,14 @@
 module Data.Vector.Unboxed.Mutable (
   Vector,
 
 module Data.Vector.Unboxed.Mutable (
   Vector,
 
-  new, length, slice, read, write, fill, fillIndexed,
+  new, length, slice, read, write, unstream, fill, fillIndexed,
   dataOf
 ) where
 
 import qualified Data.Vector.Unboxed.Prim as Prim
 import           Data.Vector.Unboxed.Unbox ( Unbox )
 
   dataOf
 ) where
 
 import qualified Data.Vector.Unboxed.Prim as Prim
 import           Data.Vector.Unboxed.Unbox ( Unbox )
 
+import           Data.Vector.Stream.Size ( upperBound )
 import qualified Data.Vector.Stream as Stream
 import           Data.Vector.Stream ( Stream )
 
 import qualified Data.Vector.Stream as Stream
 import           Data.Vector.Stream ( Stream )
 
@@ -59,6 +60,41 @@ write (Vector i n v) j x
   = assert (j < n)
   $ Prim.write v (i+j) x
 
   = assert (j < n)
   $ Prim.write v (i+j) x
 
+unstream :: Unbox a => Stream a -> ST s (Vector s a)
+{-# INLINE unstream #-}
+unstream s = case upperBound (Stream.size s) of
+               Just n  -> unstream_known   s n
+               Nothing -> unstream_unknown s
+
+gROWTH_FACTOR :: Double
+gROWTH_FACTOR = 1.6
+
+unstream_known :: Unbox a => Stream a -> Int -> ST s (Vector s a)
+{-# INLINE unstream_known #-}
+unstream_known s n
+  = do
+      v  <- new n
+      n' <- fill v s
+      return $ slice v 0 n'
+
+unstream_unknown :: Unbox a => Stream a -> ST s (Vector s a)
+{-# INLINE unstream_unknown #-}
+unstream_unknown s
+  = do
+      v <- Prim.new 0
+      (w, n, _) <- Stream.foldM put (v, 0, 0) s
+      return $ Vector 0 n w
+  where
+    {-# INLINE put #-}
+    put (v, i, n) x = do
+                        (v', n') <- enlarge v i n
+                        Prim.write v' i x
+                        return (v', i+1, n')
+
+    {-# INLINE enlarge #-}
+    enlarge v i n | i < n     = return (v, n)
+                  | otherwise = Prim.grow v n gROWTH_FACTOR
+
 fill :: Unbox a => Vector s a -> Stream a -> ST s Int
 {-# INLINE fill #-}
 fill !v s = Stream.foldM put 0 s
 fill :: Unbox a => Vector s a -> Stream a -> ST s Int
 {-# INLINE fill #-}
 fill !v s = Stream.foldM put 0 s
index 4bc224c..12ca5cd 100644 (file)
@@ -3,7 +3,8 @@
 module Data.Vector.Unboxed.Prim (
   Unbox(..), Vector, MutableVector,
 
 module Data.Vector.Unboxed.Prim (
   Unbox(..), Vector, MutableVector,
 
-  size, new, unsafeFreeze, at, read, write
+  size, new, unsafeFreeze, at, read, write,
+  copy, grow
 ) where
 
 import Data.Vector.Unboxed.Unbox
 ) where
 
 import Data.Vector.Unboxed.Unbox
@@ -12,6 +13,9 @@ import GHC.Prim (
     ByteArray#, MutableByteArray#,
     newByteArray#, unsafeFreezeByteArray#,
   )
     ByteArray#, MutableByteArray#,
     newByteArray#, unsafeFreezeByteArray#,
   )
+import GHC.Float (
+    int2Double, double2Int
+  )
 import GHC.ST (
     ST(..)
   )
 import GHC.ST (
     ST(..)
   )
@@ -53,3 +57,28 @@ write :: Unbox a => MutableVector s a -> Int -> a -> ST s ()
 write (MutableVector arr#) (I# i#) x = ST $ \s# ->
   case write# arr# i# x s# of s2# -> (# s2#, () #)
 
 write (MutableVector arr#) (I# i#) x = ST $ \s# ->
   case write# arr# i# x s# of s2# -> (# s2#, () #)
 
+copy :: Unbox a => MutableVector s a -> Int
+                -> MutableVector s a -> Int -> Int
+                -> ST s ()
+{-# INLINE copy #-}
+copy mv i mw j n = do_copy i j n
+  where
+    do_copy i j 0 = return ()
+    do_copy i j n = do
+                      x <- read mw j
+                      write mv i x
+                      do_copy (i+1) (j+1) (n-1)
+
+grow :: Unbox a => MutableVector s a -> Int -> Double
+                -> ST s (MutableVector s a, Int)
+{-# INLINE grow #-}
+grow v n r
+  = do
+      w <- new m
+      copy w 0 v 0 n
+      return (w, m)
+  where
+    n' = double2Int (int2Double n * r)
+    m | n' <= n   = n+1
+      | otherwise = n'
+
index 581d040..45fd9a4 100644 (file)
@@ -12,6 +12,7 @@ Build-Type:     Simple
 
 Library
   Exposed-Modules:
 
 Library
   Exposed-Modules:
+        Data.Vector.Stream.Size
         Data.Vector.Stream
 
         Data.Vector.Unboxed.Unbox
         Data.Vector.Stream
 
         Data.Vector.Unboxed.Unbox