015a397e8f5437f83fed3cda54568ac16eaaa4a4
[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 qualified Data.Vector.Stream as Stream
13 import Data.Vector.Stream ( Step(..), 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 :: Unbox a
25 => Int -> (forall s. Mut.Vector s a -> ST s (Mut.Vector s a)) -> Vector a
26 {-# INLINE new #-}
27 new n init = runST (
28 do
29 mv <- Mut.new n
30 mv' <- init mv
31 let (mprim, i, n') = Mut.dataOf mv'
32 prim <- Prim.unsafeFreeze mprim
33 return $ Vector i n' prim
34 )
35
36 stream :: Unbox a => Vector a -> Stream a
37 {-# INLINE_STREAM stream #-}
38 stream (Vector i n arr) = Stream get i n
39 where
40 n' = n+i
41
42 {-# INLINE get #-}
43 get j | j < n' = Yield (Prim.at arr j) (j+1)
44 | otherwise = Done
45
46 unstream :: Unbox a => Stream a -> Vector a
47 {-# INLINE_STREAM unstream #-}
48 unstream s@(Stream _ _ n) = new n (\mv ->
49 do
50 n' <- Mut.fill mv s
51 return $ Mut.slice mv 0 n'
52 )
53
54 {-# RULES
55
56 "stream/unstream [Vector.Unboxed]" forall s.
57 stream (unstream s) = s
58
59 #-}
60
61 length :: Unbox a => Vector a -> Int
62 {-# INLINE length #-}
63 length (Vector _ n _) = n
64
65 slice :: Unbox a => Vector a -> Int -> Int -> Vector a
66 {-# INLINE slice #-}
67 slice (Vector i n arr) j m
68 = assert (j + m <= n && j >= 0 && m >= 0)
69 $ Vector (i+j) m arr
70
71 unsafeAt :: Unbox a => Vector a -> Int -> a
72 {-# INLINE unsafeAt #-}
73 unsafeAt (Vector i _ arr) j = Prim.at arr (i+j)
74
75 at :: Unbox a => Vector a -> Int -> a
76 {-# INLINE at #-}
77 at v i = assert (i >= 0 && i < length v)
78 $ unsafeAt v i
79
80 infixr ++
81 (++) :: Unbox a => Vector a -> Vector a -> Vector a
82 {-# INLINE (++) #-}
83 v ++ w = unstream (stream v Stream.++ stream w)
84
85 map :: (Unbox a, Unbox b) => (a -> b) -> Vector a -> Vector b
86 {-# INLINE map #-}
87 map f = unstream . Stream.map f . stream
88
89 filter :: Unbox a => (a -> Bool) -> Vector a -> Vector a
90 {-# INLINE filter #-}
91 filter f = unstream . Stream.filter f . stream
92
93 zipWith :: (Unbox a, Unbox b, Unbox c)
94 => (a -> b -> c) -> Vector a -> Vector b -> Vector c
95 {-# INLINE zipWith #-}
96 zipWith f v w = unstream
97 $ Stream.zipWith f (stream v) (stream w)
98
99 foldl' :: Unbox a => (a -> b -> b) -> b -> Vector a -> b
100 {-# INLINE foldl' #-}
101 foldl' f z = Stream.foldl' f z . stream
102