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