{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE RankNTypes #-}
{-# LANGUAGE MagicHash #-}
{-# LANGUAGE GADTs #-}
{-# LANGUAGE TypeApplications #-}
{-# LANGUAGE ScopedTypeVariables #-}

module Data.Compact.Serialize
    ( writeCompact, unsafeReadCompact
    , hPutCompact, hUnsafeGetCompact
    ) where

import Type.Reflection
import Control.Monad
import Data.Monoid
import Data.Word
import System.IO

import Foreign.Ptr
import Foreign.Storable
import Foreign.Marshal.Alloc

import qualified Data.Binary as B
import qualified Data.ByteString.Lazy as BSL

import GHC.Compact
import GHC.Compact.Serialized

newtype CompactFile a = CompactFile (SerializedCompact a)

-- NB: This instance cannot be put on SerializedCompact as
-- ghc-compact does not have a binary dependency
instance (Typeable a) => B.Binary (CompactFile a) where
    get :: Get (CompactFile a)
get = do
        Word64
magic <- Get Word64
forall t. Binary t => Get t
B.get
        Bool -> Get () -> Get ()
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
when (Word64
magic Word64 -> Word64 -> Bool
forall a. Eq a => a -> a -> Bool
/= Word64
magicNumber) (Get () -> Get ()) -> Get () -> Get ()
forall a b. (a -> b) -> a -> b
$
            String -> Get ()
forall a. String -> Get a
forall (m :: * -> *) a. MonadFail m => String -> m a
fail String
"Data.Compact.Serialized: bad magic number"
        SomeTypeRep TypeRep a
tyrep <- Get SomeTypeRep
forall t. Binary t => Get t
B.get
        case TypeRep a
tyrep TypeRep a -> TypeRep a -> Maybe (a :~~: a)
forall k1 k2 (a :: k1) (b :: k2).
TypeRep a -> TypeRep b -> Maybe (a :~~: b)
`eqTypeRep` forall a. Typeable a => TypeRep a
forall {k} (a :: k). Typeable a => TypeRep a
typeRep @a of
          Just a :~~: a
HRefl -> SerializedCompact a -> CompactFile a
forall a. SerializedCompact a -> CompactFile a
CompactFile (SerializedCompact a -> CompactFile a)
-> Get (SerializedCompact a) -> Get (CompactFile a)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Get (SerializedCompact a)
forall a. Get (SerializedCompact a)
getSerializedCompact
          Maybe (a :~~: a)
Nothing -> String -> Get (CompactFile a)
forall a. String -> Get a
forall (m :: * -> *) a. MonadFail m => String -> m a
fail (String -> Get (CompactFile a)) -> String -> Get (CompactFile a)
forall a b. (a -> b) -> a -> b
$
            String
"Data.Compact.Serialized: expected " String -> String -> String
forall a. [a] -> [a] -> [a]
++ TypeRep a -> String
forall a. Show a => a -> String
show (forall a. Typeable a => TypeRep a
forall {k} (a :: k). Typeable a => TypeRep a
typeRep @a) String -> String -> String
forall a. [a] -> [a] -> [a]
++
            String
" but got " String -> String -> String
forall a. [a] -> [a] -> [a]
++ TypeRep a -> String
forall a. Show a => a -> String
show TypeRep a
tyrep
    put :: CompactFile a -> Put
put (CompactFile SerializedCompact a
a) = Word64 -> Put
forall t. Binary t => t -> Put
B.put Word64
magicNumber Put -> Put -> Put
forall a b. PutM a -> PutM b -> PutM b
forall (m :: * -> *) a b. Monad m => m a -> m b -> m b
>> TypeRep a -> Put
forall t. Binary t => t -> Put
B.put (forall a. Typeable a => TypeRep a
forall {k} (a :: k). Typeable a => TypeRep a
typeRep @a) Put -> Put -> Put
forall a b. PutM a -> PutM b -> PutM b
forall (m :: * -> *) a b. Monad m => m a -> m b -> m b
>> SerializedCompact a -> Put
forall a. SerializedCompact a -> Put
putSerializedCompact SerializedCompact a
a

-- Serves as a very rudimentary integrity check.
-- Chosen at random by roll of a 2^64-sided die; it was a sight to behold.
magicNumber :: Word64
magicNumber :: Word64
magicNumber = Word64
0x7c155e7a53f094f2

putPtr :: Ptr a -> B.Put
putPtr :: forall a. Ptr a -> Put
putPtr = forall t. Binary t => t -> Put
B.put @Word64 (Word64 -> Put) -> (Ptr a -> Word64) -> Ptr a -> Put
forall b c a. (b -> c) -> (a -> b) -> a -> c
. WordPtr -> Word64
forall a b. (Integral a, Num b) => a -> b
fromIntegral (WordPtr -> Word64) -> (Ptr a -> WordPtr) -> Ptr a -> Word64
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Ptr a -> WordPtr
forall a. Ptr a -> WordPtr
ptrToWordPtr

getPtr :: B.Get (Ptr a)
getPtr :: forall a. Get (Ptr a)
getPtr = WordPtr -> Ptr a
forall a. WordPtr -> Ptr a
wordPtrToPtr (WordPtr -> Ptr a) -> (Word64 -> WordPtr) -> Word64 -> Ptr a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Word64 -> WordPtr
forall a b. (Integral a, Num b) => a -> b
fromIntegral (Word64 -> Ptr a) -> Get Word64 -> Get (Ptr a)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> forall t. Binary t => Get t
B.get @Word64

getList :: B.Get a -> B.Get [a]
getList :: forall a. Get a -> Get [a]
getList Get a
getElem = do
    Int
n <- forall t. Binary t => Get t
B.get @Int
    Int -> Get a -> Get [a]
forall (m :: * -> *) a. Applicative m => Int -> m a -> m [a]
replicateM Int
n Get a
getElem

putList :: (a -> B.Put) -> [a] -> B.Put
putList :: forall a. (a -> Put) -> [a] -> Put
putList a -> Put
putElem [a]
xs = do
    forall t. Binary t => t -> Put
B.put @Int ([a] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [a]
xs)
    (a -> Put) -> [a] -> Put
forall (t :: * -> *) (m :: * -> *) a b.
(Foldable t, Monad m) =>
(a -> m b) -> t a -> m ()
mapM_ a -> Put
putElem [a]
xs

getSerializedCompact :: B.Get (SerializedCompact a)
getSerializedCompact :: forall a. Get (SerializedCompact a)
getSerializedCompact = [(Ptr (), Word)] -> Ptr () -> SerializedCompact a
forall a. [(Ptr (), Word)] -> Ptr () -> SerializedCompact a
SerializedCompact ([(Ptr (), Word)] -> Ptr () -> SerializedCompact a)
-> Get [(Ptr (), Word)] -> Get (Ptr () -> SerializedCompact a)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Get (Ptr (), Word) -> Get [(Ptr (), Word)]
forall a. Get a -> Get [a]
getList Get (Ptr (), Word)
forall a. Get (Ptr a, Word)
getBlock Get (Ptr () -> SerializedCompact a)
-> Get (Ptr ()) -> Get (SerializedCompact a)
forall a b. Get (a -> b) -> Get a -> Get b
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> Get (Ptr ())
forall a. Get (Ptr a)
getPtr
  where
    getBlock :: B.Get (Ptr a, Word)
    getBlock :: forall a. Get (Ptr a, Word)
getBlock = (,) (Ptr a -> Word -> (Ptr a, Word))
-> Get (Ptr a) -> Get (Word -> (Ptr a, Word))
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Get (Ptr a)
forall a. Get (Ptr a)
getPtr Get (Word -> (Ptr a, Word)) -> Get Word -> Get (Ptr a, Word)
forall a b. Get (a -> b) -> Get a -> Get b
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> Get Word
forall t. Binary t => Get t
B.get

putSerializedCompact :: SerializedCompact a -> B.Put
putSerializedCompact :: forall a. SerializedCompact a -> Put
putSerializedCompact (SerializedCompact [(Ptr (), Word)]
a Ptr ()
b) = ((Ptr (), Word) -> Put) -> [(Ptr (), Word)] -> Put
forall a. (a -> Put) -> [a] -> Put
putList (Ptr (), Word) -> Put
forall a. (Ptr a, Word) -> Put
putBlock [(Ptr (), Word)]
a Put -> Put -> Put
forall a. Semigroup a => a -> a -> a
<> Ptr () -> Put
forall a. Ptr a -> Put
putPtr Ptr ()
b
  where
    putBlock :: (Ptr a, Word) -> B.Put
    putBlock :: forall a. (Ptr a, Word) -> Put
putBlock (Ptr a
ptr, Word
len) = Ptr a -> Put
forall a. Ptr a -> Put
putPtr Ptr a
ptr Put -> Put -> Put
forall a. Semigroup a => a -> a -> a
<> Word -> Put
forall t. Binary t => t -> Put
B.put Word
len

-- | Write a compact region to a file.  The resulting file can
-- be read back into memory using 'unsafeReadCompact'.
--
writeCompact :: Typeable a => FilePath -> Compact a -> IO ()
writeCompact :: forall a. Typeable a => String -> Compact a -> IO ()
writeCompact String
fname Compact a
compact =
    String -> IOMode -> (Handle -> IO ()) -> IO ()
forall r. String -> IOMode -> (Handle -> IO r) -> IO r
withFile String
fname IOMode
WriteMode ((Handle -> IO ()) -> IO ()) -> (Handle -> IO ()) -> IO ()
forall a b. (a -> b) -> a -> b
$ \Handle
h -> Handle -> Compact a -> IO ()
forall a. Typeable a => Handle -> Compact a -> IO ()
hPutCompact Handle
h Compact a
compact

-- | Write a compact region to a 'Handle'.  The compact region
-- can be read out of the handle by using 'hUnsafeGetCompact'.
--
hPutCompact :: Typeable a => Handle -> Compact a -> IO ()
hPutCompact :: forall a. Typeable a => Handle -> Compact a -> IO ()
hPutCompact Handle
hdl Compact a
compact =
    Compact a -> (SerializedCompact a -> IO ()) -> IO ()
forall a c. Compact a -> (SerializedCompact a -> IO c) -> IO c
withSerializedCompact Compact a
compact ((SerializedCompact a -> IO ()) -> IO ())
-> (SerializedCompact a -> IO ()) -> IO ()
forall a b. (a -> b) -> a -> b
$ \SerializedCompact a
scompact -> do
        let bs :: ByteString
bs = CompactFile a -> ByteString
forall a. Binary a => a -> ByteString
B.encode (CompactFile a -> ByteString) -> CompactFile a -> ByteString
forall a b. (a -> b) -> a -> b
$ SerializedCompact a -> CompactFile a
forall a. SerializedCompact a -> CompactFile a
CompactFile SerializedCompact a
scompact
        -- By writing out the length of the metadata segment, we
        -- can read out a single word, read out the metadata segment,
        -- and then immediately start blasting further data from
        -- the handle into the memory region where the compact
        -- is going to go.  Otherwise, we have to indirect through
        -- a lazy bytestring which has a cost.
        Handle -> Int -> IO ()
forall a. Storable a => Handle -> a -> IO ()
hPutStorable Handle
hdl (Int64 -> Int
forall a b. (Integral a, Num b) => a -> b
fromIntegral (ByteString -> Int64
BSL.length ByteString
bs) :: Int)
        Handle -> ByteString -> IO ()
BSL.hPut Handle
hdl ByteString
bs
        let putBlock :: (Ptr (), Word) -> IO ()
putBlock (Ptr ()
ptr, Word
len) = Handle -> Ptr () -> Int -> IO ()
forall a. Handle -> Ptr a -> Int -> IO ()
hPutBuf Handle
hdl Ptr ()
ptr (Word -> Int
forall a b. (Integral a, Num b) => a -> b
fromIntegral Word
len)
        ((Ptr (), Word) -> IO ()) -> [(Ptr (), Word)] -> IO ()
forall (t :: * -> *) (m :: * -> *) a b.
(Foldable t, Monad m) =>
(a -> m b) -> t a -> m ()
mapM_ (Ptr (), Word) -> IO ()
putBlock (SerializedCompact a -> [(Ptr (), Word)]
forall a. SerializedCompact a -> [(Ptr (), Word)]
serializedCompactBlockList SerializedCompact a
scompact)

-- | Read out a compact region that was serialized to a file.
-- See 'hUnsafeGetCompact' for safety considerations when using this function.
--
unsafeReadCompact :: Typeable a => FilePath -> IO (Either String (Compact a))
unsafeReadCompact :: forall a. Typeable a => String -> IO (Either String (Compact a))
unsafeReadCompact String
fname =
    String
-> IOMode
-> (Handle -> IO (Either String (Compact a)))
-> IO (Either String (Compact a))
forall r. String -> IOMode -> (Handle -> IO r) -> IO r
withFile String
fname IOMode
ReadMode ((Handle -> IO (Either String (Compact a)))
 -> IO (Either String (Compact a)))
-> (Handle -> IO (Either String (Compact a)))
-> IO (Either String (Compact a))
forall a b. (a -> b) -> a -> b
$ \Handle
hdl -> Handle -> IO (Either String (Compact a))
forall a. Typeable a => Handle -> IO (Either String (Compact a))
hUnsafeGetCompact Handle
hdl

-- | Read out a compact region from a handle.
--
-- Compact regions written to handles this way are subject to some
-- restrictions:
--
--  * Our binary representation contains direct pointers to the info
--    tables of objects in the region.  This means that the info tables
--    of the receiving process must be laid out in exactly the same
--    way as from the original process; in practice, this means using
--    static linking, using the exact same binary and turning off ASLR.  This
--    API does NOT do any safety checking and will probably segfault if you
--    get it wrong.  DO NOT run this on untrusted input.
--
--  * You must read out the value at the correct type.  We will
--    check this for you and raise an error if the types do not match
--    (this is what the 'Typeable' constraint is for).
--
hUnsafeGetCompact
    :: forall a. Typeable a
    => Handle -> IO (Either String (Compact a))
hUnsafeGetCompact :: forall a. Typeable a => Handle -> IO (Either String (Compact a))
hUnsafeGetCompact Handle
hdl = do
    Int
l <- Handle -> IO Int
forall a. Storable a => Handle -> IO a
hGetStorable Handle
hdl
    ByteString
mbs <- Handle -> Int -> IO ByteString
BSL.hGet Handle
hdl (Int
l :: Int)
    case forall a.
Binary a =>
ByteString
-> Either (ByteString, Int64, String) (ByteString, Int64, a)
B.decodeOrFail @(CompactFile a) ByteString
mbs of
      Left (ByteString
_, Int64
_, String
err) -> Either String (Compact a) -> IO (Either String (Compact a))
forall a. a -> IO a
forall (m :: * -> *) a. Monad m => a -> m a
return (Either String (Compact a) -> IO (Either String (Compact a)))
-> Either String (Compact a) -> IO (Either String (Compact a))
forall a b. (a -> b) -> a -> b
$ String -> Either String (Compact a)
forall a b. a -> Either a b
Left String
err
      Right (ByteString
rest, Int64
_, CompactFile SerializedCompact a
scompact)
        | Bool -> Bool
not (ByteString -> Bool
BSL.null ByteString
rest) -> Either String (Compact a) -> IO (Either String (Compact a))
forall a. a -> IO a
forall (m :: * -> *) a. Monad m => a -> m a
return (Either String (Compact a) -> IO (Either String (Compact a)))
-> (String -> Either String (Compact a))
-> String
-> IO (Either String (Compact a))
forall b c a. (b -> c) -> (a -> b) -> a -> c
. String -> Either String (Compact a)
forall a b. a -> Either a b
Left
            (String -> IO (Either String (Compact a)))
-> String -> IO (Either String (Compact a))
forall a b. (a -> b) -> a -> b
$ String
"Had " String -> String -> String
forall a. [a] -> [a] -> [a]
++ Int64 -> String
forall a. Show a => a -> String
show (ByteString -> Int64
BSL.length ByteString
rest) String -> String -> String
forall a. [a] -> [a] -> [a]
++ String
" bytes of leftover metadata"
        | Bool
otherwise  -> do
          Maybe (Compact a)
res <- SerializedCompact a
-> (Ptr Any -> Word -> IO ()) -> IO (Maybe (Compact a))
forall a b.
SerializedCompact a
-> (Ptr b -> Word -> IO ()) -> IO (Maybe (Compact a))
importCompact SerializedCompact a
scompact ((Ptr Any -> Word -> IO ()) -> IO (Maybe (Compact a)))
-> (Ptr Any -> Word -> IO ()) -> IO (Maybe (Compact a))
forall a b. (a -> b) -> a -> b
$ \Ptr Any
ptr Word
l ->
                    IO Int -> IO ()
forall (f :: * -> *) a. Functor f => f a -> f ()
void (IO Int -> IO ()) -> IO Int -> IO ()
forall a b. (a -> b) -> a -> b
$ Handle -> Ptr Any -> Int -> IO Int
forall a. Handle -> Ptr a -> Int -> IO Int
hGetBuf Handle
hdl Ptr Any
ptr (Word -> Int
forall a b. (Integral a, Num b) => a -> b
fromIntegral Word
l)
          case Maybe (Compact a)
res of
            Maybe (Compact a)
Nothing -> Either String (Compact a) -> IO (Either String (Compact a))
forall a. a -> IO a
forall (m :: * -> *) a. Monad m => a -> m a
return (Either String (Compact a) -> IO (Either String (Compact a)))
-> Either String (Compact a) -> IO (Either String (Compact a))
forall a b. (a -> b) -> a -> b
$ String -> Either String (Compact a)
forall a b. a -> Either a b
Left String
"Failed to import compact"
            Just Compact a
compact -> Either String (Compact a) -> IO (Either String (Compact a))
forall a. a -> IO a
forall (m :: * -> *) a. Monad m => a -> m a
return (Either String (Compact a) -> IO (Either String (Compact a)))
-> Either String (Compact a) -> IO (Either String (Compact a))
forall a b. (a -> b) -> a -> b
$ Compact a -> Either String (Compact a)
forall a b. b -> Either a b
Right Compact a
compact

hPutStorable :: forall a. Storable a => Handle -> a -> IO ()
hPutStorable :: forall a. Storable a => Handle -> a -> IO ()
hPutStorable Handle
h a
a =
    (Ptr a -> IO ()) -> IO ()
forall a b. Storable a => (Ptr a -> IO b) -> IO b
alloca ((Ptr a -> IO ()) -> IO ()) -> (Ptr a -> IO ()) -> IO ()
forall a b. (a -> b) -> a -> b
$ \Ptr a
ptr -> do
        Ptr a -> a -> IO ()
forall a. Storable a => Ptr a -> a -> IO ()
poke Ptr a
ptr a
a
        Handle -> Ptr a -> Int -> IO ()
forall a. Handle -> Ptr a -> Int -> IO ()
hPutBuf Handle
h Ptr a
ptr (a -> Int
forall a. Storable a => a -> Int
sizeOf (a
forall a. HasCallStack => a
undefined :: a))

hGetStorable :: forall a. Storable a => Handle -> IO a
hGetStorable :: forall a. Storable a => Handle -> IO a
hGetStorable Handle
h =
    (Ptr a -> IO a) -> IO a
forall a b. Storable a => (Ptr a -> IO b) -> IO b
alloca ((Ptr a -> IO a) -> IO a) -> (Ptr a -> IO a) -> IO a
forall a b. (a -> b) -> a -> b
$ \Ptr a
ptr -> do
        Handle -> Ptr a -> Int -> IO Int
forall a. Handle -> Ptr a -> Int -> IO Int
hGetBuf Handle
h Ptr a
ptr (a -> Int
forall a. Storable a => a -> Int
sizeOf (a
forall a. HasCallStack => a
undefined :: a))
        Ptr a -> IO a
forall a. Storable a => Ptr a -> IO a
peek Ptr a
ptr