GHC/Event/Array.hs

    {-# LANGUAGE Trustworthy #-}
    {-# LANGUAGE BangPatterns, CPP, NoImplicitPrelude #-}

    module GHC.Event.Array
        (
          Array
        , capacity
        , clear
        , concat
        , copy
        , duplicate
        , empty
        , ensureCapacity
        , findIndex
        , forM_
        , length
        , loop
        , new
        , removeAt
        , snoc
        , unsafeLoad
        , unsafeRead
        , unsafeWrite
        , useAsPtr
        ) where

    import Data.Bits ((.|.), shiftR)
    import Data.IORef (IORef, atomicModifyIORef', newIORef, readIORef, writeIORef)
    import Data.Maybe
    import Foreign.C.Types (CSize(..))
    import Foreign.ForeignPtr (ForeignPtr, withForeignPtr)
    import Foreign.Ptr (Ptr, nullPtr, plusPtr)
    import Foreign.Storable (Storable(..))
    import GHC.Base hiding (empty)
    import GHC.ForeignPtr (mallocPlainForeignPtrBytes, newForeignPtr_)
    import GHC.Num (Num(..))
    import GHC.Real (fromIntegral)
    import GHC.Show (show)

    #include "MachDeps.h"

    #define BOUNDS_CHECKING 1

    #if defined(BOUNDS_CHECKING)
    -- This fugly hack is brought by GHC's apparent reluctance to deal
    -- with MagicHash and UnboxedTuples when inferring types. Eek!
    #define CHECK_BOUNDS(_func_,_len_,_k_) \
    if (_k_) < 0 || (_k_) >= (_len_) then errorWithoutStackTrace ("GHC.Event.Array." ++ (_func_) ++ ": bounds error, index " ++ show (_k_) ++ ", capacity " ++ show (_len_)) else
    #else
    #define CHECK_BOUNDS(_func_,_len_,_k_)
    #endif

    -- Invariant: size <= capacity
    newtype Array a = Array (IORef (AC a))

    -- The actual array content.
    data AC a = AC
        !(ForeignPtr a)  -- Elements
        !Int      -- Number of elements (length)
        !Int      -- Maximum number of elements (capacity)

    empty :: IO (Array a)
    empty = do
      p <- newForeignPtr_ nullPtr
      Array `fmap` newIORef (AC p 0 0)

    allocArray :: Storable a => Int -> IO (ForeignPtr a)
    allocArray n = allocHack undefined
     where
      allocHack :: Storable a => a -> IO (ForeignPtr a)
      allocHack dummy = mallocPlainForeignPtrBytes (n * sizeOf dummy)

    reallocArray :: Storable a => ForeignPtr a -> Int -> Int -> IO (ForeignPtr a)
    reallocArray p newSize oldSize = reallocHack undefined p
     where
      reallocHack :: Storable a => a -> ForeignPtr a -> IO (ForeignPtr a)
      reallocHack dummy src = do
          let size = sizeOf dummy
          dst <- mallocPlainForeignPtrBytes (newSize * size)
          withForeignPtr src $ \s ->
            when (s /= nullPtr && oldSize > 0) .
              withForeignPtr dst $ \d -> do
                _ <- memcpy d s (fromIntegral (oldSize * size))
                return ()
          return dst

    new :: Storable a => Int -> IO (Array a)
    new c = do
        es <- allocArray cap
        fmap Array (newIORef (AC es 0 cap))
      where
        cap = firstPowerOf2 c

    duplicate :: Storable a => Array a -> IO (Array a)
    duplicate a = dupHack undefined a
     where
      dupHack :: Storable b => b -> Array b -> IO (Array b)
      dupHack dummy (Array ref) = do
        AC es len cap <- readIORef ref
        ary <- allocArray cap
        withForeignPtr ary $ \dest ->
          withForeignPtr es $ \src -> do
            _ <- memcpy dest src (fromIntegral (len * sizeOf dummy))
            return ()
        Array `fmap` newIORef (AC ary len cap)

    length :: Array a -> IO Int
    length (Array ref) = do
        AC _ len _ <- readIORef ref
        return len

    capacity :: Array a -> IO Int
    capacity (Array ref) = do
        AC _ _ cap <- readIORef ref
        return cap

    unsafeRead :: Storable a => Array a -> Int -> IO a
    unsafeRead (Array ref) ix = do
        AC es _ cap <- readIORef ref
        CHECK_BOUNDS("unsafeRead",cap,ix)
          withForeignPtr es $ \p ->
            peekElemOff p ix

    unsafeWrite :: Storable a => Array a -> Int -> a -> IO ()
    unsafeWrite (Array ref) ix a = do
        ac <- readIORef ref
        unsafeWrite' ac ix a

    unsafeWrite' :: Storable a => AC a -> Int -> a -> IO ()
    unsafeWrite' (AC es _ cap) ix a = do
        CHECK_BOUNDS("unsafeWrite'",cap,ix)
          withForeignPtr es $ \p ->
            pokeElemOff p ix a

    unsafeLoad :: Array a -> (Ptr a -> Int -> IO Int) -> IO Int
    unsafeLoad (Array ref) load = do
        AC es _ cap <- readIORef ref
        len' <- withForeignPtr es $ \p -> load p cap
        writeIORef ref (AC es len' cap)
        return len'

    ensureCapacity :: Storable a => Array a -> Int -> IO ()
    ensureCapacity (Array ref) c = do
        ac@(AC _ _ cap) <- readIORef ref
        ac'@(AC _ _ cap') <- ensureCapacity' ac c
        when (cap' /= cap) $
          writeIORef ref ac'

    ensureCapacity' :: Storable a => AC a -> Int -> IO (AC a)
    ensureCapacity' ac@(AC es len cap) c = do
        if c > cap
          then do
            es' <- reallocArray es cap' cap
            return (AC es' len cap')
          else
            return ac
      where
        cap' = firstPowerOf2 c

    useAsPtr :: Array a -> (Ptr a -> Int -> IO b) -> IO b
    useAsPtr (Array ref) f = do
        AC es len _ <- readIORef ref
        withForeignPtr es $ \p -> f p len

    snoc :: Storable a => Array a -> a -> IO ()
    snoc (Array ref) e = do
        ac@(AC _ len _) <- readIORef ref
        let len' = len + 1
        ac'@(AC es _ cap) <- ensureCapacity' ac len'
        unsafeWrite' ac' len e
        writeIORef ref (AC es len' cap)

    clear :: Array a -> IO ()
    clear (Array ref) = do
      atomicModifyIORef' ref $ \(AC es _ cap) ->
            (AC es 0 cap, ())

    forM_ :: Storable a => Array a -> (a -> IO ()) -> IO ()
    forM_ ary g = forHack ary g undefined
      where
        forHack :: Storable b => Array b -> (b -> IO ()) -> b -> IO ()
        forHack (Array ref) f dummy = do
          AC es len _ <- readIORef ref
          let size = sizeOf dummy
              offset = len * size
          withForeignPtr es $ \p -> do
            let go n | n >= offset = return ()
                     | otherwise = do
                  f =<< peek (p `plusPtr` n)
                  go (n + size)
            go 0

    loop :: Storable a => Array a -> b -> (b -> a -> IO (b,Bool)) -> IO ()
    loop ary z g = loopHack ary z g undefined
      where
        loopHack :: Storable b => Array b -> c -> (c -> b -> IO (c,Bool)) -> b
                 -> IO ()
        loopHack (Array ref) y f dummy = do
          AC es len _ <- readIORef ref
          let size = sizeOf dummy
              offset = len * size
          withForeignPtr es $ \p -> do
            let go n k
                    | n >= offset = return ()
                    | otherwise = do
                          (k',cont) <- f k =<< peek (p `plusPtr` n)
                          when cont $ go (n + size) k'
            go 0 y

    findIndex :: Storable a => (a -> Bool) -> Array a -> IO (Maybe (Int,a))
    findIndex = findHack undefined
     where
      findHack :: Storable b => b -> (b -> Bool) -> Array b -> IO (Maybe (Int,b))
      findHack dummy p (Array ref) = do
        AC es len _ <- readIORef ref
        let size   = sizeOf dummy
            offset = len * size
        withForeignPtr es $ \ptr ->
          let go !n !i
                | n >= offset = return Nothing
                | otherwise = do
                    val <- peek (ptr `plusPtr` n)
                    if p val
                      then return $ Just (i, val)
                      else go (n + size) (i + 1)
          in  go 0 0

    concat :: Storable a => Array a -> Array a -> IO ()
    concat (Array d) (Array s) = do
      da@(AC _ dlen _) <- readIORef d
      sa@(AC _ slen _) <- readIORef s
      writeIORef d =<< copy' da dlen sa 0 slen

    -- | Copy part of the source array into the destination array. The
    -- destination array is resized if not large enough.
    copy :: Storable a => Array a -> Int -> Array a -> Int -> Int -> IO ()
    copy (Array d) dstart (Array s) sstart maxCount = do
      da <- readIORef d
      sa <- readIORef s
      writeIORef d =<< copy' da dstart sa sstart maxCount

    -- | Copy part of the source array into the destination array. The
    -- destination array is resized if not large enough.
    copy' :: Storable a => AC a -> Int -> AC a -> Int -> Int -> IO (AC a)
    copy' d dstart s sstart maxCount = copyHack d s undefined
     where
      copyHack :: Storable b => AC b -> AC b -> b -> IO (AC b)
      copyHack dac@(AC _ oldLen _) (AC src slen _) dummy = do
        when (maxCount < 0 || dstart < 0 || dstart > oldLen || sstart < 0 ||
              sstart > slen) $ errorWithoutStackTrace "copy: bad offsets or lengths"
        let size = sizeOf dummy
            count = min maxCount (slen - sstart)
        if count == 0
          then return dac
          else do
            AC dst dlen dcap <- ensureCapacity' dac (dstart + count)
            withForeignPtr dst $ \dptr ->
              withForeignPtr src $ \sptr -> do
                _ <- memcpy (dptr `plusPtr` (dstart * size))
                            (sptr `plusPtr` (sstart * size))
                            (fromIntegral (count * size))
                return $ AC dst (max dlen (dstart + count)) dcap

    removeAt :: Storable a => Array a -> Int -> IO ()
    removeAt a i = removeHack a undefined
     where
      removeHack :: Storable b => Array b -> b -> IO ()
      removeHack (Array ary) dummy = do
        AC fp oldLen cap <- readIORef ary
        when (i < 0 || i >= oldLen) $ errorWithoutStackTrace "removeAt: invalid index"
        let size   = sizeOf dummy
            newLen = oldLen - 1
        when (newLen > 0 && i < newLen) .
          withForeignPtr fp $ \ptr -> do
            _ <- memmove (ptr `plusPtr` (size * i))
                         (ptr `plusPtr` (size * (i+1)))
                         (fromIntegral (size * (newLen-i)))
            return ()
        writeIORef ary (AC fp newLen cap)

    {-The firstPowerOf2 function works by setting all bits on the right-hand
    side of the most significant flagged bit to 1, and then incrementing
    the entire value at the end so it "rolls over" to the nearest power of
    two.
    -}

    -- | Computes the next-highest power of two for a particular integer,
    -- @n@.  If @n@ is already a power of two, returns @n@.  If @n@ is
    -- zero, returns zero, even though zero is not a power of two.
    firstPowerOf2 :: Int -> Int
    firstPowerOf2 !n =
        let !n1 = n - 1
            !n2 = n1 .|. (n1 `shiftR` 1)
            !n3 = n2 .|. (n2 `shiftR` 2)
            !n4 = n3 .|. (n3 `shiftR` 4)
            !n5 = n4 .|. (n4 `shiftR` 8)
            !n6 = n5 .|. (n5 `shiftR` 16)
    #if WORD_SIZE_IN_BITS == 32
        in n6 + 1
    #elif WORD_SIZE_IN_BITS == 64
            !n7 = n6 .|. (n6 `shiftR` 32)
        in n7 + 1
    #else
    # error firstPowerOf2 not defined on this architecture
    #endif

    foreign import ccall unsafe "string.h memcpy"
        memcpy :: Ptr a -> Ptr a -> CSize -> IO (Ptr a)

    foreign import ccall unsafe "string.h memmove"
        memmove :: Ptr a -> Ptr a -> CSize -> IO (Ptr a)