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