Boxed Vectors
authorRoman Leshchinskiy <rl@cse.unsw.edu.au>
Sun, 6 Jul 2008 16:44:38 +0000 (16:44 +0000)
committerRoman Leshchinskiy <rl@cse.unsw.edu.au>
Sun, 6 Jul 2008 16:44:38 +0000 (16:44 +0000)
Data/Vector.hs [new file with mode: 0644]
Data/Vector/Mutable.hs [new file with mode: 0644]
Data/Vector/Prim.hs [new file with mode: 0644]
vector.cabal

diff --git a/Data/Vector.hs b/Data/Vector.hs
new file mode 100644 (file)
index 0000000..dc6faac
--- /dev/null
@@ -0,0 +1,110 @@
+{-# LANGUAGE RankNTypes, CPP #-}
+
+#include "phases.h"
+
+module Data.Vector
+where
+
+import qualified Data.Vector.Prim    as Prim
+import qualified Data.Vector.Mutable as Mut
+
+import qualified Data.Vector.Stream  as Stream
+import           Data.Vector.Stream ( Step(..), Stream(..) )
+
+import Control.Exception ( assert )
+import Control.Monad.ST  ( ST, runST )
+
+import Prelude hiding ( length, (++) )
+
+data Vector a = Vector {-# UNPACK #-} !Int
+                       {-# UNPACK #-} !Int
+                       {-# UNPACK #-} !(Prim.Vector a)
+
+new :: Int -> (forall s. Mut.Vector s a -> ST s (Mut.Vector s a)) -> Vector a
+{-# INLINE new #-}
+new n init = runST (
+  do
+    mv <- Mut.new n
+    mv' <- init mv
+    let (mprim, i, n') = Mut.dataOf mv'
+    prim <- Prim.unsafeFreeze mprim
+    return $ Vector i n' prim
+  )
+
+new' :: Int -> a -> (forall s. Mut.Vector s a -> ST s (Mut.Vector s a)) -> Vector a
+{-# INLINE new' #-}
+new' n x init = runST (
+  do
+    mv <- Mut.new' n x
+    mv' <- init mv
+    let (mprim, i, n') = Mut.dataOf mv'
+    prim <- Prim.unsafeFreeze mprim
+    return $ Vector i n' prim
+  )
+
+stream :: Vector a -> Stream a
+{-# INLINE_STREAM stream #-}
+stream (Vector i n arr) = Stream get i n
+  where
+    n' = n+i
+
+    {-# INLINE get #-}
+    get j | j < n'    = Prim.at' arr j $ \x -> Yield x (j+1)
+          | otherwise = Done
+
+unstream :: Stream a -> Vector a
+{-# INLINE_STREAM unstream #-}
+unstream s@(Stream _ _ n) = new n (\mv ->
+  do
+    n' <- Mut.fill mv s
+    return $ Mut.slice mv 0 n'
+  )
+
+{-# RULES
+
+"stream/unstream [Vector]" forall s.
+  stream (unstream s) = s
+
+ #-}
+
+length :: Vector a -> Int
+{-# INLINE length #-}
+length (Vector _ n _) = n
+
+slice :: Vector a -> Int -> Int -> Vector a
+{-# INLINE slice #-}
+slice (Vector i n arr) j m
+  = assert (j + m <= n && j >= 0 && m >= 0)
+  $ Vector (i+j) m arr
+
+unsafeAt ::Vector a -> Int -> a
+{-# INLINE unsafeAt #-}
+unsafeAt (Vector i _ arr) j = Prim.at arr (i+j)
+
+at :: Vector a -> Int -> a
+{-# INLINE at #-}
+at v i = assert (i >= 0 && i < length v)
+       $ unsafeAt v i
+
+infixr ++
+(++) :: Vector a -> Vector a -> Vector a
+{-# INLINE (++) #-}
+v ++ w = unstream (stream v Stream.++ stream w)
+
+map :: (a -> b) -> Vector a -> Vector b
+{-# INLINE map #-}
+map f = unstream . Stream.map f . stream
+
+filter :: (a -> Bool) -> Vector a -> Vector a
+{-# INLINE filter #-}
+filter f = unstream . Stream.filter f . stream
+
+zipWith :: (a -> b -> c) -> Vector a -> Vector b -> Vector c
+{-# INLINE zipWith #-}
+zipWith f v w = unstream
+              $ Stream.zipWith f (stream v) (stream w)
+
+foldl' :: (a -> b -> b) -> b -> Vector a -> b
+{-# INLINE foldl' #-}
+foldl' f z = Stream.foldl' f z . stream
+
diff --git a/Data/Vector/Mutable.hs b/Data/Vector/Mutable.hs
new file mode 100644 (file)
index 0000000..c958d4b
--- /dev/null
@@ -0,0 +1,65 @@
+{-# LANGUAGE BangPatterns #-}
+
+module Data.Vector.Mutable (
+  Vector,
+
+  new, new', length, slice, read, write, fill,
+  dataOf
+) where
+
+import qualified Data.Vector.Prim as Prim
+
+import qualified Data.Vector.Stream as Stream
+import           Data.Vector.Stream ( Stream )
+
+import Control.Exception ( assert )
+import Control.Monad.ST  ( ST )
+
+import Prelude hiding ( length, read )
+
+data Vector s a = Vector {-# UNPACK #-} !Int
+                         {-# UNPACK #-} !Int
+                         {-# UNPACK #-} !(Prim.MutableVector s a)
+
+dataOf :: Vector s a -> (Prim.MutableVector s a, Int, Int)
+{-# INLINE dataOf #-}
+dataOf (Vector i n v) = (v, i, n)
+
+new :: Int -> ST s (Vector s a)
+{-# INLINE new #-}
+new n = new' n (error "Data.Vector.Mutable: uninitialised element")
+
+new' :: Int -> a -> ST s (Vector s a)
+{-# INLINE new' #-}
+new' n x = assert (n >= 0)
+         $ Vector 0 n `fmap` Prim.new' n x
+
+length :: Vector s a -> Int
+{-# INLINE length #-}
+length (Vector _ n _) = n
+
+slice :: Vector s a -> Int -> Int -> Vector s a
+{-# INLINE slice #-}
+slice (Vector i n v) j m
+  = assert (j + m <= n && j >= 0 && m >= 0)
+  $ Vector (i+j) m v
+
+read :: Vector s a -> Int -> ST s a
+{-# INLINE read #-}
+read (Vector i n v) j
+  = assert (j < n)
+  $ Prim.read v (i+j)
+
+write :: Vector s a -> Int -> a -> ST s ()
+{-# INLINE write #-}
+write (Vector i n v) j x
+  = assert (j < n)
+  $ Prim.write v (i+j) x
+
+fill :: Vector s a -> Stream a -> ST s Int
+{-# INLINE fill #-}
+fill !v s = Stream.foldM put 0 s
+  where
+    {-# INLINE put #-}
+    put i x = do { write v i x; return (i+1) }
+
diff --git a/Data/Vector/Prim.hs b/Data/Vector/Prim.hs
new file mode 100644 (file)
index 0000000..df5c68f
--- /dev/null
@@ -0,0 +1,56 @@
+{-# LANGUAGE MagicHash, UnboxedTuples #-}
+
+module Data.Vector.Prim (
+  Vector, MutableVector,
+  new, new', unsafeFreeze, at, at', read, write
+) where
+
+import GHC.Prim (
+    Array#, MutableArray#,
+    newArray#, readArray#, writeArray#, indexArray#, unsafeFreezeArray#
+  )
+import GHC.ST (
+    ST(..)
+  )
+import GHC.Base (
+    Int(..)
+  )
+
+import Prelude hiding ( read )
+
+data Vector          a = Vector (Array# a)
+data MutableVector s a = MutableVector (MutableArray# s a)
+
+new :: Int -> ST s (MutableVector s a)
+{-# INLINE new #-}
+new n = new' n (error "Data.Vector: uninitialised element")
+
+new' :: Int -> a -> ST s (MutableVector s a)
+{-# INLINE new' #-}
+new' (I# n#) x = ST $ \s# ->
+  case newArray# n# x s# of
+    (# s2#, arr# #) -> (# s2#, MutableVector arr# #)
+
+unsafeFreeze :: MutableVector s a -> ST s (Vector a)
+{-# INLINE unsafeFreeze #-}
+unsafeFreeze (MutableVector arr#) = ST $ \s# ->
+  case unsafeFreezeArray# arr# s# of
+    (# s2, frozen# #) -> (# s2, Vector frozen# #)
+
+at :: Vector a -> Int -> a
+{-# INLINE at #-}
+at v i = at' v i id
+
+at' :: Vector a -> Int -> (a -> b) -> b
+{-# INLINE at' #-}
+at' (Vector arr#) (I# n#) f = case indexArray# arr# n# of (# x #) -> f x
+
+read :: MutableVector s a -> Int -> ST s a
+{-# INLINE read #-}
+read (MutableVector arr#) (I# n#) = ST $ readArray# arr# n#
+
+write :: MutableVector s a -> Int -> a -> ST s ()
+{-# INLINE write #-}
+write (MutableVector arr#) (I# n#) x = ST $ \s# ->
+  case writeArray# arr# n# x s# of s2# -> (# s2#, () #)
+
index 0a8ffc2..581d040 100644 (file)
@@ -12,12 +12,16 @@ Build-Type:     Simple
 
 Library
   Exposed-Modules:
 
 Library
   Exposed-Modules:
+        Data.Vector.Stream
+
         Data.Vector.Unboxed.Unbox
         Data.Vector.Unboxed.Prim
         Data.Vector.Unboxed.Mutable
         Data.Vector.Unboxed
         Data.Vector.Unboxed.Unbox
         Data.Vector.Unboxed.Prim
         Data.Vector.Unboxed.Mutable
         Data.Vector.Unboxed
-        Data.Vector.Stream
 
 
+        Data.Vector.Prim
+        Data.Vector.Mutable
+        Data.Vector
   Include-Dirs:
         include
 
   Include-Dirs:
         include