GHC/Event/Array.hs

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