{-
  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 RecursiveDo #-}
module Glean.RTS.Typecheck
  ( CompiledTypecheck
  , checkType
  , checkSignature
  ) where

import Control.Monad

import Glean.Angle.Types hiding (Type)
import Glean.Bytecode.Types
import Glean.RTS.Foreign.Bytecode
import Glean.RTS.Types
import Glean.RTS.Bytecode.Code
import Glean.RTS.Bytecode.Gen.Issue

-- | Type tag for Subroutine
data CompiledTypecheck

typecheck
  :: SysCalls
  -> Register 'DataPtr
  -> Register 'DataPtr
  -> Register 'BinaryOutputPtr
  -> Type
  -> Code ()
typecheck syscalls@SysCalls{..} input inputend out = tc
  where
    tc ByteTy = do
      size <- constant 1
      local $ \ptr -> do
        move input ptr
        inputBytes input inputend size
        outputBytes ptr input out
    tc NatTy = local $ \reg -> do
      inputNat input inputend reg
      outputNat reg out
    tc StringTy =
      local $ \ptr -> do
        move input ptr
        inputSkipUntrustedString input inputend
        outputBytes ptr input out
    tc (ArrayTy elty) = local $ \size -> do
      inputNat input inputend size
      outputNat size out
      case derefType elty of
        ByteTy -> local $ \ptr -> do
          move input ptr
          inputBytes input inputend size
          outputBytes ptr input out
        _ -> mdo
          jumpIf0 size end
          loop <- label
          tc elty
          decrAndJumpIfNot0 size loop
          end <- label
          return ()
    tc (RecordTy fields) = mapM_ (tc . fieldDefType) fields
    tc (SumTy fields) = mdo
      local $ \sel -> do
        inputNat input inputend sel
        outputNat sel out
        select sel alts
      raise "selector out of range"
      alts <- forM fields $ \(FieldDef _ ty) -> do
        alt <- label
        tc ty
        jump end
        return alt
      end <- label
      return ()
    tc (SetTy ByteTy) = local $ \size -> do
      inputNat input inputend size
      local $ \set ptr -> do
        newWordSet set
        move input ptr
        inputBytes input inputend size
        insertBytesWordSet set ptr input
        byteSetToByteArray set out
        freeWordSet set
      return ()
    tc (SetTy elty) = local $ \size -> do
      inputNat input inputend size
      local $ \set -> mdo
        newSet set
        jumpIf0 size end
        loop <- label
        output $ \tempOut -> do
          typecheck syscalls input inputend tempOut elty
          insertOutputSet set tempOut
          decrAndJumpIfNot0 size loop
        end <- label
        setToArray set out
        freeSet set
      return ()
    tc (PredicateTy _ (PidRef (Pid pid) _)) = local $ \ide -> do
      t <- constant $ fromIntegral pid
      inputNat input inputend ide
      rename ide t ide
      outputNat ide out
    tc (NamedTy _ (ExpandedType _ ty)) = tc ty
    tc (MaybeTy ty) = mdo
      local $ \sel -> do
        inputNat input inputend sel
        outputNat sel out
        select sel [end,just]
      raise "maybe selector out of range"
      just <- label
      tc ty
      end <- label
      return ()
    tc (EnumeratedTy names) = tcEnum $ fromIntegral $ length names
    tc BooleanTy = tcEnum 2
    tc _ = error "typecheck"

    tcEnum arity = mdo
      k <- constant arity
      local $ \sel -> do
        inputNat input inputend sel
        outputNat sel out
        jumpIfLt sel k end
        raise "selector out of range"
      end <- label
      return ()

data SysCalls = SysCalls {
    rename
    :: Register 'Word
    -> Register 'Word
    -> Register 'Word
    -> Code()

  , newSet
    :: Register 'Word -- (output) set token
    -> Code ()

  , insertOutputSet
    :: Register 'Word -- set token
    -> Register 'BinaryOutputPtr
    -> Code ()

  , setToArray
    :: Register 'Word -- set token
    -> Register 'BinaryOutputPtr -- (output) array
    -> Code ()

  , freeSet
    :: Register 'Word -- set token (invalid after this call)
    -> Code ()

  , newWordSet
    :: Register 'Word -- (output) set token
    -> Code ()

  , insertBytesWordSet
    :: Register 'Word -- set token
    -> Register 'DataPtr
    -> Register 'DataPtr
    -> Code ()

  , wordSetToArray
    :: Register 'Word -- set token
    -> Register 'BinaryOutputPtr -- (output) array
    -> Code ()

  , byteSetToByteArray
    :: Register 'Word -- set token
    -> Register 'BinaryOutputPtr -- (output) array
    -> Code ()

  , freeWordSet
    :: Register 'Word -- set token (invalid after this call)
    -> Code ()
  }

-- | Generate a subroutine which typechecks and substitutes a value. It has
-- the following arguments:
--
-- std::function<Id(Id id, Id type)> - fact substitution
-- binary::Input * - value
-- binary::Output * - substituted value
--
checkType :: Type -> IO (Subroutine CompiledTypecheck)
checkType ty = checkSignature ty $ RecordTy []

-- | Generate a subroutine which typechecks and substitutes a clause. It has
-- the following arguments:
--
-- std::function<Id(Id id, Id type)> - fact substitution
-- set syscalls.
-- const void * - begin of clause/key
-- const void * - end of key/begin of value
-- const void * - end of clause/value
--
-- It returns the substituted clause in the first output and size of the
-- substituted key in the first local register.
--
checkSignature :: Type -> Type -> IO (Subroutine CompiledTypecheck)
checkSignature key_ty val_ty =
  fmap snd $ generate Optimised $
    \rename_
      newSet_ insertOutputSet_ setToArray_ freeSet_
      newWordSet_ insertBytesWordSet_ wordSetToArray_
      byteSetToByteArray_ freeWordSet_
      clause_begin key_end clause_end -> output $ \out -> do
    let syscalls = SysCalls
          { rename = \id pid reg ->
              callFun_2_1 rename_ id pid reg
          , newSet = callFun_0_1 newSet_
          , insertOutputSet = \set out ->
              callFun_2_0 insertOutputSet_ set (castRegister out)
          , setToArray = \set arr ->
              callFun_1_1 setToArray_ set (castRegister arr)
          , freeSet = callFun_1_0 freeSet_
          , newWordSet = callFun_0_1 newWordSet_
          , insertBytesWordSet = \set start end ->
              callFun_3_0
              insertBytesWordSet_ set (castRegister start) (castRegister end)
          , wordSetToArray = \set arr ->
              callFun_1_1 wordSetToArray_ set (castRegister arr)
          , byteSetToByteArray = \set arr ->
              callFun_1_1 byteSetToByteArray_ set (castRegister arr)
          , freeWordSet = callFun_1_0 freeWordSet_
          }
    -- We return the key size in the first local register
    local $ \key_size -> mdo
    typecheck syscalls clause_begin key_end out key_ty
    check "key" clause_begin key_end
    getOutputSize out key_size
    typecheck syscalls clause_begin clause_end out val_ty
    check "value" clause_begin clause_end
    ret
  where
  check str reg1 reg2 = mdo
    jumpIfNe (castRegister reg1) (castRegister reg2) bad
    -- The basic block layout algorithm will make the happy path sequential and
    -- delete this jump.
    jump ok
    bad <- label
    raise $ "extra bytes at end of " <> str
    ok <- label
    return ()