System/Event/Array.hs

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