-- Copyright 2019 Google LLC
--
-- Use of this source code is governed by a BSD-style
-- license that can be found in the LICENSE file or at
-- https://developers.google.com/open-source/licenses/bsd

module Serialize (HasPtrs (..), takePtrSnapshot, restorePtrSnapshot) where

import Prelude hiding (pi, abs)
import Control.Monad
import qualified Data.ByteString as BS
import Data.ByteString.Internal (memcpy)
import Data.ByteString.Unsafe (unsafeUseAsCString)
import Data.Int
import Data.Store hiding (size)
import Foreign.Ptr
import Foreign.Marshal.Array
import GHC.Generics (Generic)

import Types.Primitives

foreign import ccall "malloc_dex"           dexMalloc    :: Int64  -> IO (Ptr ())
foreign import ccall "dex_allocation_size"  dexAllocSize :: Ptr () -> IO Int64

data WithSnapshot a = WithSnapshot a [PtrSnapshot]  deriving Generic
type RawPtr = Ptr ()

class HasPtrs a where
  traversePtrs :: Applicative f => (PtrType -> RawPtr -> f RawPtr) -> a -> f a

takePtrSnapshot :: PtrType -> PtrLitVal -> IO PtrLitVal
takePtrSnapshot _ NullPtr = return NullPtr
takePtrSnapshot (CPU, ptrTy) (PtrLitVal ptrVal) = case ptrTy of
  PtrType eltTy -> do
    childPtrs <- loadPtrPtrs ptrVal
    PtrSnapshot <$> PtrArray <$> mapM (takePtrSnapshot eltTy) childPtrs
  _ -> PtrSnapshot . ByteArray <$> loadPtrBytes ptrVal
takePtrSnapshot (GPU, _) _ = error "Snapshots of GPU memory not implemented"
takePtrSnapshot _ (PtrSnapshot _) = error "Already a snapshot"
{-# SCC takePtrSnapshot #-}

loadPtrBytes :: RawPtr -> IO BS.ByteString
loadPtrBytes ptr = do
  numBytes <- fromIntegral <$> dexAllocSize ptr
  liftM BS.pack $ peekArray numBytes $ castPtr ptr

loadPtrPtrs :: RawPtr -> IO [PtrLitVal]
loadPtrPtrs ptr = do
  numBytes <- fromIntegral <$> dexAllocSize ptr
  childPtrs <- peekArray (numBytes `div` ptrSize) $ castPtr ptr
  forM childPtrs \childPtr ->
    if childPtr == nullPtr
      then return NullPtr
      else return $ PtrLitVal childPtr

restorePtrSnapshot :: PtrLitVal -> IO PtrLitVal
restorePtrSnapshot NullPtr = return NullPtr
restorePtrSnapshot (PtrSnapshot snapshot) = case snapshot of
  PtrArray  children -> do
    childrenPtrs <- forM children \child ->
      restorePtrSnapshot child >>= \case
        NullPtr -> return nullPtr
        PtrLitVal p -> return p
        PtrSnapshot _ -> error "expected a pointer literal"
    PtrLitVal <$> storePtrPtrs childrenPtrs
  ByteArray bytes -> PtrLitVal <$> storePtrBytes bytes
restorePtrSnapshot (PtrLitVal _) = error "not a snapshot"
{-# SCC restorePtrSnapshot #-}

storePtrBytes :: BS.ByteString -> IO RawPtr
storePtrBytes xs = do
  let numBytes = BS.length xs
  destPtr <- dexMalloc $ fromIntegral numBytes
  -- this is safe because we don't modify srcPtr's memory or let it escape
  unsafeUseAsCString xs \srcPtr ->
    memcpy (castPtr destPtr) (castPtr srcPtr) numBytes
  return destPtr

storePtrPtrs :: [RawPtr] -> IO RawPtr
storePtrPtrs ptrs = do
  ptr <- dexMalloc $ fromIntegral $ length ptrs * ptrSize
  pokeArray (castPtr ptr) ptrs
  return ptr

-- === instances ===

instance Store a => Store (WithSnapshot a)