6686eaf2c8caddb040bc3dd15402e986a5058603
[darcs-mirrors/vector.git] / Data / Vector / Unboxed.hs
1 {-# LANGUAGE RankNTypes #-}
2
3 module Data.Vector.Unboxed
4 where
5
6 import Data.Vector.Unboxed.Unbox
7 import qualified Data.Vector.Unboxed.Mutable as Mut
8
9 import qualified Data.Vector.Stream as Stream
10 import Data.Vector.Stream ( Step(..), Stream(..) )
11
12 import Control.Exception ( assert )
13 import Control.Monad.ST ( ST, runST )
14
15 import Prelude hiding ( length )
16
17 data Vector a = Vector !Int
18 !Int
19 !(Array a)
20
21 new :: Unbox a
22 => Int -> (forall s. Mut.Vector s a -> ST s (Mut.Vector s a)) -> Vector a
23 {-# INLINE new #-}
24 new n init = runST (
25 do
26 mv <- Mut.new n
27 mv' <- init mv
28 let (marr, i, n') = Mut.dataOf mv'
29 arr <- unsafeFreezeArray marr
30 return $ Vector i n' arr
31 )
32
33 stream :: Unbox a => Vector a -> Stream a
34 {-# INLINE_STREAM stream #-}
35 stream (Vector i n arr) = Stream get i n
36 where
37 n' = n+i
38
39 {-# INLINE get #-}
40 get j | j < n' = Yield (indexArray arr j) (j+1)
41 | otherwise = Done
42
43 unstream :: Unbox a => Stream a -> Vector a
44 {-# INLINE_STREAM unstream #-}
45 unstream s@(Stream _ _ n) = new n (\mv ->
46 do
47 n' <- Mut.fill mv s
48 return $ Mut.slice mv 0 n'
49 )
50
51 {-# RULES
52
53 "stream/unstream [Vector.Unboxed]" forall s.
54 stream (unstream s) = s
55
56 #-}
57
58 length :: Unbox a => Vector a -> Int
59 {-# INLINE length #-}
60 length (Vector _ n _) = n
61
62 slice :: Unbox a => Vector a -> Int -> Int -> Vector a
63 {-# INLINE slice #-}
64 slice (Vector i n arr) j m
65 = assert (j + m <= n && j >= 0 && m >= 0)
66 $ Vector (i+j) m arr
67
68 unsafeAt :: Unbox a => Vector a -> Int -> a
69 {-# INLINE unsafeAt #-}
70 unsafeAt (Vector i _ arr) j = indexArray arr (i+j)
71
72 at :: Unbox a => Vector a -> Int -> a
73 {-# INLINE at #-}
74 at v i = assert (i >= 0 && i < length v)
75 $ unsafeAt v i
76
77
78 map :: (Unbox a, Unbox b) => (a -> b) -> Vector a -> Vector b
79 {-# INLINE map #-}
80 map f = unstream . Stream.map f . stream
81
82 filter :: Unbox a => (a -> Bool) -> Vector a -> Vector a
83 {-# INLINE filter #-}
84 filter f = unstream . Stream.filter f . stream
85
86 zipWith :: (Unbox a, Unbox b, Unbox c)
87 => (a -> b -> c) -> Vector a -> Vector b -> Vector c
88 {-# INLINE zipWith #-}
89 zipWith f v w = unstream
90 $ Stream.zipWith f (stream v) (stream w)
91