C.hs

{-# LANGUAGE CPP                 #-}
{-# LANGUAGE BangPatterns        #-}
{-# LANGUAGE BlockArguments      #-}
{-# LANGUAGE LambdaCase          #-}
{-# LANGUAGE PatternSynonyms     #-}
{-# LANGUAGE RankNTypes          #-}
{-# LANGUAGE RecordWildCards     #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE TypeApplications    #-}
{-# LANGUAGE ViewPatterns        #-}
-- | How to call foreign functions following the C-style calling convention.

#ifdef FFI_ENABLED

module Cryptol.Eval.FFI.C (callForeignC) where
import           Cryptol.Eval.FFI.ForeignSrc

import           Cryptol.Eval
import           Cryptol.TypeCheck.AST
import           Cryptol.TypeCheck.FFI.FFIType

import           Control.Exception(bracket_)
import           Data.Foldable
import           Data.IORef
import           Data.Proxy
import           Data.Ratio
import           Data.Traversable
import           Data.Word
import           Foreign
import           Foreign.C.Types
import           GHC.Float
import           LibBF                         (bfFromDouble, bfToDouble,
                                                pattern NearEven)
import           Numeric.GMP.Raw.Unsafe
import           Numeric.GMP.Utils

import           Cryptol.Backend
import           Cryptol.Backend.Concrete
import           Cryptol.Backend.FloatHelpers
import           Cryptol.Backend.Monad
import           Cryptol.Backend.SeqMap
import           Cryptol.Eval.Type
import           Cryptol.Eval.Value
import           Cryptol.TypeCheck.Solver.InfNat
import           Cryptol.Utils.RecordMap


-- | Methods for obtaining a return value. The producer of this type must supply
-- both 1) a polymorphic IO object directly containing a return value that the
-- consumer can instantiate at any 'FFIRet' type, and 2) an effectful function
-- that takes some output arguments and modifies what they are pointing at to
-- store a return value. The consumer can choose which one to use.
data GetRet = GetRet
  { getRetAsValue   :: forall a. FFIRet a => IO a
  , getRetAsOutArgs :: [SomeFFIArg] -> IO () }

-- | Operations needed for returning a basic reference type.
data BasicRefRet a = BasicRefRet
  { -- | Initialize the object before passing to foreign function.
    initBasicRefRet    :: Ptr a -> IO ()
    -- | Free the object after returning from foreign function and obtaining
    -- return value.
  , clearBasicRefRet   :: Ptr a -> IO ()
    -- | Convert the object to a Cryptol value.
  , marshalBasicRefRet :: a -> Eval (GenValue Concrete) }


-- | Call a foreign function, which follows the C style calling convetion.
callForeignC ::
  Name                           {- ^ Name of foregin function -} ->
  FFIFunType FFIType             {- ^ FFI type -} ->
  ForeignImpl                    {- ^ Address of foreign worker -} ->
  TypeEnv                        {- ^ Values for numeric type parametres -} ->
  [(FFIType, GenValue Concrete)] {- ^ Function arguments -} ->
  Eval (GenValue Concrete)
callForeignC name FFIFunType {..} impl tenv typesAndVals =
  marshalArgs typesAndVals \inArgs -> do
    tyArgs <- traverse marshalTyArg ffiTParams
    let tyInArgs = tyArgs ++ inArgs
    marshalRet ffiRetType GetRet
      { getRetAsValue = callForeignImpl impl tyInArgs
      , getRetAsOutArgs = callForeignImpl impl . (tyInArgs ++) }
  where
  -- Look up the value of a type parameter in the type environment and marshal
  -- it.
  marshalTyArg :: TParam -> Eval SomeFFIArg
  marshalTyArg tp
    | n <= toInteger (maxBound :: CSize) =
      pure $ SomeFFIArg @CSize $ fromInteger n
    | otherwise = raiseError Concrete $ FFITypeNumTooBig name tp n
    where n = evalFinType $ TVar $ TVBound tp

  -- Marshal the given value as the given FFIType and call the given function
  -- with the results. A single Cryptol argument may correspond to any number of
  -- C arguments, so the callback takes a list.
  --
  -- NOTE: the result must be used only in the callback since it may have a
  -- limited lifetime (e.g. pointer returned by alloca).
  marshalArg ::
    FFIType ->
    GenValue Concrete ->
    ([SomeFFIArg] -> Eval a) ->
    Eval a

  marshalArg FFIBool val f = f [SomeFFIArg @Word8 (fromBool (fromVBit val))]

  marshalArg (FFIBasic (FFIBasicVal t)) val f =
    getMarshalBasicValArg t \doExport ->
    do arg <- doExport val
       f [SomeFFIArg arg]

  marshalArg (FFIBasic (FFIBasicRef t)) val f =
    getMarshalBasicRefArg t \doExport  ->
    -- Since we need to do Eval actions in an IO callback, we need to manually
    -- unwrap and wrap the Eval datatype
    Eval \stk ->
      doExport val \arg ->
        with arg \ptr ->
          runEval stk (f [SomeFFIArg ptr])

  marshalArg (FFIArray (map evalFinType -> sizes) bt) val f =
    case bt of

      FFIBasicVal t ->
        getMarshalBasicValArg t \doExport  ->
          -- Since we need to do Eval actions in an IO callback,
          -- we need to manually unwrap and wrap the Eval datatype
          Eval \stk ->
            marshalArrayArg stk \v k ->
              k =<< runEval stk (doExport v)

      FFIBasicRef t -> Eval \stk ->
        getMarshalBasicRefArg t \doExport ->
        marshalArrayArg stk doExport

    where marshalArrayArg stk doExport =
            allocaArray (fromInteger (product sizes)) \ptr -> do
              -- Traverse the nested sequences and write the elements to the
              -- array in order.
              -- ns is the dimensions of the values we are currently
              -- processing.
              -- vs is the values we are currently processing.
              -- nvss is the stack of previous ns and vs that we keep track of
              -- that we push onto when we start processing a nested sequence
              -- and pop off when we finish processing the current ones.
              -- i is the index into the array.

              let
                  -- write next element of multi-dimensional array
                  write (n:ns) (v:vs) nvss !i =
                    do vs' <- traverse (runEval stk)
                                       (enumerateSeqMap n (fromVSeq v))
                       write ns vs' ((n, vs):nvss) i

                  -- write next element in flat array
                  write [] (v:vs) nvss !i =
                    doExport v \rep ->
                      do pokeElemOff ptr i rep
                         write [] vs nvss (i + 1)

                  -- finished with flat array, do next element of multi-d array
                  write ns [] ((n, vs):nvss) !i = write (n:ns) vs nvss i

                  -- done
                  write _ _ [] _ = pure ()


              write sizes [val] [] 0
              runEval stk $ f [SomeFFIArg ptr]

  marshalArg (FFITuple types) val f =
    do vals <- sequence (fromVTuple val)
       marshalArgs (types `zip` vals) f

  marshalArg (FFIRecord typeMap) val f =
    do vals <- traverse (`lookupRecord` val) (displayOrder typeMap)
       marshalArgs (displayElements typeMap `zip` vals) f

  -- Call marshalArg on a bunch of arguments and collect the results together
  -- (in the order of the arguments).
  marshalArgs ::
    [(FFIType, GenValue Concrete)] ->
    ([SomeFFIArg] -> Eval a) ->
    Eval a
  marshalArgs tysVs f = go tysVs []
    where
    go [] args = f (concat (reverse args))
    go ((t, v):tvs) prevArgs =
      marshalArg t v \currArgs ->
      go tvs (currArgs : prevArgs)

  -- Given an FFIType and a GetRet, obtain a return value and convert it to a
  -- Cryptol value. The return value is obtained differently depending on the
  -- FFIType.
  marshalRet :: FFIType -> GetRet -> Eval (GenValue Concrete)
  marshalRet FFIBool gr =
    do rep <- io (getRetAsValue gr @Word8)
       pure (VBit (toBool rep))

  marshalRet (FFIBasic (FFIBasicVal t)) gr =
    getMarshalBasicValRet t \doImport ->
      do rep <- io (getRetAsValue gr)
         doImport rep

  marshalRet (FFIBasic (FFIBasicRef t)) gr =
    getBasicRefRet t \how ->
    Eval             \stk ->
    alloca           \ptr ->
    bracket_ (initBasicRefRet how ptr) (clearBasicRefRet how ptr)
      do getRetAsOutArgs gr [SomeFFIArg ptr]
         rep <- peek ptr
         runEval stk (marshalBasicRefRet how rep)

  marshalRet (FFIArray (map evalFinType -> sizes) bt) gr =
    Eval \stk -> do
    let totalSize = fromInteger (product sizes)
        getResult marshal ptr = do
          getRetAsOutArgs gr [SomeFFIArg ptr]
          let tyv = case bt of
                FFIBasicVal bv -> case bv of
                  FFIWord len _ -> TVSeq len TVBit
                  FFIFloat e p _ -> TVFloat e p
                FFIBasicRef br -> case br of
                  FFIInteger Nothing -> TVInteger
                  FFIInteger (Just z) -> TVIntMod $ evalFinType z
                  FFIRational -> TVRational

          let build (n:ns) !i = do
                -- We need to be careful to actually run this here and not just
                -- stick the IO action into the sequence with io, or else we
                -- will read from the array after it is deallocated.
                vs <- for [0 .. fromInteger n - 1] \j ->
                  build ns (i * fromInteger n + j)
                runEval stk $ 
                  mkSeq Concrete (Nat n) tyv (finiteSeqMap Concrete (map pure vs))
              build [] !i = do
                e <- peekElemOff ptr i
                runEval stk (marshal e)

          build sizes 0

    case bt of

      FFIBasicVal t ->
        getMarshalBasicValRet t \doImport ->
        allocaArray totalSize (getResult doImport)

      FFIBasicRef t ->
        getBasicRefRet t      \how ->
        allocaArray totalSize \ptr ->
          do let forEach f = for_ [0 .. totalSize - 1] (f . advancePtr ptr)
             bracket_ (forEach (initBasicRefRet how))
                      (forEach (clearBasicRefRet how))
                      (getResult (marshalBasicRefRet how) ptr)

  marshalRet (FFITuple types) gr = VTuple <$> marshalMultiRet types gr

  marshalRet (FFIRecord typeMap) gr =
    VRecord . recordFromFields . zip (displayOrder typeMap) <$>
      marshalMultiRet (displayElements typeMap) gr

  -- Obtain multiple return values as output arguments for a composite return
  -- type. Each return value is fully evaluated but put back in an Eval since
  -- VTuple and VRecord expect it.
  marshalMultiRet :: [FFIType] -> GetRet -> Eval [Eval (GenValue Concrete)]
  -- Since IO callbacks are involved we just do the whole thing in IO and wrap
  -- it in an Eval at the end. This should be fine since we are not changing
  -- the (Cryptol) call stack.
  marshalMultiRet types gr = Eval \stk -> do
    -- We use this IORef hack here since we are calling marshalRet recursively
    -- but marshalRet doesn't let us return any extra information from the
    -- callback through to the result of the function. So we remember the result
    -- as a side effect.
    vals <- newIORef []
    let go [] args = getRetAsOutArgs gr args
        go (t:ts) prevArgs = do
          val <- runEval stk $ marshalRet t $ getRetFromAsOutArgs \currArgs ->
            go ts (prevArgs ++ currArgs)
          modifyIORef' vals (val :)
    go types []
    map pure <$> readIORef vals

  -- | Call the callback with a 'BasicRefRet' for the given type.
  getBasicRefRet :: FFIBasicRefType ->
    (forall a. Storable a => BasicRefRet a -> b) -> b
  getBasicRefRet (FFIInteger mbMod) f = f BasicRefRet
    { initBasicRefRet = mpz_init
    , clearBasicRefRet = mpz_clear
    , marshalBasicRefRet = \mpz -> do
        n <- io $ peekInteger' mpz
        VInteger <$>
          case mbMod of
            Nothing -> pure n
            Just m  -> intToZn Concrete (evalFinType m) n }
  getBasicRefRet FFIRational f = f BasicRefRet
    { initBasicRefRet = mpq_init
    , clearBasicRefRet = mpq_clear
    , marshalBasicRefRet = \mpq -> do
        r <- io $ peekRational' mpq
        pure $ VRational $ SRational (numerator r) (denominator r) }

  -- Evaluate a finite numeric type expression.
  evalFinType :: Type -> Integer
  evalFinType = finNat' . evalNumType tenv



-- | Given a way to 'getRetAsOutArgs', create a 'GetRet', where the
-- 'getRetAsValue' simply allocates a temporary space to call 'getRetAsOutArgs'
-- on. This is useful for return types that we know how to obtain directly as a
-- value but need to obtain as an output argument when multiple return values
-- are involved.
getRetFromAsOutArgs :: ([SomeFFIArg] -> IO ()) -> GetRet
getRetFromAsOutArgs f = GetRet
  { getRetAsValue = alloca \ptr -> do
      f [SomeFFIArg ptr]
      peek ptr
  , getRetAsOutArgs = f }

-- | Given a 'FFIBasicValType', call the callback with a marshalling function
-- that marshals values to the 'FFIArg' type corresponding to the
-- 'FFIBasicValType'. The callback must be able to handle marshalling functions
-- that marshal to any 'FFIArg' type.
getMarshalBasicValArg ::
  FFIBasicValType ->
  (forall rep.
      FFIArg rep =>
      (GenValue Concrete -> Eval rep) ->
      result) ->
  result

getMarshalBasicValArg (FFIWord _ s) f = withWordType s \(_ :: p t) ->
  f @t $ fmap (fromInteger . bvVal) . fromVWord Concrete "getMarshalBasicValArg"

getMarshalBasicValArg (FFIFloat _ _ s) f =
  case s of
    -- LibBF can only convert to 'Double' directly, so we do that first then
    -- convert to 'Float', which should not result in any loss of precision if
    -- the original data was 32-bit anyways.
    FFIFloat32 -> f $ pure . CFloat . double2Float . toDouble
    FFIFloat64 -> f $ pure . CDouble . toDouble
  where
  toDouble = fst . bfToDouble NearEven . bfValue . fromVFloat

-- | Given a 'FFIBasicValType', call the callback with an unmarshalling function
-- from the 'FFIRet' type corresponding to the 'FFIBasicValType' to Cryptol
-- values. The callback must be able to handle unmarshalling functions from any
-- 'FFIRet' type.
getMarshalBasicValRet :: FFIBasicValType ->
  (forall a. FFIRet a => (a -> Eval (GenValue Concrete)) -> b) -> b
getMarshalBasicValRet (FFIWord n s) f = withWordType s \(_ :: p t) ->
  f @t $ word Concrete n . toInteger
getMarshalBasicValRet (FFIFloat e p s) f =
  case s of
    FFIFloat32 -> f $ toValue . \case CFloat x -> float2Double x
    FFIFloat64 -> f $ toValue . \case CDouble x -> x
  where toValue = pure . VFloat . BF e p . bfFromDouble

-- | Call the callback with the Word type corresponding to the given
-- 'FFIWordSize'.
withWordType :: FFIWordSize ->
  (forall a. (FFIArg a, FFIRet a, Integral a) => Proxy a -> b) -> b
withWordType FFIWord8  f = f $ Proxy @Word8
withWordType FFIWord16 f = f $ Proxy @Word16
withWordType FFIWord32 f = f $ Proxy @Word32
withWordType FFIWord64 f = f $ Proxy @Word64

-- | Given a 'FFIBasicRefType', call the callback with a marshalling function
-- that takes a Cryptol value and calls its callback with the 'Storable' type
-- corresponding to the 'FFIBasicRefType'.
getMarshalBasicRefArg :: FFIBasicRefType ->
  (forall rep.
      Storable rep =>
      (GenValue Concrete -> (rep -> IO val) -> IO val) ->
      result) ->
  result
getMarshalBasicRefArg (FFIInteger _) f = f \val g ->
  withInInteger' (fromVInteger val) g
getMarshalBasicRefArg FFIRational f = f \val g -> do
  let SRational {..} = fromVRational val
  withInRational' (sNum % sDenom) g

#else
module Cryptol.Eval.FFI.C where
#endif