Initial revision
authorRoman Leshchinskiy <rl@cse.unsw.edu.au>
Sat, 5 Jul 2008 01:23:54 +0000 (01:23 +0000)
committerRoman Leshchinskiy <rl@cse.unsw.edu.au>
Sat, 5 Jul 2008 01:23:54 +0000 (01:23 +0000)
Data/Vector/Stream.hs [new file with mode: 0644]
Data/Vector/Unboxed.hs [new file with mode: 0644]
Data/Vector/Unboxed/Mutable.hs [new file with mode: 0644]
Data/Vector/Unboxed/Unbox.hs [new file with mode: 0644]
include/phases.h [new file with mode: 0644]

diff --git a/Data/Vector/Stream.hs b/Data/Vector/Stream.hs
new file mode 100644 (file)
index 0000000..f25e2d5
--- /dev/null
@@ -0,0 +1,123 @@
+{-# LANGUAGE ExistentialQuantification, BangPatterns #-}
+
+module Data.Vector.Stream (
+  Step(..), Stream(..),
+
+  empty, singleton, replicate,
+  map, filter, zipWith,
+  foldr, foldl, foldl',
+  mapM_, foldM
+) where
+
+import Prelude hiding ( replicate, map, filter, zipWith,
+                        foldr, foldl,
+                        mapM_ )
+
+data Step s a = Yield a s
+              | Skip    s
+              | Done
+
+data Stream a = forall s. Stream (s -> Step s a) s Int
+
+empty :: Stream a
+{-# INLINE_STREAM empty #-}
+empty = Stream (const Done) () 0
+
+singleton :: a -> Stream a
+{-# INLINE_STREAM singleton #-}
+singleton x = Stream step True 1
+  where
+    {-# INLINE step #-}
+    step True  = Yield x False
+    step False = Done
+
+replicate :: Int -> a -> Stream a
+{-# INLINE_STREAM replicate #-}
+replicate n x = Stream step n (max n 0)
+  where
+    {-# INLINE step #-}
+    step i | i > 0     = Yield x (i-1)
+           | otherwise = Done
+
+map :: (a -> b) -> Stream a -> Stream b
+{-# INLINE_STREAM map #-}
+map f (Stream step s n) = Stream step' s n
+  where
+    {-# INLINE step' #-}
+    step' s = case step s of
+                Yield x s' -> Yield (f x) s'
+                Skip    s' -> Skip        s'
+                Done       -> Done
+
+filter :: (a -> Bool) -> Stream a -> Stream a
+{-# INLINE_STREAM filter #-}
+filter f (Stream step s n) = Stream step' s n
+  where
+    {-# INLINE step' #-}
+    step' s = case step s of
+                Yield x s' | f x       -> Yield x s'
+                           | otherwise -> Skip    s'
+                Skip    s'             -> Skip    s'
+                Done                   -> Done
+
+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)
+  where
+    {-# INLINE step #-}
+    step (sa, sb, Nothing) = case stepa sa of
+                               Yield x sa' -> Skip (sa', sb, Just x)
+                               Skip    sa' -> Skip (sa', sb, Nothing)
+                               Done        -> Done
+
+    step (sa, sb, Just x)  = case stepb sb of
+                               Yield y sb' -> Yield (f x y) (sa, sb', Nothing)
+                               Skip    sb' -> Skip          (sa, sb', Just x)
+                               Done        -> Done
+
+foldl :: (a -> b -> b) -> b -> Stream a -> b
+{-# INLINE_STREAM foldl #-}
+foldl f z (Stream step s _) = foldl_go z s
+  where
+    foldl_go z s = case step s of
+                     Yield x s' -> foldl_go (f x z) s'
+                     Skip    s' -> foldl_go z       s'
+                     Done       -> z
+
+foldl' :: (a -> b -> b) -> b -> Stream a -> b
+{-# INLINE_STREAM foldl' #-}
+foldl' f z (Stream step s _) = foldl_go z s
+  where
+    foldl_go !z s = case step s of
+                      Yield x s' -> foldl_go (f x z) s'
+                      Skip    s' -> foldl_go z       s'
+                      Done       -> z
+
+foldr :: (a -> b -> b) -> b -> Stream a -> b
+{-# INLINE_STREAM foldr #-}
+foldr f z (Stream step s _) = foldr_go s
+  where
+    foldr_go s = case step s of
+                   Yield x s' -> f x (foldr_go s')
+                   Skip    s' -> foldr_go s'
+                   Done       -> z
+
+mapM_ :: Monad m => (a -> m ()) -> Stream a -> m ()
+{-# INLINE_STREAM mapM_ #-}
+mapM_ m (Stream step s _) = mapM_go s
+   where
+     mapM_go s = case step s of
+                   Yield x s' -> do { m x; mapM_go s' }
+                   Skip    s' -> mapM_go s'
+                   Done       -> return ()
+
+foldM :: Monad m => (a -> b -> m a) -> a -> Stream b -> m a
+{-# INLINE_STREAM foldM #-}
+foldM m z (Stream step s _) = foldM_go z s
+  where
+    foldM_go z s = case step s of
+                     Yield x s' -> do { z' <- m z x; foldM_go z' s' }
+                     Skip    s' -> foldM_go z s'
+                     Done       -> return z
+
diff --git a/Data/Vector/Unboxed.hs b/Data/Vector/Unboxed.hs
new file mode 100644 (file)
index 0000000..6686eaf
--- /dev/null
@@ -0,0 +1,91 @@
+{-# LANGUAGE RankNTypes #-}
+
+module Data.Vector.Unboxed
+where
+
+import Data.Vector.Unboxed.Unbox
+import qualified Data.Vector.Unboxed.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 !Int
+                       !Int
+                       !(Array a)
+
+new :: Unbox a
+    => 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 (marr, i, n') = Mut.dataOf mv'
+    arr <- unsafeFreezeArray marr
+    return $ Vector i n' arr
+  )
+
+stream :: Unbox a => 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'    = Yield (indexArray arr j) (j+1)
+          | otherwise = Done
+
+unstream :: Unbox a => 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.Unboxed]" forall s.
+  stream (unstream s) = s
+
+ #-}
+
+length :: Unbox a => Vector a -> Int
+{-# INLINE length #-}
+length (Vector _ n _) = n
+
+slice :: Unbox a => 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 :: Unbox a => Vector a -> Int -> a
+{-# INLINE unsafeAt #-}
+unsafeAt (Vector i _ arr) j = indexArray arr (i+j)
+
+at :: Unbox a => Vector a -> Int -> a
+{-# INLINE at #-}
+at v i = assert (i >= 0 && i < length v)
+       $ unsafeAt v i
+
+
+map :: (Unbox a, Unbox b) => (a -> b) -> Vector a -> Vector b
+{-# INLINE map #-}
+map f = unstream . Stream.map f . stream
+
+filter :: Unbox a => (a -> Bool) -> Vector a -> Vector a
+{-# INLINE filter #-}
+filter f = unstream . Stream.filter f . stream
+
+zipWith :: (Unbox a, Unbox b, Unbox c)
+        => (a -> b -> c) -> Vector a -> Vector b -> Vector c
+{-# INLINE zipWith #-}
+zipWith f v w = unstream
+              $ Stream.zipWith f (stream v) (stream w)
+
diff --git a/Data/Vector/Unboxed/Mutable.hs b/Data/Vector/Unboxed/Mutable.hs
new file mode 100644 (file)
index 0000000..bca78fa
--- /dev/null
@@ -0,0 +1,96 @@
+{-# LANGUAGE BangPatterns #-}
+
+module Data.Vector.Unboxed.Mutable (
+  Vector,
+
+  new, length, slice, read, write, fill, fillIndexed,
+  dataOf
+) where
+
+import Data.Vector.Unboxed.Unbox
+
+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 !Int                  -- ^ start
+                         !Int                  -- ^ length
+                         !(MutableArray s a)   -- ^ data
+
+dataOf :: Vector s a -> (MutableArray s a, Int, Int)
+{-# INLINE dataOf #-}
+dataOf (Vector i n arr) = (arr, i, n)
+
+new :: Unbox a => Int -> ST s (Vector s a)
+{-# INLINE new #-}
+new n = assert (n >= 0)
+      $ do
+          arr <- newArray n
+          return $ Vector 0 n arr
+
+length :: Unbox a => Vector s a -> Int
+{-# INLINE length #-}
+length (Vector _ n _) = n
+
+slice :: Unbox a => Vector s a -> Int -> Int -> Vector s a
+{-# INLINE slice #-}
+slice (Vector i n arr) j m
+  = assert (j + m <= n && j >= 0 && m >= 0)
+  $ Vector (i+j) m arr
+
+slicel :: Unbox a => Vector s a -> Int -> Vector s a
+{-# INLINE slicel #-}
+slicel (Vector i n arr) m
+  = assert (m <= n && m >= 0)
+  $ Vector i m arr
+
+read :: Unbox a => Vector s a -> Int -> ST s a
+{-# INLINE read #-}
+read (Vector i n arr) j
+  = assert (j < n)
+  $ readArray arr (i+j)
+
+write :: Unbox a => Vector s a -> Int -> a -> ST s ()
+{-# INLINE write #-}
+write (Vector i n arr) j x
+  = assert (j < n)
+  $ writeArray arr (i+j) x
+
+fill :: Unbox a => 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) }
+
+fillIndexed :: Unbox a => Vector s a -> Stream (Int, a) -> ST s ()
+{-# INLINE fillIndexed #-}
+fillIndexed !v s = Stream.mapM_ put s
+  where
+    {-# INLINE put #-}
+    put (i,x) = write v i x
+
+copyTo :: Unbox a => Vector s a -> Vector s a -> ST s ()
+{-# INLINE copyTo #-}
+copyTo !v !w = assert (length v == length w)
+             $ copy_loop 0
+  where
+    n = length v
+
+    copy_loop i | i < n     = do
+                                x <- read v i
+                                write w i x
+                                copy_loop (i+1)
+                | otherwise = return ()
+
+clone :: Unbox a => Vector s a -> ST s (Vector s a)
+{-# INLINE clone #-}
+clone v = do
+            w <- new (length v)
+            v `copyTo` w
+            return w
+
diff --git a/Data/Vector/Unboxed/Unbox.hs b/Data/Vector/Unboxed/Unbox.hs
new file mode 100644 (file)
index 0000000..a4c277b
--- /dev/null
@@ -0,0 +1,68 @@
+{-# LANGUAGE MagicHash, UnboxedTuples, ScopedTypeVariables #-}
+
+module Data.Vector.Unboxed.Unbox (
+  Unbox(..), Array, MutableArray,
+
+  arraySize, newArray, unsafeFreezeArray, indexArray, readArray, writeArray
+) where
+
+import GHC.Prim (
+    ByteArray#, MutableByteArray#, State#,
+    newByteArray#, unsafeFreezeByteArray#,
+
+    Int#, indexIntArray#, readIntArray#, writeIntArray#
+  )
+import GHC.ST (
+    ST(..)
+  )
+import GHC.Base (
+    Int(..)
+  )
+import Data.Array.Base (
+    wORD_SCALE )
+
+data Array          a = Array        ByteArray#
+data MutableArray s a = MutableArray (MutableByteArray# s)
+
+class Unbox a where
+  arraySize#  :: a -> Int# -> Int#
+  indexArray# :: ByteArray# -> Int# -> a
+  readArray#  :: MutableByteArray# s -> Int# -> State# s -> (# State# s, a #)
+  writeArray# :: MutableByteArray# s -> Int# -> a -> State# s -> State# s
+
+instance Unbox Int where
+  arraySize#  _                  = wORD_SCALE
+  indexArray# arr# i#            = I# (indexIntArray# arr# i#)
+  readArray#  arr# i# s#         = case readIntArray# arr# i# s# of
+                                     (# s1#, n# #) -> (# s1#, I# n# #)
+  writeArray# arr# i# (I# n#) s# = writeIntArray# arr# i# n# s#
+
+arraySize :: Unbox a => a -> Int -> Int
+{-# INLINE arraySize #-}
+arraySize a (I# i#) = I# (arraySize# a i#)
+
+newArray :: forall s a. Unbox a => Int -> ST s (MutableArray s a)
+{-# INLINE newArray #-}
+newArray (I# n#) = ST $ \s# ->
+  case newByteArray# (arraySize# (undefined :: a) n#) s# of
+    (# s2#, arr# #) -> (# s2#, MutableArray arr# #)
+
+unsafeFreezeArray :: Unbox a => MutableArray s a -> ST s (Array a)
+{-# INLINE unsafeFreezeArray #-}
+unsafeFreezeArray (MutableArray arr#) = ST $ \s# ->
+  case unsafeFreezeByteArray# arr# s# of
+    (# s2, frozen# #) -> (# s2, Array frozen# #)
+
+indexArray :: Unbox a => Array a -> Int -> a
+{-# INLINE indexArray #-}
+indexArray (Array arr#) (I# i#) = indexArray# arr# i#
+
+readArray :: Unbox a => MutableArray s a -> Int -> ST s a
+{-# INLINE readArray #-}
+readArray (MutableArray arr#) (I# i#) = ST $ readArray# arr# i#
+
+writeArray :: Unbox a => MutableArray s a -> Int -> a -> ST s ()
+{-# INLINE writeArray #-}
+writeArray (MutableArray arr#) (I# i#) x = ST $ \s# ->
+  case writeArray# arr# i# x s# of s2# -> (# s2#, () #)
+
diff --git a/include/phases.h b/include/phases.h
new file mode 100644 (file)
index 0000000..be72156
--- /dev/null
@@ -0,0 +1,2 @@
+#define INLINE_STREAM INLINE [1]
+