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