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