1 {-# LANGUAGE RankNTypes, CPP #-}
3 #include "phases.h"
5 module Data.Vector
6 where
8 import qualified Data.Vector.Prim as Prim
9 import qualified Data.Vector.Mutable as Mut
11 import qualified Data.Vector.Stream as Stream
12 import Data.Vector.Stream ( Stream )
14 import Control.Exception ( assert )
15 import Control.Monad.ST ( ST, runST )
17 import Prelude hiding ( length, (++) )
19 data Vector a = Vector {-# UNPACK #-} !Int
20 {-# UNPACK #-} !Int
21 {-# UNPACK #-} !(Prim.Vector a)
23 new :: 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 )
34 new' :: Int -> a -> (forall s. Mut.Vector s a -> ST s (Mut.Vector s a)) -> Vector a
35 {-# INLINE new' #-}
36 new' n x init = runST (
37 do
38 mv <- Mut.new' n x
39 mv' <- init mv
40 let (mprim, i, n') = Mut.dataOf mv'
41 prim <- Prim.unsafeFreeze mprim
42 return \$ Vector i n' prim
43 )
45 stream :: Vector a -> Stream a
46 {-# INLINE_STREAM stream #-}
47 stream (Vector i n arr) = Stream.unfold get i n
48 where
49 n' = n+i
51 {-# INLINE get #-}
52 get j | j < n' = Prim.at' arr j \$ \x -> Just (x, j+1)
53 | otherwise = Nothing
55 unstream :: Stream a -> Vector a
56 {-# INLINE_STREAM unstream #-}
57 unstream s = new (Stream.bound s) (\mv ->
58 do
59 n <- Mut.fill mv s
60 return \$ Mut.slice mv 0 n
61 )
63 {-# RULES
65 "stream/unstream [Vector]" forall s.
66 stream (unstream s) = s
68 #-}
70 length :: Vector a -> Int
71 {-# INLINE length #-}
72 length (Vector _ n _) = n
74 slice :: Vector a -> Int -> Int -> Vector a
75 {-# INLINE slice #-}
76 slice (Vector i n arr) j m
77 = assert (j + m <= n && j >= 0 && m >= 0)
78 \$ Vector (i+j) m arr
80 unsafeAt ::Vector a -> Int -> a
81 {-# INLINE unsafeAt #-}
82 unsafeAt (Vector i _ arr) j = Prim.at arr (i+j)
84 at :: Vector a -> Int -> a
85 {-# INLINE at #-}
86 at v i = assert (i >= 0 && i < length v)
87 \$ unsafeAt v i
89 infixr ++
90 (++) :: Vector a -> Vector a -> Vector a
91 {-# INLINE (++) #-}
92 v ++ w = unstream (stream v Stream.++ stream w)
94 map :: (a -> b) -> Vector a -> Vector b
95 {-# INLINE map #-}
96 map f = unstream . Stream.map f . stream
98 filter :: (a -> Bool) -> Vector a -> Vector a
99 {-# INLINE filter #-}
100 filter f = unstream . Stream.filter f . stream
102 zipWith :: (a -> b -> c) -> Vector a -> Vector b -> Vector c
103 {-# INLINE zipWith #-}
104 zipWith f v w = unstream
105 \$ Stream.zipWith f (stream v) (stream w)
107 foldl' :: (a -> b -> b) -> b -> Vector a -> b
108 {-# INLINE foldl' #-}
109 foldl' f z = Stream.foldl' f z . stream