{-
  Copyright (c) Meta Platforms, Inc. and affiliates.
  All rights reserved.

  This source code is licensed under the BSD-style license found in the
  LICENSE file in the root directory of this source tree.
-}

{-# LANGUAGE DerivingStrategies #-}
module Glean.RTS (
  Fid(..), lowestFid,
  Pid(..), lowestPid,

  withValue,
  toValue, fromValue,
  encodeByteArray,
  mangleString,

  ByteStringRef(..), unsafeWithByteStringRef, derefByteString,
  MangledStringRef(..), unsafeWithMangledStringRef, demangle,
  Decoder, DecodingException(..), withDecoder,
  dByte, dNat,
  dTrustedStringRef, dString,
  dArray, dSet, dByteStringRef, dBytes,
  dSelector, dFact,

  glean_push_value_byte,
  glean_push_value_nat,
  glean_push_value_string,
  glean_push_value_array,
  glean_push_value_set,
  glean_push_value_bytes,
  glean_push_value_selector,
  glean_push_value_fact,

  glean_push_fact,

  glean_pop_value_byte,
  glean_pop_value_nat,
  glean_pop_value_string,
  glean_pop_value_array,
  glean_pop_value_set,
  glean_pop_value_bytes,
  glean_pop_value_selector,
  glean_pop_value_fact,
) where

import Data.Word
import Foreign.C.String
import Foreign.C.Types
import Foreign.Marshal.Alloc
import Foreign.Ptr
import Foreign.Storable

import Control.Exception
import Control.Monad
import Control.Monad.ST
import Control.Monad.ST.Unsafe
import Data.ByteString (ByteString)
import qualified Data.ByteString as BS
import qualified Data.ByteString.Unsafe as BS
import System.IO.Unsafe (unsafePerformIO)

import qualified Util.FFI as FFI
import Util.List (uniq)

import qualified Glean.FFI as FFI
import Glean.RTS.Builder
import Glean.RTS.Types
import Glean.RTS.Term

withValue :: Value -> (Builder -> IO a) -> IO a
withValue v f = withBuilder $ \b -> do
  encodeValue b v
  f b

encodeValue :: Builder -> Value -> IO ()
encodeValue b (Byte n) = FFI.call $ glean_push_value_byte b n
encodeValue b (Nat n) = FFI.call $ glean_push_value_nat b n
encodeValue b (Array xs) = do
  FFI.call $ glean_push_value_array b $ fromIntegral $ length xs
  mapM_ (encodeValue b) xs
encodeValue b (ByteArray xs) = encodeByteArray b xs
encodeValue b (Tuple xs) = mapM_ (encodeValue b) xs
encodeValue b (Alt n x) = do
  FFI.call $ glean_push_value_selector b $ fromIntegral n
  encodeValue b x
encodeValue b (String s) = FFI.unsafeWithBytes s $ \p n ->
  FFI.call $ glean_push_value_string b (castPtr p) n
encodeValue b (Ref id) = FFI.call $ glean_push_value_fact b id

encodeByteArray :: Builder -> ByteString -> IO ()
encodeByteArray b xs =
  BS.unsafeUseAsCStringLen xs $ \(p,n) -> do
    FFI.call $ glean_push_value_array b (fromIntegral n)
    FFI.call $ glean_push_value_bytes b (castPtr p) (fromIntegral n)

toValue :: Rep Pid -> ByteString -> Value
toValue ty bs = runST $ withDecoder bs $ \d -> decodeValue d ty

fromValue :: Value -> ByteString
fromValue val = unsafePerformIO $ withValue val finishBuilder

-- | A reference to a range of bytes with a context-dependent lifetime.
data ByteStringRef = ByteStringRef
  { byteStringRefPtr :: {-# UNPACK #-} !(Ptr Word8)
  , byteStringRefSize :: {-# UNPACK #-} !CSize
  }

-- | Execute the action with a reference to the contents of the 'ByteString'.
unsafeWithByteStringRef :: ByteString -> (ByteStringRef -> IO a) -> IO a
unsafeWithByteStringRef s f =
  FFI.unsafeWithBytes s $ \p n -> f $ ByteStringRef (castPtr p) n

-- | Construct a 'ByteString' by copying a range of bytes.
derefByteString :: ByteStringRef -> IO ByteString
derefByteString (ByteStringRef p n) = FFI.copyByteString (castPtr p) n

-- | A reference to a range of bytes containing an RTS-mangled string.
newtype MangledStringRef = MangledStringRef
  { mangledBytes  :: ByteStringRef
      -- ^ Byte range including terminator.
  }

-- | Construct a 'MangledStringRef' from a 'ByteString' containing the mangled
-- bytes. The references can't be used outside the function.
unsafeWithMangledStringRef :: ByteString -> (MangledStringRef -> IO a) -> IO a
unsafeWithMangledStringRef s f =
  unsafeWithByteStringRef s $ f . MangledStringRef

-- | Yield the internal representation of a UTF-8 string.
mangleString :: ByteString -> ByteString
mangleString s = BS.intercalate "\0\1" (BS.split 0 s) <> "\0\0"

-- | A stateful 'Decoder' for binary values.
data Decoder s = Decoder
  { decoderPtr :: {-# UNPACK #-} !(Ptr (Ptr ()))
  , decoderEnd :: {-# UNPACK #-} !(Ptr ())
  , decoderBuf :: {-# UNPACK #-} !FFI.FFIResultBuf
  }

newtype DecodingException = DecodingException String
  deriving newtype Show

instance Exception DecodingException

-- | Create a temporary 'Decoder' for a 'ByteString'.
withDecoder :: ByteString -> (Decoder s -> ST s a) -> ST s a
withDecoder bs f = unsafeIOToST $
  FFI.withFFIResultBuf $ \ !buf ->
  FFI.unsafeWithBytes bs $ \start len ->
  alloca $ \pstart -> do
    poke pstart start
    let end = start `plusPtr` fromIntegral len
    res <- unsafeSTToIO $ f Decoder
      { decoderPtr = pstart
      , decoderEnd = end
      , decoderBuf = buf
      }
    pos <- peek pstart
    when (pos /= end) $
      let extra = show (end `minusPtr` pos) in
      throwIO $ DecodingException $ "extra "<> extra <>" bytes at end of value"
    return res

dffi
  :: Storable a
  => (Ptr (Ptr ()) -> Ptr () -> Ptr a -> IO CString) -> Decoder s -> ST s a
{-# INLINE dffi #-}
dffi f d =
  unsafeIOToST $ FFI.ffiBuf (decoderBuf d) $ f (decoderPtr d) (decoderEnd d)

dffi2
  :: (Storable a, Storable b)
  => (Ptr (Ptr ()) -> Ptr () -> Ptr a -> Ptr b -> IO CString)
  -> Decoder s
  -> ST s (a,b)
{-# INLINE dffi2 #-}
dffi2 f d =
  unsafeIOToST $ FFI.ffiBuf2 (decoderBuf d) $ f (decoderPtr d) (decoderEnd d)

-- | Obtain the next byte from the 'Decoder'.
dByte :: Decoder s -> ST s Word8
dByte = dffi glean_pop_value_byte

-- | Obtain a packed natural number from the 'Decoder'.
dNat :: Decoder s -> ST s Word64
dNat = dffi glean_pop_value_nat

-- | Obtain the length of an array from the 'Decoder'.
dArray :: Decoder s -> ST s CSize
dArray = dffi glean_pop_value_array

dSet :: Decoder s -> ST s CSize
dSet = dffi glean_pop_value_set

-- | Obtain a fixed number of bytes from the 'Decoder'. The 'ByteStringRef'
-- has the same lifetime as the 'Decoder'.
dByteStringRef :: Decoder s -> CSize -> ST s ByteStringRef
dByteStringRef d n = do
  p <- dffi (\s e -> glean_pop_value_bytes_ref s e n) d
  return $ ByteStringRef p n

-- | Obtain the given number of bytes from the 'Decoder'.
dBytes :: Decoder s -> CSize -> ST s ByteString
dBytes d n = do
  -- TODO: reimplement based on dByteStringRef
  p <- dffi (\s e -> glean_pop_value_bytes s e n) d
  unsafeIOToST $ FFI.unsafeMallocedByteString (castPtr p) n

-- | Obtain a selector from the 'Decoder'.
dSelector :: Decoder s -> ST s CSize
dSelector = dffi glean_pop_value_selector

-- | Obtain a reference to an RTS-mangled string from the 'Decoder'. The
-- reference has the same lifetime as the 'Decoder'.
dTrustedStringRef :: Decoder s -> ST s (MangledStringRef, CSize)
dTrustedStringRef Decoder{..} = unsafeIOToST $ do
  start <- peek decoderPtr
  n <- glean_pop_value_trusted_string_ref decoderPtr decoderEnd
  end <- peek decoderPtr
  return
    ( MangledStringRef
        { mangledBytes = ByteStringRef
          { byteStringRefPtr = castPtr start
          , byteStringRefSize = fromIntegral $ end `minusPtr` start
          }
        },
      n )


-- | Obtain a UTF8-encoded string from the 'Decoder'.
dString :: Decoder s -> ST s ByteString
dString d = do
  -- TODO: reimplement based on dStringRef
  (p,n) <- dffi2 glean_pop_value_string d
  unsafeIOToST $ FFI.unsafeMallocedByteString (castPtr p) n

-- | Obtain a fact id from the 'Decoder'.
dFact :: Decoder s -> ST s Fid
dFact = dffi glean_pop_value_fact

demangle :: MangledStringRef -> Ptr Word8 -> IO CSize
demangle (MangledStringRef (ByteStringRef p n)) q =
  glean_string_demangle_trusted p n q

-- | Decode a binary value. We allocate a single slot to hold results and reuse
-- it for every decoding step, which is a lot more efficient.
decodeValue :: Decoder s -> Rep Pid -> ST s (Term Fid)
decodeValue d ty = case ty of
  ByteRep -> Byte <$> dByte d
  NatRep -> Nat <$> dNat d
  ArrayRep elty -> do
    size <- dArray d
    case elty of
      ByteRep -> ByteArray <$> dBytes d size
      _ -> Array <$> replicateM (fromIntegral size) (decodeValue d elty)
  TupleRep tys -> Tuple <$> mapM (decodeValue d) tys
  SumRep tys -> do
    sel <- dSelector d
    Alt (fromIntegral sel) <$> decodeValue d (tys !! fromIntegral sel)
  SetRep elty -> do
    size <- dSet d
    Array . uniq <$> replicateM (fromIntegral size) (decodeValue d elty)
  StringRep -> String <$> dString d
  PredicateRep _ -> Ref <$> dFact d

foreign import ccall unsafe glean_push_value_byte
  :: Builder -> Word8 -> IO CString
foreign import ccall unsafe glean_push_value_nat
  :: Builder -> Word64 -> IO CString
foreign import ccall unsafe glean_push_value_array
  :: Builder -> CSize -> IO CString
foreign import ccall unsafe glean_push_value_set
  :: Builder -> CSize -> IO CString
foreign import ccall unsafe glean_push_value_bytes
  :: Builder -> Ptr Word8 -> CSize -> IO CString
foreign import ccall unsafe glean_push_value_selector
  :: Builder -> CSize -> IO CString
foreign import ccall unsafe glean_push_value_string
  :: Builder -> Ptr Word8 -> CSize -> IO CString
foreign import ccall unsafe glean_push_value_fact
  :: Builder -> Fid -> IO CString

foreign import ccall unsafe glean_pop_value_byte
  :: Ptr (Ptr ()) -> Ptr () -> Ptr Word8 -> IO CString
foreign import ccall unsafe glean_pop_value_nat
  :: Ptr (Ptr ()) -> Ptr () -> Ptr Word64 -> IO CString
foreign import ccall unsafe glean_pop_value_array
  :: Ptr (Ptr ()) -> Ptr () -> Ptr CSize -> IO CString
foreign import ccall unsafe glean_pop_value_set
  :: Ptr (Ptr ()) -> Ptr () -> Ptr CSize -> IO CString
foreign import ccall unsafe glean_pop_value_bytes_ref
  :: Ptr (Ptr ()) -> Ptr () -> CSize -> Ptr (Ptr Word8) -> IO CString
foreign import ccall unsafe glean_pop_value_bytes
  :: Ptr (Ptr ()) -> Ptr () -> CSize -> Ptr (Ptr Word8) -> IO CString
foreign import ccall unsafe glean_pop_value_selector
  :: Ptr (Ptr ()) -> Ptr () -> Ptr CSize -> IO CString
foreign import ccall unsafe glean_pop_value_trusted_string_ref
  :: Ptr (Ptr ()) -> Ptr () -> IO CSize
foreign import ccall unsafe glean_pop_value_string
  :: Ptr (Ptr ()) -> Ptr () -> Ptr (Ptr Word8) -> Ptr CSize -> IO CString
foreign import ccall unsafe glean_pop_value_fact
  :: Ptr (Ptr ()) -> Ptr () -> Ptr Fid -> IO CString

foreign import ccall unsafe glean_push_fact
  :: Builder -> Pid -> Builder -> CSize -> IO CString

foreign import ccall unsafe glean_string_demangle_trusted
  :: Ptr Word8 -> CSize -> Ptr Word8 -> IO CSize