GHCi/BinaryArray.hs

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

    import Prelude
    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)