9b0bfb90d32247b2d59af0246c31394f3e986ddc
[darcs-mirrors/vector.git] / Data / Vector.hs
1 {-# LANGUAGE RankNTypes, CPP #-}
2
3 #include "phases.h"
4
5 module Data.Vector
6 where
7
8 import qualified Data.Vector.Prim as Prim
9 import qualified Data.Vector.Mutable as Mut
10
11 import qualified Data.Vector.Stream as Stream
12 import Data.Vector.Stream ( Stream )
13
14 import Control.Exception ( assert )
15 import Control.Monad.ST ( ST, runST )
16
17 import Prelude hiding ( length, (++) )
18
19 data Vector a = Vector {-# UNPACK #-} !Int
20 {-# UNPACK #-} !Int
21 {-# UNPACK #-} !(Prim.Vector a)
22
23 new :: Int -> (forall s. Mut.Vector s a -> ST s (Mut.Vector s a)) -> Vector a
24 {-# INLINE new #-}
25 new n init = runST (
26 do
27 mv <- Mut.new n
28 mv' <- init mv
29 let (mprim, i, n') = Mut.dataOf mv'
30 prim <- Prim.unsafeFreeze mprim
31 return $ Vector i n' prim
32 )
33
34 new' :: Int -> a -> (forall s. Mut.Vector s a -> ST s (Mut.Vector s a)) -> Vector a
35 {-# INLINE new' #-}
36 new' n x init = runST (
37 do
38 mv <- Mut.new' n x
39 mv' <- init mv
40 let (mprim, i, n') = Mut.dataOf mv'
41 prim <- Prim.unsafeFreeze mprim
42 return $ Vector i n' prim
43 )
44
45 stream :: Vector a -> Stream a
46 {-# INLINE_STREAM stream #-}
47 stream (Vector i n arr) = Stream.unfold get i n
48 where
49 n' = n+i
50
51 {-# INLINE get #-}
52 get j | j < n' = Prim.at' arr j $ \x -> Just (x, j+1)
53 | otherwise = Nothing
54
55 unstream :: Stream a -> Vector a
56 {-# INLINE_STREAM unstream #-}
57 unstream s = new (Stream.bound s) (\mv ->
58 do
59 n <- Mut.fill mv s
60 return $ Mut.slice mv 0 n
61 )
62
63 {-# RULES
64
65 "stream/unstream [Vector]" forall s.
66 stream (unstream s) = s
67
68 #-}
69
70 length :: Vector a -> Int
71 {-# INLINE length #-}
72 length (Vector _ n _) = n
73
74 slice :: Vector a -> Int -> Int -> Vector a
75 {-# INLINE slice #-}
76 slice (Vector i n arr) j m
77 = assert (j + m <= n && j >= 0 && m >= 0)
78 $ Vector (i+j) m arr
79
80 unsafeAt ::Vector a -> Int -> a
81 {-# INLINE unsafeAt #-}
82 unsafeAt (Vector i _ arr) j = Prim.at arr (i+j)
83
84 at :: Vector a -> Int -> a
85 {-# INLINE at #-}
86 at v i = assert (i >= 0 && i < length v)
87 $ unsafeAt v i
88
89 infixr ++
90 (++) :: Vector a -> Vector a -> Vector a
91 {-# INLINE (++) #-}
92 v ++ w = unstream (stream v Stream.++ stream w)
93
94 map :: (a -> b) -> Vector a -> Vector b
95 {-# INLINE map #-}
96 map f = unstream . Stream.map f . stream
97
98 filter :: (a -> Bool) -> Vector a -> Vector a
99 {-# INLINE filter #-}
100 filter f = unstream . Stream.filter f . stream
101
102 zipWith :: (a -> b -> c) -> Vector a -> Vector b -> Vector c
103 {-# INLINE zipWith #-}
104 zipWith f v w = unstream
105 $ Stream.zipWith f (stream v) (stream w)
106
107 foldl' :: (a -> b -> b) -> b -> Vector a -> b
108 {-# INLINE foldl' #-}
109 foldl' f z = Stream.foldl' f z . stream
110