{-
  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.
-}

module Glean.RTS.Foreign.Bytecode
  ( Subroutine
  , SubroutineCode(..)
  , subroutine
  , inspect
  , size
  ) where

import Control.Monad
import Data.ByteString (ByteString)
import qualified Data.Vector.Storable as V
import Data.Word (Word64)
import Foreign (withArray, withArrayLen)
import Foreign.C.String
import Foreign.C.Types
import Foreign.ForeignPtr
import Foreign.Ptr
import Foreign.Storable (peekElemOff)
import System.IO.Unsafe (unsafePerformIO, unsafeDupablePerformIO)

import Glean.FFI
import Util.FFI

data SubroutineCode = SubroutineCode
  { -- | Instructions
    subInsns :: !(V.Vector Word64)

    -- | Number of inputs
  , subInputs :: {-# UNPACK #-} !Word64

    -- | Number of outputs
  , subOutputs :: {-# UNPACK #-} !Word64

    -- | Number of local registers
  , subLocals :: {-# UNPACK #-} !Word64

    -- | Constants (copied into local registers at beginning)
  , subConstants :: !(V.Vector Word64)

    -- | Literals
  , subLiterals :: [ByteString]
  }

-- | A bytecode subroutine. The s parameter describes the kind of
-- subroutine this is (e.g. CompiledTypecheck, CompiledQuery), to
-- avoid us mixing up the two.
newtype Subroutine s = Subroutine (ForeignPtr (Subroutine s))

instance Object (Subroutine s) where
  wrap = Subroutine
  unwrap (Subroutine p) = p
  destroy = glean_subroutine_free

-- | Create a subroutine
subroutine
  :: V.Vector Word64 -- ^ instructions
  -> Word64 -- ^ number of inputs
  -> Word64 -- ^ number of outputs
  -> Word64 -- ^ number of local registers
  -> [Word64] -- ^ constants
  -> [ByteString] -- ^ literals
  -> IO (Subroutine s)
subroutine code inputs outputs locals consts lits =
  V.unsafeWith code $ \code_ptr ->
  withArrayLen consts $ \consts_len consts_ptr ->
  withMany (\s -> unsafeWithBytes s . curry) lits $ \ps ->
  let (lit_ptrs, lit_sizes) = unzip ps
  in
  withArray lit_ptrs $ \p_lit_ptrs ->
  withArray lit_sizes $ \p_lit_sizes ->
  construct $ invoke $ glean_subroutine_new
    code_ptr
    (fromIntegral $ V.length code)
    (fromIntegral inputs)
    (fromIntegral outputs)
    (fromIntegral locals)
    consts_ptr
    (fromIntegral consts_len)
    p_lit_ptrs
    p_lit_sizes
    (fromIntegral $ length ps)

size :: Subroutine s -> Int
size sub = unsafeDupablePerformIO $ with sub $ \p_sub -> do
  fromIntegral <$> glean_subroutine_size p_sub

inspect :: Subroutine s -> SubroutineCode
inspect sub = unsafePerformIO $ with sub $ \p_sub -> do
  ( insns_ptr, insns_size
    , inputs, outputs, locals
    , consts_ptr, consts_size
    , lit_count) <- invoke $ glean_subroutine_inspect p_sub
  insns <- V.generateM (fromIntegral insns_size) $ peekElemOff insns_ptr
  consts <- V.generateM (fromIntegral consts_size) $ peekElemOff consts_ptr
  lits <- forM [1 .. lit_count] $ \i -> do
    (lit_ptr, lit_size) <- invoke $ glean_subroutine_literal p_sub (i-1)
    copyByteString lit_ptr lit_size
  return SubroutineCode
    { subInsns = insns
    , subInputs = inputs
    , subOutputs = outputs
    , subLocals = locals
    , subConstants = consts
    , subLiterals = lits
    }

foreign import ccall unsafe glean_subroutine_new
  :: Ptr Word64
  -> CSize
  -> CSize
  -> CSize
  -> CSize
  -> Ptr Word64
  -> CSize
  -> Ptr (Ptr ())
  -> Ptr CSize
  -> CSize
  -> Ptr (Ptr (Subroutine s))
  -> IO CString
foreign import ccall unsafe "&glean_subroutine_free" glean_subroutine_free
  :: FunPtr (Ptr (Subroutine s) -> IO ())

foreign import ccall unsafe glean_subroutine_inspect
  :: Ptr (Subroutine s)
  -> Ptr (Ptr Word64)
  -> Ptr CSize
  -> Ptr Word64
  -> Ptr Word64
  -> Ptr Word64
  -> Ptr (Ptr Word64)
  -> Ptr CSize
  -> Ptr CSize
  -> IO ()

foreign import ccall unsafe glean_subroutine_size
  :: Ptr (Subroutine s)
  -> IO CSize

foreign import ccall unsafe glean_subroutine_literal
  :: Ptr (Subroutine s)
  -> CSize
  -> Ptr (Ptr ())
  -> Ptr CSize
  -> IO ()