GHCi/BinaryArray.hs

    {-# LANGUAGE BangPatterns, MagicHash, UnboxedTuples, FlexibleContexts #-}
    -- | Efficient serialisation for GHCi Instruction arrays
    --
    -- Author: Ben Gamari
    --
    module GHCi.BinaryArray(putArray, getArray) where

    import Foreign.Ptr
    import Data.Binary
    import Data.Binary.Put (putBuilder)
    import qualified Data.Binary.Get.Internal as Binary
    import qualified Data.ByteString.Builder as BB
    import qualified Data.ByteString.Builder.Internal as BB
    import qualified Data.Array.Base as A
    import qualified Data.Array.IO.Internals as A
    import qualified Data.Array.Unboxed as A
    import GHC.Exts
    import GHC.IO

    -- | An efficient serialiser of 'A.UArray'.
    putArray :: Binary i => A.UArray i a -> Put
    putArray (A.UArray l u _ arr#) = do
        put l
        put u
        putBuilder $ byteArrayBuilder arr#

    byteArrayBuilder :: ByteArray# -> BB.Builder
    byteArrayBuilder arr# = BB.builder $ go 0 (I# (sizeofByteArray# arr#))
      where
        go :: Int -> Int -> BB.BuildStep a -> BB.BuildStep a
        go !inStart !inEnd k (BB.BufferRange outStart outEnd)
          -- There is enough room in this output buffer to write all remaining array
          -- contents
          | inRemaining <= outRemaining = do
              copyByteArrayToAddr arr# inStart outStart inRemaining
              k (BB.BufferRange (outStart `plusPtr` inRemaining) outEnd)
          -- There is only enough space for a fraction of the remaining contents
          | otherwise = do
              copyByteArrayToAddr arr# inStart outStart outRemaining
              let !inStart' = inStart + outRemaining
              return $! BB.bufferFull 1 outEnd (go inStart' inEnd k)
          where
            inRemaining  = inEnd - inStart
            outRemaining = outEnd `minusPtr` outStart

        copyByteArrayToAddr :: ByteArray# -> Int -> Ptr a -> Int -> IO ()
        copyByteArrayToAddr src# (I# src_off#) (Ptr dst#) (I# len#) =
            IO $ \s -> case copyByteArrayToAddr# src# src_off# dst# len# s of
                         s' -> (# s', () #)

    -- | An efficient deserialiser of 'A.UArray'.
    getArray :: (Binary i, A.Ix i, A.MArray A.IOUArray a IO) => Get (A.UArray i a)
    getArray = do
        l <- get
        u <- get
        arr@(A.IOUArray (A.STUArray _ _ _ arr#)) <-
            return $ unsafeDupablePerformIO $ A.newArray_ (l,u)
        let go 0 _ = return ()
            go !remaining !off = do
                Binary.readNWith n $ \ptr ->
                  copyAddrToByteArray ptr arr# off n
                go (remaining - n) (off + n)
              where n = min chunkSize remaining
        go (I# (sizeofMutableByteArray# arr#)) 0
        return $! unsafeDupablePerformIO $ unsafeFreezeIOUArray arr
      where
        chunkSize = 10*1024

        copyAddrToByteArray :: Ptr a -> MutableByteArray# RealWorld
                            -> Int -> Int -> IO ()
        copyAddrToByteArray (Ptr src#) dst# (I# dst_off#) (I# len#) =
            IO $ \s -> case copyAddrToByteArray# src# dst# dst_off# len# s of
                         s' -> (# s', () #)

    -- this is inexplicably not exported in currently released array versions
    unsafeFreezeIOUArray :: A.IOUArray ix e -> IO (A.UArray ix e)
    unsafeFreezeIOUArray (A.IOUArray marr) = stToIO (A.unsafeFreezeSTUArray marr)