Update base for new Safe Haskell design
[packages/base.git] / GHC / Event / Array.hs
1 {-# LANGUAGE Unsafe #-}
2 {-# LANGUAGE BangPatterns, CPP, ForeignFunctionInterface, NoImplicitPrelude #-}
3
4 module GHC.Event.Array
5 (
6 Array
7 , capacity
8 , clear
9 , concat
10 , copy
11 , duplicate
12 , empty
13 , ensureCapacity
14 , findIndex
15 , forM_
16 , length
17 , loop
18 , new
19 , removeAt
20 , snoc
21 , unsafeLoad
22 , unsafeRead
23 , unsafeWrite
24 , useAsPtr
25 ) where
26
27 import Control.Monad hiding (forM_)
28 import Data.Bits ((.|.), shiftR)
29 import Data.IORef (IORef, atomicModifyIORef, newIORef, readIORef, writeIORef)
30 import Data.Maybe
31 import Foreign.C.Types (CSize(..))
32 import Foreign.ForeignPtr (ForeignPtr, withForeignPtr)
33 import Foreign.Ptr (Ptr, nullPtr, plusPtr)
34 import Foreign.Storable (Storable(..))
35 import GHC.Base
36 import GHC.Err (undefined)
37 import GHC.ForeignPtr (mallocPlainForeignPtrBytes, newForeignPtr_)
38 import GHC.Num (Num(..))
39 import GHC.Real (fromIntegral)
40 import GHC.Show (show)
41
42 #include "MachDeps.h"
43
44 #define BOUNDS_CHECKING 1
45
46 #if defined(BOUNDS_CHECKING)
47 -- This fugly hack is brought by GHC's apparent reluctance to deal
48 -- with MagicHash and UnboxedTuples when inferring types. Eek!
49 #define CHECK_BOUNDS(_func_,_len_,_k_) \
50 if (_k_) < 0 || (_k_) >= (_len_) then error ("GHC.Event.Array." ++ (_func_) ++ ": bounds error, index " ++ show (_k_) ++ ", capacity " ++ show (_len_)) else
51 #else
52 #define CHECK_BOUNDS(_func_,_len_,_k_)
53 #endif
54
55 -- Invariant: size <= capacity
56 newtype Array a = Array (IORef (AC a))
57
58 -- The actual array content.
59 data AC a = AC
60 !(ForeignPtr a) -- Elements
61 !Int -- Number of elements (length)
62 !Int -- Maximum number of elements (capacity)
63
64 empty :: IO (Array a)
65 empty = do
66 p <- newForeignPtr_ nullPtr
67 Array `fmap` newIORef (AC p 0 0)
68
69 allocArray :: Storable a => Int -> IO (ForeignPtr a)
70 allocArray n = allocHack undefined
71 where
72 allocHack :: Storable a => a -> IO (ForeignPtr a)
73 allocHack dummy = mallocPlainForeignPtrBytes (n * sizeOf dummy)
74
75 reallocArray :: Storable a => ForeignPtr a -> Int -> Int -> IO (ForeignPtr a)
76 reallocArray p newSize oldSize = reallocHack undefined p
77 where
78 reallocHack :: Storable a => a -> ForeignPtr a -> IO (ForeignPtr a)
79 reallocHack dummy src = do
80 let size = sizeOf dummy
81 dst <- mallocPlainForeignPtrBytes (newSize * size)
82 withForeignPtr src $ \s ->
83 when (s /= nullPtr && oldSize > 0) .
84 withForeignPtr dst $ \d -> do
85 _ <- memcpy d s (fromIntegral (oldSize * size))
86 return ()
87 return dst
88
89 new :: Storable a => Int -> IO (Array a)
90 new c = do
91 es <- allocArray cap
92 fmap Array (newIORef (AC es 0 cap))
93 where
94 cap = firstPowerOf2 c
95
96 duplicate :: Storable a => Array a -> IO (Array a)
97 duplicate a = dupHack undefined a
98 where
99 dupHack :: Storable b => b -> Array b -> IO (Array b)
100 dupHack dummy (Array ref) = do
101 AC es len cap <- readIORef ref
102 ary <- allocArray cap
103 withForeignPtr ary $ \dest ->
104 withForeignPtr es $ \src -> do
105 _ <- memcpy dest src (fromIntegral (len * sizeOf dummy))
106 return ()
107 Array `fmap` newIORef (AC ary len cap)
108
109 length :: Array a -> IO Int
110 length (Array ref) = do
111 AC _ len _ <- readIORef ref
112 return len
113
114 capacity :: Array a -> IO Int
115 capacity (Array ref) = do
116 AC _ _ cap <- readIORef ref
117 return cap
118
119 unsafeRead :: Storable a => Array a -> Int -> IO a
120 unsafeRead (Array ref) ix = do
121 AC es _ cap <- readIORef ref
122 CHECK_BOUNDS("unsafeRead",cap,ix)
123 withForeignPtr es $ \p ->
124 peekElemOff p ix
125
126 unsafeWrite :: Storable a => Array a -> Int -> a -> IO ()
127 unsafeWrite (Array ref) ix a = do
128 ac <- readIORef ref
129 unsafeWrite' ac ix a
130
131 unsafeWrite' :: Storable a => AC a -> Int -> a -> IO ()
132 unsafeWrite' (AC es _ cap) ix a = do
133 CHECK_BOUNDS("unsafeWrite'",cap,ix)
134 withForeignPtr es $ \p ->
135 pokeElemOff p ix a
136
137 unsafeLoad :: Storable a => Array a -> (Ptr a -> Int -> IO Int) -> IO Int
138 unsafeLoad (Array ref) load = do
139 AC es _ cap <- readIORef ref
140 len' <- withForeignPtr es $ \p -> load p cap
141 writeIORef ref (AC es len' cap)
142 return len'
143
144 ensureCapacity :: Storable a => Array a -> Int -> IO ()
145 ensureCapacity (Array ref) c = do
146 ac@(AC _ _ cap) <- readIORef ref
147 ac'@(AC _ _ cap') <- ensureCapacity' ac c
148 when (cap' /= cap) $
149 writeIORef ref ac'
150
151 ensureCapacity' :: Storable a => AC a -> Int -> IO (AC a)
152 ensureCapacity' ac@(AC es len cap) c = do
153 if c > cap
154 then do
155 es' <- reallocArray es cap' cap
156 return (AC es' len cap')
157 else
158 return ac
159 where
160 cap' = firstPowerOf2 c
161
162 useAsPtr :: Array a -> (Ptr a -> Int -> IO b) -> IO b
163 useAsPtr (Array ref) f = do
164 AC es len _ <- readIORef ref
165 withForeignPtr es $ \p -> f p len
166
167 snoc :: Storable a => Array a -> a -> IO ()
168 snoc (Array ref) e = do
169 ac@(AC _ len _) <- readIORef ref
170 let len' = len + 1
171 ac'@(AC es _ cap) <- ensureCapacity' ac len'
172 unsafeWrite' ac' len e
173 writeIORef ref (AC es len' cap)
174
175 clear :: Storable a => Array a -> IO ()
176 clear (Array ref) = do
177 !_ <- atomicModifyIORef ref $ \(AC es _ cap) ->
178 let e = AC es 0 cap in (e, e)
179 return ()
180
181 forM_ :: Storable a => Array a -> (a -> IO ()) -> IO ()
182 forM_ ary g = forHack ary g undefined
183 where
184 forHack :: Storable b => Array b -> (b -> IO ()) -> b -> IO ()
185 forHack (Array ref) f dummy = do
186 AC es len _ <- readIORef ref
187 let size = sizeOf dummy
188 offset = len * size
189 withForeignPtr es $ \p -> do
190 let go n | n >= offset = return ()
191 | otherwise = do
192 f =<< peek (p `plusPtr` n)
193 go (n + size)
194 go 0
195
196 loop :: Storable a => Array a -> b -> (b -> a -> IO (b,Bool)) -> IO ()
197 loop ary z g = loopHack ary z g undefined
198 where
199 loopHack :: Storable b => Array b -> c -> (c -> b -> IO (c,Bool)) -> b
200 -> IO ()
201 loopHack (Array ref) y f dummy = do
202 AC es len _ <- readIORef ref
203 let size = sizeOf dummy
204 offset = len * size
205 withForeignPtr es $ \p -> do
206 let go n k
207 | n >= offset = return ()
208 | otherwise = do
209 (k',cont) <- f k =<< peek (p `plusPtr` n)
210 when cont $ go (n + size) k'
211 go 0 y
212
213 findIndex :: Storable a => (a -> Bool) -> Array a -> IO (Maybe (Int,a))
214 findIndex = findHack undefined
215 where
216 findHack :: Storable b => b -> (b -> Bool) -> Array b -> IO (Maybe (Int,b))
217 findHack dummy p (Array ref) = do
218 AC es len _ <- readIORef ref
219 let size = sizeOf dummy
220 offset = len * size
221 withForeignPtr es $ \ptr ->
222 let go !n !i
223 | n >= offset = return Nothing
224 | otherwise = do
225 val <- peek (ptr `plusPtr` n)
226 if p val
227 then return $ Just (i, val)
228 else go (n + size) (i + 1)
229 in go 0 0
230
231 concat :: Storable a => Array a -> Array a -> IO ()
232 concat (Array d) (Array s) = do
233 da@(AC _ dlen _) <- readIORef d
234 sa@(AC _ slen _) <- readIORef s
235 writeIORef d =<< copy' da dlen sa 0 slen
236
237 -- | Copy part of the source array into the destination array. The
238 -- destination array is resized if not large enough.
239 copy :: Storable a => Array a -> Int -> Array a -> Int -> Int -> IO ()
240 copy (Array d) dstart (Array s) sstart maxCount = do
241 da <- readIORef d
242 sa <- readIORef s
243 writeIORef d =<< copy' da dstart sa sstart maxCount
244
245 -- | Copy part of the source array into the destination array. The
246 -- destination array is resized if not large enough.
247 copy' :: Storable a => AC a -> Int -> AC a -> Int -> Int -> IO (AC a)
248 copy' d dstart s sstart maxCount = copyHack d s undefined
249 where
250 copyHack :: Storable b => AC b -> AC b -> b -> IO (AC b)
251 copyHack dac@(AC _ oldLen _) (AC src slen _) dummy = do
252 when (maxCount < 0 || dstart < 0 || dstart > oldLen || sstart < 0 ||
253 sstart > slen) $ error "copy: bad offsets or lengths"
254 let size = sizeOf dummy
255 count = min maxCount (slen - sstart)
256 if count == 0
257 then return dac
258 else do
259 AC dst dlen dcap <- ensureCapacity' dac (dstart + count)
260 withForeignPtr dst $ \dptr ->
261 withForeignPtr src $ \sptr -> do
262 _ <- memcpy (dptr `plusPtr` (dstart * size))
263 (sptr `plusPtr` (sstart * size))
264 (fromIntegral (count * size))
265 return $ AC dst (max dlen (dstart + count)) dcap
266
267 removeAt :: Storable a => Array a -> Int -> IO ()
268 removeAt a i = removeHack a undefined
269 where
270 removeHack :: Storable b => Array b -> b -> IO ()
271 removeHack (Array ary) dummy = do
272 AC fp oldLen cap <- readIORef ary
273 when (i < 0 || i >= oldLen) $ error "removeAt: invalid index"
274 let size = sizeOf dummy
275 newLen = oldLen - 1
276 when (newLen > 0 && i < newLen) .
277 withForeignPtr fp $ \ptr -> do
278 _ <- memmove (ptr `plusPtr` (size * i))
279 (ptr `plusPtr` (size * (i+1)))
280 (fromIntegral (size * (newLen-i)))
281 return ()
282 writeIORef ary (AC fp newLen cap)
283
284 {-The firstPowerOf2 function works by setting all bits on the right-hand
285 side of the most significant flagged bit to 1, and then incrementing
286 the entire value at the end so it "rolls over" to the nearest power of
287 two.
288 -}
289
290 -- | Computes the next-highest power of two for a particular integer,
291 -- @n@. If @n@ is already a power of two, returns @n@. If @n@ is
292 -- zero, returns zero, even though zero is not a power of two.
293 firstPowerOf2 :: Int -> Int
294 firstPowerOf2 !n =
295 let !n1 = n - 1
296 !n2 = n1 .|. (n1 `shiftR` 1)
297 !n3 = n2 .|. (n2 `shiftR` 2)
298 !n4 = n3 .|. (n3 `shiftR` 4)
299 !n5 = n4 .|. (n4 `shiftR` 8)
300 !n6 = n5 .|. (n5 `shiftR` 16)
301 #if WORD_SIZE_IN_BITS == 32
302 in n6 + 1
303 #elif WORD_SIZE_IN_BITS == 64
304 !n7 = n6 .|. (n6 `shiftR` 32)
305 in n7 + 1
306 #else
307 # error firstPowerOf2 not defined on this architecture
308 #endif
309
310 foreign import ccall unsafe "string.h memcpy"
311 memcpy :: Ptr a -> Ptr a -> CSize -> IO (Ptr a)
312
313 foreign import ccall unsafe "string.h memmove"
314 memmove :: Ptr a -> Ptr a -> CSize -> IO (Ptr a)