{-# 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)