bca78fab2e2c8b544fa7852f238bb36981ad1063
[darcs-mirrors/vector.git] / Data / Vector / Unboxed / Mutable.hs
1 {-# LANGUAGE BangPatterns #-}
2
3 module Data.Vector.Unboxed.Mutable (
4 Vector,
5
6 new, length, slice, read, write, fill, fillIndexed,
7 dataOf
8 ) where
9
10 import Data.Vector.Unboxed.Unbox
11
12 import qualified Data.Vector.Stream as Stream
13 import Data.Vector.Stream ( Stream )
14
15 import Control.Exception ( assert )
16 import Control.Monad.ST ( ST )
17
18 import Prelude hiding ( length, read )
19
20 data Vector s a = Vector !Int -- ^ start
21 !Int -- ^ length
22 !(MutableArray s a) -- ^ data
23
24 dataOf :: Vector s a -> (MutableArray s a, Int, Int)
25 {-# INLINE dataOf #-}
26 dataOf (Vector i n arr) = (arr, i, n)
27
28 new :: Unbox a => Int -> ST s (Vector s a)
29 {-# INLINE new #-}
30 new n = assert (n >= 0)
31 $ do
32 arr <- newArray n
33 return $ Vector 0 n arr
34
35 length :: Unbox a => Vector s a -> Int
36 {-# INLINE length #-}
37 length (Vector _ n _) = n
38
39 slice :: Unbox a => Vector s a -> Int -> Int -> Vector s a
40 {-# INLINE slice #-}
41 slice (Vector i n arr) j m
42 = assert (j + m <= n && j >= 0 && m >= 0)
43 $ Vector (i+j) m arr
44
45 slicel :: Unbox a => Vector s a -> Int -> Vector s a
46 {-# INLINE slicel #-}
47 slicel (Vector i n arr) m
48 = assert (m <= n && m >= 0)
49 $ Vector i m arr
50
51 read :: Unbox a => Vector s a -> Int -> ST s a
52 {-# INLINE read #-}
53 read (Vector i n arr) j
54 = assert (j < n)
55 $ readArray arr (i+j)
56
57 write :: Unbox a => Vector s a -> Int -> a -> ST s ()
58 {-# INLINE write #-}
59 write (Vector i n arr) j x
60 = assert (j < n)
61 $ writeArray arr (i+j) x
62
63 fill :: Unbox a => Vector s a -> Stream a -> ST s Int
64 {-# INLINE fill #-}
65 fill !v s = Stream.foldM put 0 s
66 where
67 {-# INLINE put #-}
68 put i x = do { write v i x; return (i+1) }
69
70 fillIndexed :: Unbox a => Vector s a -> Stream (Int, a) -> ST s ()
71 {-# INLINE fillIndexed #-}
72 fillIndexed !v s = Stream.mapM_ put s
73 where
74 {-# INLINE put #-}
75 put (i,x) = write v i x
76
77 copyTo :: Unbox a => Vector s a -> Vector s a -> ST s ()
78 {-# INLINE copyTo #-}
79 copyTo !v !w = assert (length v == length w)
80 $ copy_loop 0
81 where
82 n = length v
83
84 copy_loop i | i < n = do
85 x <- read v i
86 write w i x
87 copy_loop (i+1)
88 | otherwise = return ()
89
90 clone :: Unbox a => Vector s a -> ST s (Vector s a)
91 {-# INLINE clone #-}
92 clone v = do
93 w <- new (length v)
94 v `copyTo` w
95 return w
96