On close, wait to receive status on client.

-- Copyright (c) 2016, Google Inc.
-- All rights reserved.
--
-- Redistribution and use in source and binary forms, with or without
-- modification, are permitted provided that the following conditions are met:
--     * Redistributions of source code must retain the above copyright
--       notice, this list of conditions and the following disclaimer.
--     * Redistributions in binary form must reproduce the above copyright
--       notice, this list of conditions and the following disclaimer in the
--       documentation and/or other materials provided with the distribution.
--     * Neither the name of Google Inc. nor the
--       names of its contributors may be used to endorse or promote products
--       derived from this software without specific prior written permission.
--
-- THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" AND
-- ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED
-- WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
-- DISCLAIMED. IN NO EVENT SHALL Google Inc. BE LIABLE FOR ANY
-- DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES
-- (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES;
-- LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND
-- ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
-- (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS
-- SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
--
--------------------------------------------------------------------------------
{-# LANGUAGE ScopedTypeVariables, OverloadedStrings, ExistentialQuantification, RecordWildCards #-}
module Network.Grpc.Core.Call where

import           Control.Concurrent
import           Control.Exception
import           Control.Monad

import qualified Data.ByteString              as B
import qualified Data.ByteString.Lazy         as L
import           Data.Monoid ((<>), Last(..))
import           Data.IORef

import           Foreign.C.Types              as C
import qualified Foreign.Marshal.Alloc        as C
import qualified Foreign.Marshal.Utils        as C
import qualified Foreign.Ptr                  as C
import           Foreign.Ptr (Ptr)
import qualified Foreign.Storable             as C

-- transformers
import           Control.Monad.IO.Class
import           Control.Monad.Trans.Class
import           Control.Monad.Trans.Except
import           Control.Monad.Trans.Reader

import qualified Network.Grpc.CompletionQueue as CQ
import Network.Grpc.Lib.PropagationBits
{#import Network.Grpc.Lib.ByteBuffer#}
{#import Network.Grpc.Lib.Metadata#}
{#import Network.Grpc.Lib.TimeSpec#}
{#import Network.Grpc.Lib.Core#}

#include <grpc/grpc.h>
#include "hs_grpc.h"

{#context lib = "grpc" prefix = "grpc" #}

data Deadline
  = AbsoluteDeadline TimeSpec
  | RelativeDeadline Int --seconds

data ClientContext = ClientContext
  { ccChannel :: Channel
  , ccCQ :: CompletionQueue
  , ccWorker :: CQ.Worker
  }

data CallOptions = CallOptions
  { coDeadline :: Maybe Deadline
  , coParentContext :: Maybe () -- todo
  , coPropagationMask :: Maybe PropagationMask
  , coMetadata :: [Metadata]
  }

instance Monoid CallOptions where
  mempty = CallOptions Nothing Nothing Nothing []
  mappend (CallOptions a b c d) (CallOptions a' b' c' d') =
    CallOptions
      (getLast (Last a <> Last a'))
      (getLast (Last b <> Last b'))
      (getLast (Last c <> Last c'))
      (d <> d')

withAbsoluteDeadline :: TimeSpec -> CallOptions
withAbsoluteDeadline deadline =
  mempty { coDeadline = Just (AbsoluteDeadline deadline) }

withRelativeDeadlineSeconds :: Int -> CallOptions
withRelativeDeadlineSeconds seconds =
  mempty { coDeadline = Just (RelativeDeadline seconds) }

withParentContext :: () -> CallOptions
withParentContext ctx =
  mempty { coParentContext = Just ctx }

withParentContextPropagating :: () -> PropagationMask -> CallOptions
withParentContextPropagating ctx prop =
  mempty { coParentContext   = Just ctx
         , coPropagationMask = Just prop }

withMetadata :: [Metadata] -> CallOptions
withMetadata md =
  mempty { coMetadata = md }

resolveDeadline :: CallOptions -> IO TimeSpec
resolveDeadline co =
  case coDeadline co of
    Nothing -> return gprInfFuture
    Just (AbsoluteDeadline deadline) -> return deadline
    Just (RelativeDeadline seconds) ->
      secondsFromNow (fromIntegral seconds)

newClientContext :: Channel -> IO ClientContext
newClientContext chan = do
  cq <- completionQueueCreate reservedPtr
  cqt <- CQ.startCompletionQueueThread cq
  return (ClientContext chan cq cqt)

destroyClientContext :: ClientContext -> IO ()
destroyClientContext (ClientContext _ cq w) = do
  completionQueueShutdown cq
  CQ.waitWorkerTermination w

type MethodName = B.ByteString
type Arg = B.ByteString

-- | Run the IO function once, cache it in the MVar.
onceMVar :: MVar (Maybe a) -> IO a -> IO a
onceMVar mvar io = modifyMVar mvar $ \x ->
  case x of
    Just x' -> return (x, x')
    Nothing -> do
      x' <- io
      return (Just x', x')

reservedPtr :: Ptr ()
reservedPtr = C.nullPtr

data UnaryResult a = UnaryResult [Metadata] [Metadata] a deriving Show

callUnary :: ClientContext -> CallOptions -> MethodName -> Arg -> IO (RpcReply (UnaryResult L.ByteString))
callUnary ctx@(ClientContext chan cq _) co method arg = do
  deadline <- resolveDeadline co
  bracket (grpcChannelCreateCall (cChannel chan) C.nullPtr propagateDefaults cq method (cHost chan) deadline) grpcCallDestroy $ \call0 -> newMVar call0 >>= \mcall -> do
    crw <- newClientReaderWriter ctx mcall

    sendInitOp <- opSendInitialMetadata (coMetadata co)
    sendCloseOp <- opSendCloseFromClient
    sendMessageOp <- opSendMessage arg
    recvStatusOp <- opRecvStatusOnClient crw
    recvMessageOp <- opRecvMessage
    recvInitialMetadataOp <- opRecvInitialMetadata crw

    res <- callBatch crw [
          OpX sendInitOp
        , OpX sendCloseOp
        , OpX recvInitialMetadataOp
        , OpX recvMessageOp
        , OpX sendMessageOp
        , OpX recvStatusOp
        ]
    case res of
      RpcOk _ -> do
        (RpcStatus trailMD status statusDetails) <- opRead recvStatusOp
        case status of
          StatusOk -> do
            initMD <- opRead recvInitialMetadataOp
            answ <- opRead recvMessageOp
            let answ' = maybe L.empty id answ
            return (RpcOk (UnaryResult initMD trailMD answ'))
          _ -> return (RpcError (StatusError status statusDetails))
      RpcError err -> return (RpcError err)

callDownstream :: ClientContext -> CallOptions -> MethodName -> Arg -> IO (RpcReply (Client B.ByteString L.ByteString))
callDownstream ctx@(ClientContext chan cq _) co method arg = do
  deadline <- resolveDeadline co
  mcall <- grpcChannelCreateCall (cChannel chan) C.nullPtr propagateDefaults cq method (cHost chan) deadline >>= newMVar
  crw <- newClientReaderWriter ctx mcall

  sendInitOp <- opSendInitialMetadata (coMetadata co)
  sendCloseOp <- opSendCloseFromClient
  sendMessageOp <- opSendMessage arg

  res <- callBatch crw [
        OpX sendInitOp
      , OpX sendCloseOp
      , OpX sendMessageOp
      ]
  case res of
    RpcOk _ -> do
      res' <- callBatchStatusOnClient crw
      case res' of
        RpcOk _ -> return (RpcOk (Client crw defaultEncoder defaultDecoder))
        RpcError err -> return (RpcError err)
    RpcError err -> return (RpcError err)

callUpstream :: ClientContext -> CallOptions -> MethodName -> IO (RpcReply (Client B.ByteString L.ByteString))
callUpstream ctx@(ClientContext chan cq _) co method = do
  deadline <- resolveDeadline co
  mcall <- grpcChannelCreateCall (cChannel chan) C.nullPtr propagateDefaults cq method (cHost chan) deadline >>= newMVar
  crw <- newClientReaderWriter ctx mcall

  sendInitOp <- opSendInitialMetadata (coMetadata co)

  res <- callBatch crw [ OpX sendInitOp ]
  case res of
    RpcOk _ -> do
      res' <- callBatchStatusOnClient crw
      case res' of
        RpcOk _ -> return (RpcOk (Client crw defaultEncoder defaultDecoder))
        RpcError err -> return (RpcError err)
    RpcError err -> return (RpcError err)

callBidi :: ClientContext -> CallOptions -> MethodName -> IO (RpcReply (Client B.ByteString L.ByteString))
callBidi ctx@(ClientContext chan cq _) co method = do
  deadline <- resolveDeadline co
  mcall <- grpcChannelCreateCall (cChannel chan) C.nullPtr propagateDefaults cq method (cHost chan) deadline >>= newMVar

  crw <- newClientReaderWriter ctx mcall
  sendInitOp <- opSendInitialMetadata (coMetadata co)

  res <- callBatch crw [ OpX sendInitOp ]
  case res of
    RpcOk _ -> do
      res' <- callBatchStatusOnClient crw
      case res' of
        RpcOk _ -> return (RpcOk (Client crw defaultEncoder defaultDecoder))
        RpcError err -> return (RpcError err)
    RpcError err -> return (RpcError err)

data Client req resp = Client
  { clientCrw :: ClientReaderWriter
  , clientEncoder :: Encoder req
  , clientDecoder :: Decoder resp
  }

type Decoder a = L.ByteString -> IO (RpcReply a)
type Encoder a = a -> IO (RpcReply B.ByteString)

defaultDecoder :: L.ByteString -> IO (RpcReply L.ByteString)
defaultDecoder bs = return (RpcOk bs)

defaultEncoder :: B.ByteString -> IO (RpcReply B.ByteString)
defaultEncoder bs = return (RpcOk bs)

data ClientReaderWriter = ClientReaderWriter {
  context :: ClientContext,
  callMVar_ :: MVar Call,
  initialMDRef :: !(IORef (Maybe [Metadata])),
  trailingMDRef :: !(IORef (Maybe [Metadata])),
  statusFromServer :: !(MVar RpcStatus),
  statusFromServerTag :: !(MVar CQ.EventDesc)
}

newClientReaderWriter :: ClientContext -> MVar Call -> IO ClientReaderWriter
newClientReaderWriter ctx mcall = do
  initMD <- newIORef Nothing
  trailMD <- newIORef Nothing
  status <- newEmptyMVar
  statusEvent <- newEmptyMVar
  return (ClientReaderWriter ctx mcall initMD trailMD status statusEvent)

data OpX = forall t. OpX (OpT t)

data OpArray = OpArray { opArrPtr           :: !GrpcOpPtr
                       , opArrLen           :: !CULong
                       , opArrFinishAndFree :: !(IO ())
                       }

toArray :: [OpX] -> IO OpArray
toArray ops = do
  aptr <- C.mallocBytes (length ops * {#sizeof grpc_op#})
  let ptrs = iterate (`C.plusPtr` {#sizeof grpc_op#}) aptr
      ops' = concatMap (\(OpX op) -> opAdd op) ops
      free = C.free aptr >> sequence_ (map (\(OpX op) -> opFinish op) ops)
      write op p = do
        {#set grpc_op->flags#} p 0
        {#set grpc_op->reserved#} p C.nullPtr
        op p
  zipWithM_ write ops' ptrs
  return $! OpArray aptr (fromIntegral $ length ops) free

callBatch :: ClientReaderWriter -> [OpX] -> IO (RpcReply ())
callBatch crw ops = do
  let ctx = context crw
  arr <- toArray ops
  let onBatchComplete = opArrFinishAndFree arr
  CQ.withEvent (ccWorker ctx) onBatchComplete $ \eDesc -> do
    callStatus <- withMVar (callMVar_ crw) $ \call ->
      grpcCallStartBatch call (opArrPtr arr) (opArrLen arr) (CQ.eventTag eDesc) reservedPtr
    case callStatus of
      CallOk -> do
        e <- CQ.interruptibleWaitEvent eDesc
        case e of
          QueueOpComplete OpSuccess _ -> return (RpcOk ())
          QueueOpComplete OpError _ -> return (RpcError (Error "callBatch: op error"))
          QueueTimeOut -> return (RpcError DeadlineExceeded)
          QueueShutdown -> return (RpcError (Error "queue shutdown"))
      _ -> do
         return (RpcError (CallErrorStatus callStatus))

callBatchStatusOnClient :: ClientReaderWriter -> IO (RpcReply ())
callBatchStatusOnClient crw@ClientReaderWriter{..} = do
  tag <- tryReadMVar statusFromServerTag
  case tag of
    Just _ -> return (RpcOk ()) -- already did this before
    Nothing -> do
      statusOp <- opRecvStatusOnClient crw
      arr <- toArray [ OpX statusOp ]
      let onBatchComplete = opArrFinishAndFree arr
      eDesc <- CQ.allocateEvent (ccWorker context) onBatchComplete
      putMVar statusFromServerTag eDesc
      callStatus <- withMVar callMVar_ $ \call ->
        grpcCallStartBatch call (opArrPtr arr) (opArrLen arr) (CQ.eventTag eDesc) reservedPtr
      case callStatus of
        CallOk -> return (RpcOk ())
        _ -> return (RpcError (CallErrorStatus callStatus))

data OpT out = Op
  { opAdd    :: [Ptr GrpcOp -> IO ()]
  , opValue  :: IORef out
  , opFinish :: IO ()
  }

opRead :: OpT a -> IO a
opRead op = readIORef (opValue op)

opRecvInitialMetadata :: ClientReaderWriter -> IO (OpT [Metadata])
opRecvInitialMetadata crw = do
  arr <- mallocMetadataArray
  value <- newIORef (error "opRecvInitialMetadata never finished")
  let
    add =
      [ \p -> do
        {#set grpc_op->op#} p (fromIntegral (fromEnum OpRecvInitialMetadata))
        {#set grpc_op->data.recv_initial_metadata#} p arr
      ]
    finish = do
      mds <- readMetadataArray arr
      writeIORef (initialMDRef crw) (Just mds)
      writeIORef value mds
      freeMetadataArray arr
  return (Op add value finish)

opSendMessage :: B.ByteString -> IO (OpT ())
opSendMessage bs = do
  bb <- fromByteString bs
  value <- newIORef ()
  let
    add =
      [ \p -> do
        {#set grpc_op->op#} p (fromIntegral (fromEnum OpSendMessage))
        {#set grpc_op->data.send_message#} p bb
      ]
    finish =
      byteBufferDestroy bb
  return (Op add value finish)

opRecvMessage :: IO (OpT (Maybe L.ByteString))
opRecvMessage = do
  bbptr <- C.malloc :: IO (Ptr (Ptr CByteBuffer))
  value <- newIORef Nothing
  let
    add =
      [ \p -> do
        {#set grpc_op->op#} p (fromIntegral (fromEnum OpRecvMessage))
        {#set grpc_op->data.recv_message#} p bbptr
      ]
    finish = do
      bb <- C.peek bbptr
      C.free bbptr
      writeIORef value =<< if bb /= C.nullPtr
        then do
          lbs <- toLazyByteString bb
          byteBufferDestroy bb
          return (Just lbs)
        else return Nothing
  return (Op add value finish)

opSendInitialMetadata :: [Metadata] -> IO (OpT ())
opSendInitialMetadata elems = do
  (mdArrPtr, free) <- mallocMetadata elems
  value <- newIORef ()
  let
    add =
      [ \p -> do
        {#set grpc_op->op#} p (fromIntegral (fromEnum OpSendInitialMetadata))
        {#set grpc_op->data.send_initial_metadata.count#} p (fromIntegral (length elems))
        {#set grpc_op->data.send_initial_metadata.metadata#} p (C.castPtr mdArrPtr)
        {#set grpc_op->data.send_initial_metadata.maybe_compression_level.is_set#} p 0
        {#set grpc_op->data.send_initial_metadata.maybe_compression_level.level#} p 0
      ]
    finish = do
      free
  return (Op add value finish)

data RpcStatus = RpcStatus [Metadata] StatusCode B.ByteString
  deriving Show

opRecvStatusOnClient :: ClientReaderWriter -> IO (OpT RpcStatus)
opRecvStatusOnClient (ClientReaderWriter{..}) = do
  trailingMetadataArrPtr <- mallocMetadataArray
  statusCodePtr <- C.malloc :: IO (Ptr StatusCodeT)
  ptrPtrStatusStr <- C.new C.nullPtr :: IO (Ptr (Ptr CChar))
  capacityPtr <- C.new 0 :: IO (Ptr SizeT)
  value <- newIORef (error "opRecvStatusOnClient never ran")
  let
    add =
      [ \p -> do
        {#set grpc_op->op#} p (fromIntegral (fromEnum OpRecvStatusOnClient))
        {#set grpc_op->data.recv_status_on_client.trailing_metadata#} p trailingMetadataArrPtr
        {#set grpc_op->data.recv_status_on_client.status#} p statusCodePtr
        {#set grpc_op->data.recv_status_on_client.status_details#} p ptrPtrStatusStr
        {#set grpc_op->data.recv_status_on_client.status_details_capacity#} p capacityPtr
      ]
    finish = do
      trailingMd <- readMetadataArray trailingMetadataArrPtr
      statusCode <- fmap toStatusCode (C.peek statusCodePtr)
      statusDetails <- B.packCString =<< C.peek ptrPtrStatusStr
      let status = RpcStatus trailingMd statusCode statusDetails
      writeIORef value status
      putMVar statusFromServer status
      free
    free = do
      freeMetadataArray trailingMetadataArrPtr
      C.free statusCodePtr
      C.free ptrPtrStatusStr
      C.free capacityPtr
  return (Op add value finish)

opSendCloseFromClient :: IO (OpT ())
opSendCloseFromClient = do
  value <- newIORef ()
  let
    add =
      [ \p -> do
        {#set grpc_op->op#} p (fromIntegral (fromEnum OpSendCloseFromClient))
      ]
    finish =
      return ()
  return (Op add value finish)

clientWaitForInitialMetadata :: ClientReaderWriter -> IO (RpcReply [Metadata])
clientWaitForInitialMetadata crw@(ClientReaderWriter { .. }) = do
  initMD <- readIORef initialMDRef
  case initMD of
    Just md -> return (RpcOk md)
    Nothing -> do
      recvInitialMetadataOp <- opRecvInitialMetadata crw
      res <- callBatch crw [ OpX recvInitialMetadataOp ]
      case res of
        RpcOk _ -> do
          md <- opRead recvInitialMetadataOp
          writeIORef initialMDRef (Just md)
          return (RpcOk md)
        RpcError err ->
          return (RpcError err)

clientReadInitialMetadata :: ClientReaderWriter -> IO (RpcReply (Maybe [Metadata]))
clientReadInitialMetadata (ClientReaderWriter {..}) = do
  md <- readIORef initialMDRef
  return (RpcOk md)

clientRead :: ClientReaderWriter -> IO (RpcReply (Maybe L.ByteString))
clientRead crw = do
  _ <- clientWaitForInitialMetadata crw
  recvMessage crw

recvMessage :: ClientReaderWriter -> IO (RpcReply (Maybe L.ByteString))
recvMessage crw = do
  recvMessageOp <- opRecvMessage
  res <- callBatch crw [ OpX recvMessageOp ]
  case res of
    RpcOk _ -> do
      bs <- opRead recvMessageOp
      return (RpcOk bs)
    RpcError err -> do
      putStrLn "recvMessage: callBatch failed"
      return (RpcError err)

clientWaitForStatus :: ClientReaderWriter -> IO (RpcReply RpcStatus)
clientWaitForStatus ClientReaderWriter{..} = do
  status <- readMVar statusFromServer
  return (RpcOk status)

clientWrite :: ClientReaderWriter -> B.ByteString -> IO (RpcReply ())
clientWrite crw@(ClientReaderWriter{..}) arg = do
  sendMessageOp <- opSendMessage arg
  callBatch crw [ OpX sendMessageOp ]

clientSendHalfClose :: ClientReaderWriter -> IO (RpcReply ())
clientSendHalfClose crw@(ClientReaderWriter{..}) = do
  sendCloseOp <- opSendCloseFromClient
  callBatch crw [ OpX sendCloseOp ]

clientCloseCall :: ClientReaderWriter -> IO ()
clientCloseCall ClientReaderWriter{..} = do
  tag <- tryTakeMVar statusFromServerTag
  case tag of
    Nothing -> return ()
    Just eDesc ->
      CQ.releaseEvent (ccWorker context) eDesc
  modifyMVar_ callMVar_ $ \call -> do
    grpcCallDestroy call
    return (error "grpcCallDestroy called on this Call")

type Rpc req resp a = ReaderT (Client req resp) (ExceptT RpcError IO) a

askCrw :: Rpc req resp ClientReaderWriter
askCrw = asks clientCrw

askDecoder :: Rpc req resp (Decoder resp)
askDecoder = asks clientDecoder

askEncoder :: Rpc req resp (Encoder req)
askEncoder = asks clientEncoder

joinReply :: RpcReply a -> Rpc req resp a
joinReply (RpcOk a) = return a
joinReply (RpcError err) = lift (throwE err)

withNewClient :: RpcReply (Client req resp) -> Rpc req resp a -> IO (RpcReply a)
withNewClient r_client m = do
  case r_client of
    RpcOk client -> withClient client m
    RpcError err -> return (RpcError err)

withClient :: Client req resp -> Rpc req resp a -> IO (RpcReply a)
withClient client m = do
  let m' = do
        x <- m
        _ <- waitForStatus
        throwIfErrorStatus
        return x
  e <- runExceptT (runReaderT m' client)
  case e of
    Left err -> return (RpcError err)
    Right a -> return (RpcOk a)

clientRWOp :: (ClientReaderWriter -> IO a) -> Rpc req resp a
clientRWOp act = do
  crw <- askCrw
  liftIO (act crw)

joinClientRWOp :: (ClientReaderWriter -> IO (RpcReply a)) -> Rpc req resp a
joinClientRWOp act = do
  x <- clientRWOp act
  joinReply x

branchOnStatus :: Rpc req resp a
               -> Rpc req resp a
               -> (StatusCode -> B.ByteString -> Rpc req resp a)
               -> Rpc req resp a
branchOnStatus onProcessing onSuccess onFail = do
  status <- clientRWOp (tryReadMVar . statusFromServer)
  case status of
    Nothing -> onProcessing
    Just (RpcStatus _ code msg)
      | code == StatusOk -> onSuccess
      | otherwise -> onFail code msg

throwIfErrorStatus :: Rpc req resp ()
throwIfErrorStatus =
  branchOnStatus
    (return ())
    (return ())
    (\code msg -> lift (throwE (StatusError code msg)))

initialMetadata :: Rpc req resp [Metadata]
initialMetadata = do
  status <- clientRWOp (readIORef . initialMDRef)
  case status of
    Just md -> return md
    Nothing ->
      branchOnStatus
        (joinClientRWOp clientWaitForInitialMetadata)
        (joinClientRWOp clientWaitForInitialMetadata)
        (\code msg -> lift (throwE (StatusError code msg)))

waitForStatus :: Rpc req resp RpcStatus
waitForStatus = do
  status <- clientRWOp (tryReadMVar . statusFromServer)
  case status of
    Nothing -> joinClientRWOp clientWaitForStatus
    Just status' -> return status'

receiveMessage :: Rpc req resp (Maybe resp)
receiveMessage = do
  let
    onProcessing = do
      msg <- joinClientRWOp clientRead
      case msg of
        Nothing -> return Nothing
        Just x -> do
          decoder <- askDecoder
          liftM Just (joinReply =<< liftIO (decoder x))
    onSuccess = return Nothing
    onFail code msg = lift (throwE (StatusError code msg))
  branchOnStatus onProcessing onSuccess onFail

receiveAllMessages :: Rpc req resp [resp]
receiveAllMessages = do
  decoder <- askDecoder
  let go acc = do
        value <- joinClientRWOp clientRead
        case value of
          Just x -> do
            y <- joinReply =<< (liftIO (decoder x))
            go (y:acc)
          Nothing -> return (reverse acc)
  go []

sendMessage :: req -> Rpc req resp ()
sendMessage req = do
  throwIfErrorStatus
  encoder <- askEncoder
  bs <- joinReply =<< liftIO (encoder req)
  joinClientRWOp (\crw -> clientWrite crw bs)

sendHalfClose :: Rpc req resp ()
sendHalfClose = do
  throwIfErrorStatus
  joinClientRWOp clientSendHalfClose

closeCall :: Rpc req resp ()
closeCall = do
  _ <- waitForStatus
  clientRWOp clientCloseCall
  throwIfErrorStatus

data RpcReply a
  = RpcOk a
  | RpcError RpcError
  deriving Show

data RpcError
  = DeadlineExceeded
  | Cancelled
  | Error String
  | CallErrorStatus CallError
  | StatusError StatusCode B.ByteString
  deriving Show

-- Copyright (c) 2016, Google Inc.
-- All rights reserved.
--
-- Redistribution and use in source and binary forms, with or without
-- modification, are permitted provided that the following conditions are met:
--     * Redistributions of source code must retain the above copyright
--       notice, this list of conditions and the following disclaimer.
--     * Redistributions in binary form must reproduce the above copyright
--       notice, this list of conditions and the following disclaimer in the
--       documentation and/or other materials provided with the distribution.
--     * Neither the name of Google Inc. nor the
--       names of its contributors may be used to endorse or promote products
--       derived from this software without specific prior written permission.
--
-- THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" AND
-- ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED
-- WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
-- DISCLAIMED. IN NO EVENT SHALL Google Inc. BE LIABLE FOR ANY
-- DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES
-- (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES;
-- LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND
-- ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
-- (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS
-- SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
--
--------------------------------------------------------------------------------
{-# LANGUAGE ScopedTypeVariables, OverloadedStrings, ExistentialQuantification, RecordWildCards #-}
module Network.Grpc.Core.Call where

import           Control.Concurrent
import           Control.Exception
import           Control.Monad

import qualified Data.ByteString              as B
import qualified Data.ByteString.Lazy         as L
import           Data.Monoid ((<>), Last(..))
import           Data.IORef

import           Foreign.C.Types              as C
import qualified Foreign.Marshal.Alloc        as C
import qualified Foreign.Marshal.Utils        as C
import qualified Foreign.Ptr                  as C
import           Foreign.Ptr (Ptr)
import qualified Foreign.Storable             as C

-- transformers
import           Control.Monad.IO.Class
import           Control.Monad.Trans.Class
import           Control.Monad.Trans.Except
import           Control.Monad.Trans.Reader

import qualified Network.Grpc.CompletionQueue as CQ
import Network.Grpc.Lib.PropagationBits
{#import Network.Grpc.Lib.ByteBuffer#}
{#import Network.Grpc.Lib.Metadata#}
{#import Network.Grpc.Lib.TimeSpec#}
{#import Network.Grpc.Lib.Core#}

#include <grpc/grpc.h>
#include "hs_grpc.h"

{#context lib = "grpc" prefix = "grpc" #}

data Deadline
  = AbsoluteDeadline TimeSpec
  | RelativeDeadline Int --seconds

data ClientContext = ClientContext
  { ccChannel :: Channel
  , ccCQ :: CompletionQueue
  , ccWorker :: CQ.Worker
  }

data CallOptions = CallOptions
  { coDeadline :: Maybe Deadline
  , coParentContext :: Maybe () -- todo
  , coPropagationMask :: Maybe PropagationMask
  , coMetadata :: [Metadata]
  }

instance Monoid CallOptions where
  mempty = CallOptions Nothing Nothing Nothing []
  mappend (CallOptions a b c d) (CallOptions a' b' c' d') =
    CallOptions
      (getLast (Last a <> Last a'))
      (getLast (Last b <> Last b'))
      (getLast (Last c <> Last c'))
      (d <> d')

withAbsoluteDeadline :: TimeSpec -> CallOptions
withAbsoluteDeadline deadline =
  mempty { coDeadline = Just (AbsoluteDeadline deadline) }

withRelativeDeadlineSeconds :: Int -> CallOptions
withRelativeDeadlineSeconds seconds =
  mempty { coDeadline = Just (RelativeDeadline seconds) }

withParentContext :: () -> CallOptions
withParentContext ctx =
  mempty { coParentContext = Just ctx }

withParentContextPropagating :: () -> PropagationMask -> CallOptions
withParentContextPropagating ctx prop =
  mempty { coParentContext   = Just ctx
         , coPropagationMask = Just prop }

withMetadata :: [Metadata] -> CallOptions
withMetadata md =
  mempty { coMetadata = md }

resolveDeadline :: CallOptions -> IO TimeSpec
resolveDeadline co =
  case coDeadline co of
    Nothing -> return gprInfFuture
    Just (AbsoluteDeadline deadline) -> return deadline
    Just (RelativeDeadline seconds) ->
      secondsFromNow (fromIntegral seconds)

newClientContext :: Channel -> IO ClientContext
newClientContext chan = do
  cq <- completionQueueCreate reservedPtr
  cqt <- CQ.startCompletionQueueThread cq
  return (ClientContext chan cq cqt)

destroyClientContext :: ClientContext -> IO ()
destroyClientContext (ClientContext _ cq w) = do
  completionQueueShutdown cq
  CQ.waitWorkerTermination w

type MethodName = B.ByteString
type Arg = B.ByteString

-- | Run the IO function once, cache it in the MVar.
onceMVar :: MVar (Maybe a) -> IO a -> IO a
onceMVar mvar io = modifyMVar mvar $ \x ->
  case x of
    Just x' -> return (x, x')
    Nothing -> do
      x' <- io
      return (Just x', x')

reservedPtr :: Ptr ()
reservedPtr = C.nullPtr

data UnaryResult a = UnaryResult [Metadata] [Metadata] a deriving Show

callUnary :: ClientContext -> CallOptions -> MethodName -> Arg -> IO (RpcReply (UnaryResult L.ByteString))
callUnary ctx@(ClientContext chan cq _) co method arg = do
  deadline <- resolveDeadline co
  bracket (grpcChannelCreateCall (cChannel chan) C.nullPtr propagateDefaults cq method (cHost chan) deadline) grpcCallDestroy $ \call0 -> newMVar call0 >>= \mcall -> do
    crw <- newClientReaderWriter ctx mcall

    sendInitOp <- opSendInitialMetadata (coMetadata co)
    sendCloseOp <- opSendCloseFromClient
    sendMessageOp <- opSendMessage arg
    recvStatusOp <- opRecvStatusOnClient crw
    recvMessageOp <- opRecvMessage
    recvInitialMetadataOp <- opRecvInitialMetadata crw

    res <- callBatch crw [
          OpX sendInitOp
        , OpX sendCloseOp
        , OpX recvInitialMetadataOp
        , OpX recvMessageOp
        , OpX sendMessageOp
        , OpX recvStatusOp
        ]
    case res of
      RpcOk _ -> do
        (RpcStatus trailMD status statusDetails) <- opRead recvStatusOp
        case status of
          StatusOk -> do
            initMD <- opRead recvInitialMetadataOp
            answ <- opRead recvMessageOp
            let answ' = maybe L.empty id answ
            return (RpcOk (UnaryResult initMD trailMD answ'))
          _ -> return (RpcError (StatusError status statusDetails))
      RpcError err -> return (RpcError err)

callDownstream :: ClientContext -> CallOptions -> MethodName -> Arg -> IO (RpcReply (Client B.ByteString L.ByteString))
callDownstream ctx@(ClientContext chan cq _) co method arg = do
  deadline <- resolveDeadline co
  mcall <- grpcChannelCreateCall (cChannel chan) C.nullPtr propagateDefaults cq method (cHost chan) deadline >>= newMVar
  crw <- newClientReaderWriter ctx mcall

  sendInitOp <- opSendInitialMetadata (coMetadata co)
  sendCloseOp <- opSendCloseFromClient
  sendMessageOp <- opSendMessage arg

  res <- callBatch crw [
        OpX sendInitOp
      , OpX sendCloseOp
      , OpX sendMessageOp
      ]
  case res of
    RpcOk _ -> do
      res' <- callBatchStatusOnClient crw
      case res' of
        RpcOk _ -> return (RpcOk (Client crw defaultEncoder defaultDecoder))
        RpcError err -> return (RpcError err)
    RpcError err -> return (RpcError err)

callUpstream :: ClientContext -> CallOptions -> MethodName -> IO (RpcReply (Client B.ByteString L.ByteString))
callUpstream ctx@(ClientContext chan cq _) co method = do
  deadline <- resolveDeadline co
  mcall <- grpcChannelCreateCall (cChannel chan) C.nullPtr propagateDefaults cq method (cHost chan) deadline >>= newMVar
  crw <- newClientReaderWriter ctx mcall

  sendInitOp <- opSendInitialMetadata (coMetadata co)

  res <- callBatch crw [ OpX sendInitOp ]
  case res of
    RpcOk _ -> do
      res' <- callBatchStatusOnClient crw
      case res' of
        RpcOk _ -> return (RpcOk (Client crw defaultEncoder defaultDecoder))
        RpcError err -> return (RpcError err)
    RpcError err -> return (RpcError err)

callBidi :: ClientContext -> CallOptions -> MethodName -> IO (RpcReply (Client B.ByteString L.ByteString))
callBidi ctx@(ClientContext chan cq _) co method = do
  deadline <- resolveDeadline co
  mcall <- grpcChannelCreateCall (cChannel chan) C.nullPtr propagateDefaults cq method (cHost chan) deadline >>= newMVar

  crw <- newClientReaderWriter ctx mcall
  sendInitOp <- opSendInitialMetadata (coMetadata co)

  res <- callBatch crw [ OpX sendInitOp ]
  case res of
    RpcOk _ -> do
      res' <- callBatchStatusOnClient crw
      case res' of
        RpcOk _ -> return (RpcOk (Client crw defaultEncoder defaultDecoder))
        RpcError err -> return (RpcError err)
    RpcError err -> return (RpcError err)

data Client req resp = Client
  { clientCrw :: ClientReaderWriter
  , clientEncoder :: Encoder req
  , clientDecoder :: Decoder resp
  }

type Decoder a = L.ByteString -> IO (RpcReply a)
type Encoder a = a -> IO (RpcReply B.ByteString)

defaultDecoder :: L.ByteString -> IO (RpcReply L.ByteString)
defaultDecoder bs = return (RpcOk bs)

defaultEncoder :: B.ByteString -> IO (RpcReply B.ByteString)
defaultEncoder bs = return (RpcOk bs)

data ClientReaderWriter = ClientReaderWriter {
  context :: ClientContext,
  callMVar_ :: MVar Call,
  initialMDRef :: !(IORef (Maybe [Metadata])),
  trailingMDRef :: !(IORef (Maybe [Metadata])),
  statusFromServer :: !(MVar RpcStatus),
  statusFromServerTag :: !(MVar CQ.EventDesc)
}

newClientReaderWriter :: ClientContext -> MVar Call -> IO ClientReaderWriter
newClientReaderWriter ctx mcall = do
  initMD <- newIORef Nothing
  trailMD <- newIORef Nothing
  status <- newEmptyMVar
  statusEvent <- newEmptyMVar
  return (ClientReaderWriter ctx mcall initMD trailMD status statusEvent)

data OpX = forall t. OpX (OpT t)

data OpArray = OpArray { opArrPtr           :: !GrpcOpPtr
                       , opArrLen           :: !CULong
                       , opArrFinishAndFree :: !(IO ())
                       }

toArray :: [OpX] -> IO OpArray
toArray ops = do
  aptr <- C.mallocBytes (length ops * {#sizeof grpc_op#})
  let ptrs = iterate (`C.plusPtr` {#sizeof grpc_op#}) aptr
      ops' = concatMap (\(OpX op) -> opAdd op) ops
      free = C.free aptr >> sequence_ (map (\(OpX op) -> opFinish op) ops)
      write op p = do
        {#set grpc_op->flags#} p 0
        {#set grpc_op->reserved#} p C.nullPtr
        op p
  zipWithM_ write ops' ptrs
  return $! OpArray aptr (fromIntegral $ length ops) free

callBatch :: ClientReaderWriter -> [OpX] -> IO (RpcReply ())
callBatch crw ops = do
  let ctx = context crw
  arr <- toArray ops
  let onBatchComplete = opArrFinishAndFree arr
  CQ.withEvent (ccWorker ctx) onBatchComplete $ \eDesc -> do
    callStatus <- withMVar (callMVar_ crw) $ \call ->
      grpcCallStartBatch call (opArrPtr arr) (opArrLen arr) (CQ.eventTag eDesc) reservedPtr
    case callStatus of
      CallOk -> do
        e <- CQ.interruptibleWaitEvent eDesc
        case e of
          QueueOpComplete OpSuccess _ -> return (RpcOk ())
          QueueOpComplete OpError _ -> return (RpcError (Error "callBatch: op error"))
          QueueTimeOut -> return (RpcError DeadlineExceeded)
          QueueShutdown -> return (RpcError (Error "queue shutdown"))
      _ -> do
         return (RpcError (CallErrorStatus callStatus))

callBatchStatusOnClient :: ClientReaderWriter -> IO (RpcReply ())
callBatchStatusOnClient crw@ClientReaderWriter{..} = do
  tag <- tryReadMVar statusFromServerTag
  case tag of
    Just _ -> return (RpcOk ()) -- already did this before
    Nothing -> do
      statusOp <- opRecvStatusOnClient crw
      arr <- toArray [ OpX statusOp ]
      let onBatchComplete = opArrFinishAndFree arr
      eDesc <- CQ.allocateEvent (ccWorker context) onBatchComplete
      putMVar statusFromServerTag eDesc
      callStatus <- withMVar callMVar_ $ \call ->
        grpcCallStartBatch call (opArrPtr arr) (opArrLen arr) (CQ.eventTag eDesc) reservedPtr
      case callStatus of
        CallOk -> return (RpcOk ())
        _ -> return (RpcError (CallErrorStatus callStatus))

data OpT out = Op
  { opAdd    :: [Ptr GrpcOp -> IO ()]
  , opValue  :: IORef out
  , opFinish :: IO ()
  }

opRead :: OpT a -> IO a
opRead op = readIORef (opValue op)

opRecvInitialMetadata :: ClientReaderWriter -> IO (OpT [Metadata])
opRecvInitialMetadata crw = do
  arr <- mallocMetadataArray
  value <- newIORef (error "opRecvInitialMetadata never finished")
  let
    add =
      [ \p -> do
        {#set grpc_op->op#} p (fromIntegral (fromEnum OpRecvInitialMetadata))
        {#set grpc_op->data.recv_initial_metadata#} p arr
      ]
    finish = do
      mds <- readMetadataArray arr
      writeIORef (initialMDRef crw) (Just mds)
      writeIORef value mds
      freeMetadataArray arr
  return (Op add value finish)

opSendMessage :: B.ByteString -> IO (OpT ())
opSendMessage bs = do
  bb <- fromByteString bs
  value <- newIORef ()
  let
    add =
      [ \p -> do
        {#set grpc_op->op#} p (fromIntegral (fromEnum OpSendMessage))
        {#set grpc_op->data.send_message#} p bb
      ]
    finish =
      byteBufferDestroy bb
  return (Op add value finish)

opRecvMessage :: IO (OpT (Maybe L.ByteString))
opRecvMessage = do
  bbptr <- C.malloc :: IO (Ptr (Ptr CByteBuffer))
  value <- newIORef Nothing
  let
    add =
      [ \p -> do
        {#set grpc_op->op#} p (fromIntegral (fromEnum OpRecvMessage))
        {#set grpc_op->data.recv_message#} p bbptr
      ]
    finish = do
      bb <- C.peek bbptr
      C.free bbptr
      writeIORef value =<< if bb /= C.nullPtr
        then do
          lbs <- toLazyByteString bb
          byteBufferDestroy bb
          return (Just lbs)
        else return Nothing
  return (Op add value finish)

opSendInitialMetadata :: [Metadata] -> IO (OpT ())
opSendInitialMetadata elems = do
  (mdArrPtr, free) <- mallocMetadata elems
  value <- newIORef ()
  let
    add =
      [ \p -> do
        {#set grpc_op->op#} p (fromIntegral (fromEnum OpSendInitialMetadata))
        {#set grpc_op->data.send_initial_metadata.count#} p (fromIntegral (length elems))
        {#set grpc_op->data.send_initial_metadata.metadata#} p (C.castPtr mdArrPtr)
        {#set grpc_op->data.send_initial_metadata.maybe_compression_level.is_set#} p 0
        {#set grpc_op->data.send_initial_metadata.maybe_compression_level.level#} p 0
      ]
    finish = do
      free
  return (Op add value finish)

data RpcStatus = RpcStatus [Metadata] StatusCode B.ByteString
  deriving Show

opRecvStatusOnClient :: ClientReaderWriter -> IO (OpT RpcStatus)
opRecvStatusOnClient (ClientReaderWriter{..}) = do
  trailingMetadataArrPtr <- mallocMetadataArray
  statusCodePtr <- C.malloc :: IO (Ptr StatusCodeT)
  ptrPtrStatusStr <- C.new C.nullPtr :: IO (Ptr (Ptr CChar))
  capacityPtr <- C.new 0 :: IO (Ptr SizeT)
  value <- newIORef (error "opRecvStatusOnClient never ran")
  let
    add =
      [ \p -> do
        {#set grpc_op->op#} p (fromIntegral (fromEnum OpRecvStatusOnClient))
        {#set grpc_op->data.recv_status_on_client.trailing_metadata#} p trailingMetadataArrPtr
        {#set grpc_op->data.recv_status_on_client.status#} p statusCodePtr
        {#set grpc_op->data.recv_status_on_client.status_details#} p ptrPtrStatusStr
        {#set grpc_op->data.recv_status_on_client.status_details_capacity#} p capacityPtr
      ]
    finish = do
      trailingMd <- readMetadataArray trailingMetadataArrPtr
      statusCode <- fmap toStatusCode (C.peek statusCodePtr)
      statusDetails <- B.packCString =<< C.peek ptrPtrStatusStr
      let status = RpcStatus trailingMd statusCode statusDetails
      writeIORef value status
      putMVar statusFromServer status
      free
    free = do
      freeMetadataArray trailingMetadataArrPtr
      C.free statusCodePtr
      C.free ptrPtrStatusStr
      C.free capacityPtr
  return (Op add value finish)

opSendCloseFromClient :: IO (OpT ())
opSendCloseFromClient = do
  value <- newIORef ()
  let
    add =
      [ \p -> do
        {#set grpc_op->op#} p (fromIntegral (fromEnum OpSendCloseFromClient))
      ]
    finish =
      return ()
  return (Op add value finish)

clientWaitForInitialMetadata :: ClientReaderWriter -> IO (RpcReply [Metadata])
clientWaitForInitialMetadata crw@(ClientReaderWriter { .. }) = do
  initMD <- readIORef initialMDRef
  case initMD of
    Just md -> return (RpcOk md)
    Nothing -> do
      recvInitialMetadataOp <- opRecvInitialMetadata crw
      res <- callBatch crw [ OpX recvInitialMetadataOp ]
      case res of
        RpcOk _ -> do
          md <- opRead recvInitialMetadataOp
          writeIORef initialMDRef (Just md)
          return (RpcOk md)
        RpcError err ->
          return (RpcError err)

clientReadInitialMetadata :: ClientReaderWriter -> IO (RpcReply (Maybe [Metadata]))
clientReadInitialMetadata (ClientReaderWriter {..}) = do
  md <- readIORef initialMDRef
  return (RpcOk md)

clientRead :: ClientReaderWriter -> IO (RpcReply (Maybe L.ByteString))
clientRead crw = do
  _ <- clientWaitForInitialMetadata crw
  recvMessage crw

recvMessage :: ClientReaderWriter -> IO (RpcReply (Maybe L.ByteString))
recvMessage crw = do
  recvMessageOp <- opRecvMessage
  res <- callBatch crw [ OpX recvMessageOp ]
  case res of
    RpcOk _ -> do
      bs <- opRead recvMessageOp
      return (RpcOk bs)
    RpcError err -> do
      putStrLn "recvMessage: callBatch failed"
      return (RpcError err)

clientWaitForStatus :: ClientReaderWriter -> IO (RpcReply RpcStatus)
clientWaitForStatus ClientReaderWriter{..} = do
  status <- readMVar statusFromServer
  return (RpcOk status)

clientClose :: ClientReaderWriter -> IO ()
clientClose (ClientReaderWriter{..}) = do
  withMVar callMVar_ $ \call ->
    grpcCallDestroy call

clientWrite :: ClientReaderWriter -> B.ByteString -> IO (RpcReply ())
clientWrite crw@(ClientReaderWriter{..}) arg = do
  sendMessageOp <- opSendMessage arg
  callBatch crw [ OpX sendMessageOp ]

clientSendHalfClose :: ClientReaderWriter -> IO (RpcReply ())
clientSendHalfClose crw@(ClientReaderWriter{..}) = do
  sendCloseOp <- opSendCloseFromClient
  callBatch crw [ OpX sendCloseOp ]

clientCloseCall :: ClientReaderWriter -> IO ()
clientCloseCall ClientReaderWriter{..} = do
  tag <- tryTakeMVar statusFromServerTag
  case tag of
    Nothing -> return ()
    Just eDesc ->
      CQ.releaseEvent (ccWorker context) eDesc
  withMVar callMVar_ $ \call ->
    grpcCallDestroy call

type Rpc req resp a = ReaderT (Client req resp) (ExceptT RpcError IO) a

askCrw :: Rpc req resp ClientReaderWriter
askCrw = asks clientCrw

askDecoder :: Rpc req resp (Decoder resp)
askDecoder = asks clientDecoder

askEncoder :: Rpc req resp (Encoder req)
askEncoder = asks clientEncoder

joinReply :: RpcReply a -> Rpc req resp a
joinReply (RpcOk a) = return a
joinReply (RpcError err) = lift (throwE err)

withNewClient :: RpcReply (Client req resp) -> Rpc req resp a -> IO (RpcReply a)
withNewClient r_client m = do
  case r_client of
    RpcOk client -> withClient client m
    RpcError err -> return (RpcError err)

withClient :: Client req resp -> Rpc req resp a -> IO (RpcReply a)
withClient client m = do
  e <- runExceptT (runReaderT m client)
  case e of
    Left err -> return (RpcError err)
    Right a -> return (RpcOk a)

clientRWOp :: (ClientReaderWriter -> IO a) -> Rpc req resp a
clientRWOp act = do
  crw <- askCrw
  liftIO (act crw)

joinClientRWOp :: (ClientReaderWriter -> IO (RpcReply a)) -> Rpc req resp a
joinClientRWOp act = do
  x <- clientRWOp act
  joinReply x

branchOnStatus :: Rpc req resp a
               -> Rpc req resp a
               -> (StatusCode -> B.ByteString -> Rpc req resp a)
               -> Rpc req resp a
branchOnStatus onProcessing onSuccess onFail = do
  status <- clientRWOp (tryReadMVar . statusFromServer)
  case status of
    Nothing -> onProcessing
    Just (RpcStatus _ code msg)
      | code == StatusOk -> onSuccess
      | otherwise -> onFail code msg

abortIfStatus :: Rpc req resp ()
abortIfStatus =
  branchOnStatus
    (return ())
    (return ())
    (\code msg -> lift (throwE (StatusError code msg)))

initialMetadata :: Rpc req resp [Metadata]
initialMetadata = do
  status <- clientRWOp (readIORef . initialMDRef)
  case status of
    Just md -> return md
    Nothing ->
      branchOnStatus
        (joinClientRWOp clientWaitForInitialMetadata)
        (joinClientRWOp clientWaitForInitialMetadata)
        (\code msg -> lift (throwE (StatusError code msg)))

waitForStatus :: Rpc req resp RpcStatus
waitForStatus = do
  status <- clientRWOp (tryReadMVar . statusFromServer)
  case status of
    Nothing -> joinClientRWOp clientWaitForStatus
    Just status' -> return status'

receiveMessage :: Rpc req resp (Maybe resp)
receiveMessage = do
  let
    onProcessing = do
      msg <- joinClientRWOp clientRead
      case msg of
        Nothing -> return Nothing
        Just x -> do
          decoder <- askDecoder
          liftM Just (joinReply =<< liftIO (decoder x))
    onSuccess = return Nothing
    onFail code msg = lift (throwE (StatusError code msg))
  branchOnStatus onProcessing onSuccess onFail

receiveAllMessages :: Rpc req resp [resp]
receiveAllMessages = do
  decoder <- askDecoder
  let go acc = do
        value <- joinClientRWOp clientRead
        case value of
          Just x -> do
            y <- joinReply =<< (liftIO (decoder x))
            go (y:acc)
          Nothing -> return (reverse acc)
  go []

sendMessage :: req -> Rpc req resp ()
sendMessage req = do
  abortIfStatus
  encoder <- askEncoder
  bs <- joinReply =<< liftIO (encoder req)
  joinClientRWOp (\crw -> clientWrite crw bs)

sendHalfClose :: Rpc req resp ()
sendHalfClose = do
  abortIfStatus
  joinClientRWOp clientSendHalfClose

closeCall :: Rpc req resp ()
closeCall =
  clientRWOp clientClose

data RpcReply a
  = RpcOk a
  | RpcError RpcError
  deriving Show

data RpcError
  = DeadlineExceeded
  | Cancelled
  | Error String
  | CallErrorStatus CallError
  | StatusError StatusCode B.ByteString
  deriving Show


Adjust webResourceGetMimeType.

Adjust webResourceGetMimeType.

{-# LANGUAGE CPP #-}
-- -*-haskell-*-
-----------------------------------------------------------------------------
--  Module      :  Graphics.UI.Gtk.WebKit.WebResource
--  Author      :  Andy Stewart
--  Copyright   :  (c) 2010 Andy Stewart <lazycat.manatee@gmail.com>
--
--  This library is free software; you can redistribute it and/or
--  modify it under the terms of the GNU Lesser General Public
--  License as published by the Free Software Foundation; either
--  version 2.1 of the License, or (at your option) any later version.
--
--  This library is distributed in the hope that it will be useful,
--  but WITHOUT ANY WARRANTY; without even the implied warranty of
--  MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU
--  Lesser General Public License for more details.
--
-- |
-- Maintainer  : gtk2hs-users@lists.sourceforge.net
-- Stability   : provisional
-- Portability : portable (depends on GHC)
--
-- Note:
--
-- Function `webkit_web_resource_get_data` haven't binding
-- no idea how to handle `GString`.
--
-- Access to the WebKit Web Resource
-----------------------------------------------------------------------------

module Graphics.UI.Gtk.WebKit.WebResource (
-- * Description
-- | A web resource encapsulates the data of the download as well as the URI, MIME type and frame name of
-- the resource.

-- * Types
  WebResource,
  WebResourceClass,

-- * Constructors
  webResourceNew,

-- * Methods
  webResourceGetEncoding,
  webResourceGetFrameName,
  webResourceGetMimeType,
  webResourceGetUri,
) where

import Control.Monad		(liftM)

import System.Glib.FFI
import System.Glib.UTFString
import System.Glib.GList
import System.Glib.GError
import Graphics.UI.Gtk.Gdk.Events

{#import Graphics.UI.Gtk.Abstract.Object#}	(makeNewObject)
{#import Graphics.UI.Gtk.WebKit.Types#}
{#import System.Glib.GObject#}

{#context lib="webkit" prefix ="webkit"#}

-- | Returns a new WebKitWebResource.
-- The @encoding@ can be empty.
-- The @frameName@ can be used if the resource represents contents of an
-- entire HTML frame, otherwise pass empty.
webResourceNew :: String -> Int -> String -> String -> String -> String -> IO WebResource
webResourceNew resData size uri mimeType encoding frameName =
   withCString resData $ \dataPtr ->
   withCString uri $ \uriPtr ->
   withCString mimeType $ \mimePtr ->
   withCString encoding $ \encodingPtr ->
   withCString frameName $ \framePtr ->
   wrapNewGObject mkWebResource $
     {#call web_resource_new#} dataPtr (fromIntegral size) uriPtr mimePtr encodingPtr framePtr

-- | Get encoding.
webResourceGetEncoding ::
   WebResourceClass self => self
 -> IO (Maybe String)
webResourceGetEncoding wr =
  {#call web_resource_get_encoding#} (toWebResource wr) >>= maybePeek peekCString

-- | Get frame name.
webResourceGetFrameName ::
   WebResourceClass self => self
 -> IO (Maybe String)
webResourceGetFrameName wr =
  {#call web_resource_get_frame_name#} (toWebResource wr) >>= maybePeek peekCString

-- | Get mime type.
webResourceGetMimeType ::
   WebResourceClass self => self
 -> IO (Maybe String)
webResourceGetMimeType wr =
  {#call web_resource_get_mime_type#} (toWebResource wr) >>= maybePeek peekCString

-- | Get uri.
webResourceGetUri ::
   WebResourceClass self => self
 -> IO String
webResourceGetUri wr =
  {#call web_resource_get_uri#} (toWebResource wr) >>= peekCString

{-# LANGUAGE CPP #-}
-- -*-haskell-*-
-----------------------------------------------------------------------------
--  Module      :  Graphics.UI.Gtk.WebKit.WebResource
--  Author      :  Andy Stewart
--  Copyright   :  (c) 2010 Andy Stewart <lazycat.manatee@gmail.com>
--
--  This library is free software; you can redistribute it and/or
--  modify it under the terms of the GNU Lesser General Public
--  License as published by the Free Software Foundation; either
--  version 2.1 of the License, or (at your option) any later version.
--
--  This library is distributed in the hope that it will be useful,
--  but WITHOUT ANY WARRANTY; without even the implied warranty of
--  MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU
--  Lesser General Public License for more details.
--
-- |
-- Maintainer  : gtk2hs-users@lists.sourceforge.net
-- Stability   : provisional
-- Portability : portable (depends on GHC)
--
-- Note:
--
-- Function `webkit_web_resource_get_data` haven't binding
-- no idea how to handle `GString`.
--
-- Access to the WebKit Web Resource
-----------------------------------------------------------------------------

module Graphics.UI.Gtk.WebKit.WebResource (
-- * Description
-- | A web resource encapsulates the data of the download as well as the URI, MIME type and frame name of
-- the resource.

-- * Types
  WebResource,
  WebResourceClass,

-- * Constructors
  webResourceNew,

-- * Methods
  webResourceGetEncoding,
  webResourceGetFrameName,
  webResourceGetMimeType,
  webResourceGetUri,
) where

import Control.Monad		(liftM)

import System.Glib.FFI
import System.Glib.UTFString
import System.Glib.GList
import System.Glib.GError
import Graphics.UI.Gtk.Gdk.Events

{#import Graphics.UI.Gtk.Abstract.Object#}	(makeNewObject)
{#import Graphics.UI.Gtk.WebKit.Types#}
{#import System.Glib.GObject#}

{#context lib="webkit" prefix ="webkit"#}

-- | Returns a new WebKitWebResource.
-- The @encoding@ can be empty.
-- The @frameName@ can be used if the resource represents contents of an
-- entire HTML frame, otherwise pass empty.
webResourceNew :: String -> Int -> String -> String -> String -> String -> IO WebResource
webResourceNew resData size uri mimeType encoding frameName =
   withCString resData $ \dataPtr ->
   withCString uri $ \uriPtr ->
   withCString mimeType $ \mimePtr ->
   withCString encoding $ \encodingPtr ->
   withCString frameName $ \framePtr ->
   wrapNewGObject mkWebResource $
     {#call web_resource_new#} dataPtr (fromIntegral size) uriPtr mimePtr encodingPtr framePtr

-- | Get encoding.
webResourceGetEncoding ::
   WebResourceClass self => self
 -> IO (Maybe String)
webResourceGetEncoding wr =
  {#call web_resource_get_encoding#} (toWebResource wr) >>= maybePeek peekCString

-- | Get frame name.
webResourceGetFrameName ::
   WebResourceClass self => self
 -> IO (Maybe String)
webResourceGetFrameName wr =
  {#call web_resource_get_frame_name#} (toWebResource wr) >>= maybePeek peekCString

-- | Get mime type.
webResourceGetMimeType ::
   WebResourceClass self => self
 -> IO String
webResourceGetMimeType wr =
  {#call web_resource_get_mime_type#} (toWebResource wr) >>= peekCString

-- | Get uri.
webResourceGetUri ::
   WebResourceClass self => self
 -> IO String
webResourceGetUri wr =
  {#call web_resource_get_uri#} (toWebResource wr) >>= peekCString


more descriptive error messages

more descriptive error messages

Ignore-this: 3cd2d807debfaf06f11ddb1b06860930

darcs-hash:20090721232441-115f9-1ad3cfccdfb76cae10c60475eb133e1bcc7f9370.gz

{-# LANGUAGE ForeignFunctionInterface #-}
--------------------------------------------------------------------------------
-- |
-- Module    : Foreign.CUDA.Marshal
-- Copyright : (c) 2009 Trevor L. McDonell
-- License   : BSD
--
-- Memory allocation and marshalling support for CUDA devices
--
--------------------------------------------------------------------------------

module Foreign.CUDA.Marshal
  (
    DevicePtr,

    -- ** Dynamic allocation
    malloc,
    malloc2D,
    malloc3D,
    memset,
    memset2D,
    memset3D,

    -- ** Marshalling

    -- ** Combined allocation and marshalling

    -- ** Copying
  )
  where

import Foreign.CUDA.Error
import Foreign.CUDA.Stream
import Foreign.CUDA.Internal.C2HS hiding (malloc)

#include <cuda_runtime_api.h>
{# context lib="cudart" #}

--------------------------------------------------------------------------------
-- Data Types
--------------------------------------------------------------------------------

-- |
-- A reference to data stored on the device. It is automatically freed.
--
type DevicePtr a = ForeignPtr a

withDevicePtr :: DevicePtr a -> (Ptr a -> IO b) -> IO b
withDevicePtr =  withForeignPtr

newDevicePtr   :: Ptr () -> IO (DevicePtr ())
newDevicePtr p =  (doAutoRelease free_) >>= \f -> newForeignPtr f p

--
-- Memory copy
--
{# enum cudaMemcpyKind as CopyDirection {}
    with prefix="cudaMemcpy" deriving (Eq, Show) #}

--------------------------------------------------------------------------------
-- Dynamic allocation
--------------------------------------------------------------------------------

-- |
-- Allocate the specified number of bytes in linear memory on the device. The
-- memory is suitably aligned for any kind of variable, and is not cleared.
--
malloc       :: Int64 -> IO (Either String (DevicePtr ()))
malloc bytes = do
    (rv,ptr) <- cudaMalloc bytes
    case rv of
        Success -> Right `fmap` newDevicePtr ptr
        _       -> return . Left . describe $ rv

{# fun unsafe cudaMalloc
    { alloca-  `Ptr ()' peek* ,
      cIntConv `Int64'        } -> `Status' cToEnum #}

-- |
-- Allocate at least @width * height@ bytes of linear memory on the device. The
-- function may pad the allocation to ensure corresponding pointers in each row
-- meet coalescing requirements. The actual allocation width is returned.
--
malloc2D :: (Int64, Int64)              -- ^ allocation (width,height) in bytes
         -> IO (Either String (DevicePtr (),Int64))
malloc2D (width,height) =  do
    (rv,ptr,pitch) <- cudaMallocPitch width height
    case rv of
        Success -> (\p -> Right (p,pitch)) `fmap` newDevicePtr ptr
        _       -> return . Left . describe $ rv

{# fun unsafe cudaMallocPitch
    { alloca-  `Ptr ()' peek*        ,
      alloca-  `Int64'  peekIntConv* ,
      cIntConv `Int64'               ,
      cIntConv `Int64'               } -> `Status' cToEnum #}

-- |
-- Allocate at least @width * height * depth@ bytes of linear memory on the
-- device. The function may pad the allocation to ensure hardware alignment
-- requirements are met. The actual allocation pitch is returned
--
malloc3D :: (Int64,Int64,Int64)         -- ^ allocation (width,height,depth) in bytes
         -> IO (Either String (DevicePtr (),Int64))
malloc3D = moduleErr "malloc3D" "not implemented yet"

-- |
-- Free previously allocated memory on the device
--
-- free   :: Ptr () -> IO (Maybe String)
-- free p =  nothingIfOk `fmap` cudaFree p

free_   :: Ptr () -> IO ()
free_ p =  throwIf_ (/= Success) (describe) (cudaFree p)

{# fun unsafe cudaFree
    { id `Ptr ()' } -> `Status' cToEnum #}

foreign import ccall "wrapper"
    doAutoRelease   :: (Ptr () -> IO ()) -> IO (FunPtr (Ptr () -> IO ()))

-- |
-- Initialise device memory to a given value
--
memset                  :: DevicePtr () -> Int64 -> Int -> IO (Maybe String)
memset ptr bytes symbol =  nothingIfOk `fmap` cudaMemset ptr symbol bytes

{# fun unsafe cudaMemset
    { withDevicePtr* `DevicePtr ()' ,
                     `Int'          ,
      cIntConv       `Int64'        } -> `Status' cToEnum #}

-- |
-- Initialise a matrix to a given value
--
memset2D :: DevicePtr ()                -- ^ The device memory
         -> (Int64,Int64)               -- ^ The (width,height) of the matrix in bytes
         -> Int64                       -- ^ The allocation pitch, as returned by 'malloc2D'
         -> Int                         -- ^ Value to set for each byte
         -> IO (Maybe String)
memset2D ptr (width,height) pitch symbol = nothingIfOk `fmap` cudaMemset2D ptr pitch symbol width height

{# fun unsafe cudaMemset2D
    { withDevicePtr* `DevicePtr ()' ,
      cIntConv       `Int64'        ,
                     `Int'          ,
      cIntConv       `Int64'        ,
      cIntConv       `Int64'        } -> `Status' cToEnum #}

-- |
-- Initialise the elements of a 3D array to a given value
--
memset3D :: DevicePtr ()                -- ^ The device memory
         -> (Int64,Int64,Int64)         -- ^ The (width,height,depth) of the array in bytes
         -> Int64                       -- ^ The allocation pitch, as returned by 'malloc3D'
         -> Int                         -- ^ Value to set for each byte
         -> IO (Maybe String)
memset3D = moduleErr "memset3D" "not implemented yet"

--------------------------------------------------------------------------------
-- Marshalling
--------------------------------------------------------------------------------

-- peek
-- poke

--------------------------------------------------------------------------------
-- Combined allocation and marshalling
--------------------------------------------------------------------------------

-- newArray
-- withArray

--------------------------------------------------------------------------------
-- Copying
--------------------------------------------------------------------------------

-- |
-- Copy data between host and device
--
memcpy :: Ptr a                 -- ^ destination
       -> Ptr a                 -- ^ source
       -> Int64                 -- ^ number of bytes
       -> CopyDirection
       -> IO (Maybe String)
memcpy dst src bytes dir =  nothingIfOk `fmap` cudaMemcpy dst src bytes dir

{# fun unsafe cudaMemcpy
    { castPtr   `Ptr a'         ,
      castPtr   `Ptr a'         ,
      cIntConv  `Int64'         ,
      cFromEnum `CopyDirection' } -> `Status' cToEnum #}

-- |
-- Copy data between host and device asynchronously
--
memcpyAsync :: Ptr a            -- ^ destination
            -> Ptr a            -- ^ source
            -> Int64            -- ^ number of bytes
            -> CopyDirection
            -> Stream
            -> IO (Maybe String)
memcpyAsync dst src bytes dir stream =  nothingIfOk `fmap` cudaMemcpyAsync dst src bytes dir stream

{# fun unsafe cudaMemcpyAsync
    { castPtr    `Ptr a'         ,
      castPtr    `Ptr a'         ,
      cIntConv   `Int64'         ,
      cFromEnum  `CopyDirection' ,
      cIntConv   `Stream'        } -> `Status' cToEnum #}

--------------------------------------------------------------------------------
-- Internal utilities
--------------------------------------------------------------------------------

moduleErr :: String -> String -> a
moduleErr fun msg =  error ("Foreign.CUDA." ++ fun ++ ':':' ':msg)

{-# LANGUAGE ForeignFunctionInterface #-}
--------------------------------------------------------------------------------
-- |
-- Module    : Foreign.CUDA.Marshal
-- Copyright : (c) 2009 Trevor L. McDonell
-- License   : BSD
--
-- Memory allocation and marshalling support for CUDA devices
--
--------------------------------------------------------------------------------

module Foreign.CUDA.Marshal
  (
    DevicePtr,

    -- ** Dynamic allocation
    malloc,
    malloc2D,
    malloc3D,
    memset,
    memset2D,
    memset3D,

    -- ** Marshalling

    -- ** Combined allocation and marshalling

    -- ** Copying
  )
  where

import Foreign.CUDA.Error
import Foreign.CUDA.Stream
import Foreign.CUDA.Internal.C2HS hiding (malloc)

#include <cuda_runtime_api.h>
{# context lib="cudart" #}

--------------------------------------------------------------------------------
-- Data Types
--------------------------------------------------------------------------------

-- |
-- A reference to data stored on the device. It is automatically freed.
--
type DevicePtr a = ForeignPtr a

withDevicePtr :: DevicePtr a -> (Ptr a -> IO b) -> IO b
withDevicePtr =  withForeignPtr

newDevicePtr   :: Ptr () -> IO (DevicePtr ())
newDevicePtr p =  (doAutoRelease free_) >>= \f -> newForeignPtr f p

--
-- Memory copy
--
{# enum cudaMemcpyKind as CopyDirection {}
    with prefix="cudaMemcpy" deriving (Eq, Show) #}

--------------------------------------------------------------------------------
-- Dynamic allocation
--------------------------------------------------------------------------------

-- |
-- Allocate the specified number of bytes in linear memory on the device. The
-- memory is suitably aligned for any kind of variable, and is not cleared.
--
malloc       :: Int64 -> IO (Either String (DevicePtr ()))
malloc bytes = do
    (rv,ptr) <- cudaMalloc bytes
    case rv of
        Success -> Right `fmap` newDevicePtr ptr
        _       -> return . Left . describe $ rv

{# fun unsafe cudaMalloc
    { alloca-  `Ptr ()' peek* ,
      cIntConv `Int64'        } -> `Status' cToEnum #}

-- |
-- Allocate at least @width * height@ bytes of linear memory on the device. The
-- function may pad the allocation to ensure corresponding pointers in each row
-- meet coalescing requirements. The actual allocation width is returned.
--
malloc2D :: (Int64, Int64)              -- ^ allocation (width,height) in bytes
         -> IO (Either String (DevicePtr (),Int64))
malloc2D (width,height) =  do
    (rv,ptr,pitch) <- cudaMallocPitch width height
    case rv of
        Success -> (\p -> Right (p,pitch)) `fmap` newDevicePtr ptr
        _       -> return . Left . describe $ rv

{# fun unsafe cudaMallocPitch
    { alloca-  `Ptr ()' peek*        ,
      alloca-  `Int64'  peekIntConv* ,
      cIntConv `Int64'               ,
      cIntConv `Int64'               } -> `Status' cToEnum #}

-- |
-- Allocate at least @width * height * depth@ bytes of linear memory on the
-- device. The function may pad the allocation to ensure hardware alignment
-- requirements are met. The actual allocation pitch is returned
--
malloc3D :: (Int64,Int64,Int64)         -- ^ allocation (width,height,depth) in bytes
         -> IO (Either String (DevicePtr (),Int64))
malloc3D = error "not implemented yet"

-- |
-- Free previously allocated memory on the device
--
-- free   :: Ptr () -> IO (Maybe String)
-- free p =  nothingIfOk `fmap` cudaFree p

free_   :: Ptr () -> IO ()
free_ p =  throwIf_ (/= Success) (describe) (cudaFree p)

{# fun unsafe cudaFree
    { id `Ptr ()' } -> `Status' cToEnum #}

foreign import ccall "wrapper"
    doAutoRelease   :: (Ptr () -> IO ()) -> IO (FunPtr (Ptr () -> IO ()))

-- |
-- Initialise device memory to a given value
--
memset                  :: DevicePtr () -> Int64 -> Int -> IO (Maybe String)
memset ptr bytes symbol =  nothingIfOk `fmap` cudaMemset ptr symbol bytes

{# fun unsafe cudaMemset
    { withDevicePtr* `DevicePtr ()' ,
                     `Int'          ,
      cIntConv       `Int64'        } -> `Status' cToEnum #}

-- |
-- Initialise a matrix to a given value
--
memset2D :: DevicePtr ()                -- ^ The device memory
         -> (Int64,Int64)               -- ^ The (width,height) of the matrix in bytes
         -> Int64                       -- ^ The allocation pitch, as returned by 'malloc2D'
         -> Int                         -- ^ Value to set for each byte
         -> IO (Maybe String)
memset2D ptr (width,height) pitch symbol = nothingIfOk `fmap` cudaMemset2D ptr pitch symbol width height

{# fun unsafe cudaMemset2D
    { withDevicePtr* `DevicePtr ()' ,
      cIntConv       `Int64'        ,
                     `Int'          ,
      cIntConv       `Int64'        ,
      cIntConv       `Int64'        } -> `Status' cToEnum #}

-- |
-- Initialise the elements of a 3D array to a given value
--
memset3D :: DevicePtr ()                -- ^ The device memory
         -> (Int64,Int64,Int64)         -- ^ The (width,height,depth) of the array in bytes
         -> Int64                       -- ^ The allocation pitch, as returned by 'malloc3D'
         -> Int                         -- ^ Value to set for each byte
         -> IO (Maybe String)
memset3D = error "not implemented yet"

--------------------------------------------------------------------------------
-- Marshalling
--------------------------------------------------------------------------------

-- peek
-- poke

--------------------------------------------------------------------------------
-- Combined allocation and marshalling
--------------------------------------------------------------------------------

-- newArray
-- withArray

--------------------------------------------------------------------------------
-- Copying
--------------------------------------------------------------------------------

-- |
-- Copy data between host and device
--
memcpy :: Ptr a                 -- ^ destination
       -> Ptr a                 -- ^ source
       -> Int64                 -- ^ number of bytes
       -> CopyDirection
       -> IO (Maybe String)
memcpy dst src bytes dir =  nothingIfOk `fmap` cudaMemcpy dst src bytes dir

{# fun unsafe cudaMemcpy
    { castPtr   `Ptr a'         ,
      castPtr   `Ptr a'         ,
      cIntConv  `Int64'         ,
      cFromEnum `CopyDirection' } -> `Status' cToEnum #}

-- |
-- Copy data between host and device asynchronously
--
memcpyAsync :: Ptr a            -- ^ destination
            -> Ptr a            -- ^ source
            -> Int64            -- ^ number of bytes
            -> CopyDirection
            -> Stream
            -> IO (Maybe String)
memcpyAsync dst src bytes dir stream =  nothingIfOk `fmap` cudaMemcpyAsync dst src bytes dir stream

{# fun unsafe cudaMemcpyAsync
    { castPtr    `Ptr a'         ,
      castPtr    `Ptr a'         ,
      cIntConv   `Int64'         ,
      cFromEnum  `CopyDirection' ,
      cIntConv   `Stream'        } -> `Status' cToEnum #}


this utility from the driver side as well

Ignore-this: f11aba61078a41269d2ddd7e690a4fa1

darcs-hash:20091210040032-9241b-498cea3938b2bacaaf62670935020f8e4000a458.gz

{-# LANGUAGE ForeignFunctionInterface #-}
--------------------------------------------------------------------------------
-- |
-- Module    : Foreign.CUDA.Driver.Utils
-- Copyright : (c) 2009 Trevor L. McDonell
-- License   : BSD
--
-- Utility functions
--
--------------------------------------------------------------------------------

module Foreign.CUDA.Driver.Utils
  (
    forceEither, driverVersion
  )
  where

#include <cuda.h>
{# context lib="cuda" #}

-- Friends
import Foreign.CUDA.Driver.Error
import Foreign.CUDA.Internal.C2HS

-- System
import Foreign
import Foreign.C

-- |
-- Unwrap an Either construct, throwing an error for non-right values
--
forceEither :: Either String a -> a
forceEither (Left  s) = error s
forceEither (Right r) = r

-- |
-- Return the version number of the installed CUDA driver
--
driverVersion :: IO (Either String Int)
driverVersion =  resultIfOk `fmap` cuDriverGetVersion

{# fun unsafe cuDriverGetVersion
  { alloca- `Int' peekIntConv* } -> `Status' cToEnum #}

{-# LANGUAGE ForeignFunctionInterface #-}
--------------------------------------------------------------------------------
-- |
-- Module    : Foreign.CUDA.Driver.Utils
-- Copyright : (c) 2009 Trevor L. McDonell
-- License   : BSD
--
-- Utility functions
--
--------------------------------------------------------------------------------

module Foreign.CUDA.Driver.Utils
  (
    driverVersion
  )
  where

#include <cuda.h>
{# context lib="cuda" #}

-- Friends
import Foreign.CUDA.Driver.Error
import Foreign.CUDA.Internal.C2HS

-- System
import Foreign
import Foreign.C

-- |
-- Return the version number of the installed CUDA driver
--
driverVersion :: IO (Either String Int)
driverVersion =  resultIfOk `fmap` cuDriverGetVersion

{# fun unsafe cuDriverGetVersion
  { alloca- `Int' peekIntConv* } -> `Status' cToEnum #}


build fix for CUDA < 5

build fix for CUDA < 5

{-# LANGUAGE BangPatterns             #-}
{-# LANGUAGE CPP                      #-}
{-# LANGUAGE ForeignFunctionInterface #-}
{-# LANGUAGE GADTs                    #-}
{-# LANGUAGE TemplateHaskell          #-}
--------------------------------------------------------------------------------
-- |
-- Module    : Foreign.CUDA.Runtime.Exec
-- Copyright : [2009..2014] Trevor L. McDonell
-- License   : BSD
--
-- Kernel execution control for C-for-CUDA runtime interface
--
--------------------------------------------------------------------------------

module Foreign.CUDA.Runtime.Exec (

  -- * Kernel Execution
  Fun, FunAttributes(..), FunParam(..), CacheConfig(..),
  attributes, setConfig, setParams, setCacheConfig, launch, launchKernel,

) where

#include "cbits/stubs.h"
{# context lib="cudart" #}

-- Friends
import Foreign.CUDA.Runtime.Stream                      ( Stream(..), defaultStream )
import Foreign.CUDA.Runtime.Error
import Foreign.CUDA.Internal.C2HS

-- System
import Foreign
import Foreign.C
import Control.Monad

#c
typedef struct cudaFuncAttributes cudaFuncAttributes;
#endc

--------------------------------------------------------------------------------
-- Data Types
--------------------------------------------------------------------------------

-- |
-- A @__global__@ device function.
--
-- Note that the use of a string naming a function was deprecated in CUDA 4.1
-- and removed in CUDA 5.0.
--
#if CUDART_VERSION >= 5000
type Fun = FunPtr ()
#else
type Fun = String
#endif

--
-- Function Attributes
--
{# pointer *cudaFuncAttributes as ^ foreign -> FunAttributes nocode #}

data FunAttributes = FunAttributes
  {
    constSizeBytes           :: !Int64,
    localSizeBytes           :: !Int64,
    sharedSizeBytes          :: !Int64,
    maxKernelThreadsPerBlock :: !Int,   -- ^ maximum block size that can be successively launched (based on register usage)
    numRegs                  :: !Int    -- ^ number of registers required for each thread
  }
  deriving (Show)

instance Storable FunAttributes where
  sizeOf _    = {# sizeof cudaFuncAttributes #}
  alignment _ = alignment (undefined :: Ptr ())

  poke _ _    = error "Can not poke Foreign.CUDA.Runtime.FunAttributes"
  peek p      = do
    cs <- cIntConv `fmap` {#get cudaFuncAttributes.constSizeBytes#} p
    ls <- cIntConv `fmap` {#get cudaFuncAttributes.localSizeBytes#} p
    ss <- cIntConv `fmap` {#get cudaFuncAttributes.sharedSizeBytes#} p
    tb <- cIntConv `fmap` {#get cudaFuncAttributes.maxThreadsPerBlock#} p
    nr <- cIntConv `fmap` {#get cudaFuncAttributes.numRegs#} p

    return FunAttributes
      {
        constSizeBytes           = cs,
        localSizeBytes           = ls,
        sharedSizeBytes          = ss,
        maxKernelThreadsPerBlock = tb,
        numRegs                  = nr
      }

#if CUDART_VERSION < 3000
data CacheConfig
#else
-- |
-- Cache configuration preference
--
{# enum cudaFuncCache as CacheConfig
    { }
    with prefix="cudaFuncCachePrefer" deriving (Eq, Show) #}
#endif

-- |
-- Kernel function parameters. Doubles will be converted to an internal float
-- representation on devices that do not support doubles natively.
--
data FunParam where
  IArg :: !Int             -> FunParam
  FArg :: !Float           -> FunParam
  DArg :: !Double          -> FunParam
  VArg :: Storable a => !a -> FunParam

--------------------------------------------------------------------------------
-- Execution Control
--------------------------------------------------------------------------------

-- |
-- Obtain the attributes of the named @__global__@ device function. This
-- itemises the requirements to successfully launch the given kernel.
--
{-# INLINEABLE attributes #-}
attributes :: Fun -> IO FunAttributes
attributes !fn = resultIfOk =<< cudaFuncGetAttributes fn

{-# INLINE cudaFuncGetAttributes #-}
{# fun unsafe cudaFuncGetAttributes
  { alloca-  `FunAttributes' peek*
  , withFun* `Fun'                 } -> `Status' cToEnum #}

-- |
-- Specify the grid and block dimensions for a device call. Used in conjunction
-- with 'setParams', this pushes data onto the execution stack that will be
-- popped when a function is 'launch'ed.
--
{-# INLINEABLE setConfig #-}
setConfig :: (Int,Int)          -- ^ grid dimensions
          -> (Int,Int,Int)      -- ^ block dimensions
          -> Int64              -- ^ shared memory per block (bytes)
          -> Maybe Stream       -- ^ associated processing stream
          -> IO ()
setConfig (!gx,!gy) (!bx,!by,!bz) !sharedMem !mst =
  nothingIfOk =<<
    cudaConfigureCallSimple gx gy bx by bz sharedMem (maybe defaultStream id mst)

--
-- The FFI does not support passing deferenced structures to C functions, as
-- this is highly platform/compiler dependent. Wrap our own function stub
-- accepting plain integers.
--
{-# INLINE cudaConfigureCallSimple #-}
{# fun unsafe cudaConfigureCallSimple
  {           `Int', `Int'
  ,           `Int', `Int', `Int'
  , cIntConv  `Int64'
  , useStream `Stream'            } -> `Status' cToEnum #}

-- |
-- Set the argument parameters that will be passed to the next kernel
-- invocation. This is used in conjunction with 'setConfig' to control kernel
-- execution.
--
{-# INLINEABLE setParams #-}
setParams :: [FunParam] -> IO ()
setParams = foldM_ k 0
  where
    k !offset !arg = do
      let s = size arg
      set arg s offset >>= nothingIfOk
      return (offset + s)

    size (IArg _) = sizeOf (undefined :: Int)
    size (FArg _) = sizeOf (undefined :: Float)
    size (DArg _) = sizeOf (undefined :: Double)
    size (VArg a) = sizeOf a

    set (IArg v) s o = cudaSetupArgument v s o
    set (FArg v) s o = cudaSetupArgument v s o
    set (VArg v) s o = cudaSetupArgument v s o
    set (DArg v) s o =
      cudaSetDoubleForDevice v >>= resultIfOk >>= \d ->
      cudaSetupArgument d s o

{-# INLINE cudaSetupArgument #-}
{# fun unsafe cudaSetupArgument
  `Storable a' =>
  { with'* `a'
  ,        `Int'
  ,        `Int'   } -> `Status' cToEnum #}
  where
    with' v a = with v $ \p -> a (castPtr p)

{-# INLINE cudaSetDoubleForDevice #-}
{# fun unsafe cudaSetDoubleForDevice
  { with'* `Double' peek'* } -> `Status' cToEnum #}
  where
    with' v a = with v $ \p -> a (castPtr p)
    peek'     = peek . castPtr

-- |
-- On devices where the L1 cache and shared memory use the same hardware
-- resources, this sets the preferred cache configuration for the given device
-- function. This is only a preference; the driver is free to choose a different
-- configuration as required to execute the function.
--
-- Switching between configuration modes may insert a device-side
-- synchronisation point for streamed kernel launches
--
{-# INLINEABLE setCacheConfig #-}
setCacheConfig :: Fun -> CacheConfig -> IO ()
#if CUDART_VERSION < 3000
setCacheConfig _ _       = requireSDK 'setCacheConfig 3.0
#else
setCacheConfig !fn !pref = nothingIfOk =<< cudaFuncSetCacheConfig fn pref

{-# INLINE cudaFuncSetCacheConfig #-}
{# fun unsafe cudaFuncSetCacheConfig
  { withFun*  `Fun'
  , cFromEnum `CacheConfig' } -> `Status' cToEnum #}
#endif

-- |
-- Invoke the @__global__@ kernel function on the device. This must be preceded
-- by a call to 'setConfig' and (if appropriate) 'setParams'.
--
{-# INLINEABLE launch #-}
launch :: Fun -> IO ()
launch !fn = nothingIfOk =<< cudaLaunch fn

{-# INLINE cudaLaunch #-}
{# fun unsafe cudaLaunch
  { withFun* `Fun' } -> `Status' cToEnum #}

-- |
-- Invoke a kernel on a @(gx * gy)@ grid of blocks, where each block contains
-- @(tx * ty * tz)@ threads and has access to a given number of bytes of shared
-- memory. The launch may also be associated with a specific 'Stream'.
--
{-# INLINEABLE launchKernel #-}
launchKernel
    :: Fun              -- ^ Device function symbol
    -> (Int,Int)        -- ^ grid dimensions
    -> (Int,Int,Int)    -- ^ thread block shape
    -> Int64            -- ^ shared memory per block (bytes)
    -> Maybe Stream     -- ^ (optional) execution stream
    -> [FunParam]
    -> IO ()
launchKernel !fn !grid !block !sm !mst !args = do
  setConfig grid block sm mst
  setParams args
  launch fn

--------------------------------------------------------------------------------
-- Internals
--------------------------------------------------------------------------------

-- CUDA 5.0 changed the type of a kernel function from char* to void*
--
#if CUDART_VERSION >= 5000
withFun :: Fun -> (Ptr a -> IO b) -> IO b
withFun fn action = action (castFunPtrToPtr fn)
#else
withFun :: Fun -> (Ptr CChar -> IO a) -> IO a
withFun           = withCString
#endif

{-# LANGUAGE BangPatterns             #-}
{-# LANGUAGE CPP                      #-}
{-# LANGUAGE ForeignFunctionInterface #-}
{-# LANGUAGE GADTs                    #-}
{-# LANGUAGE TemplateHaskell          #-}
--------------------------------------------------------------------------------
-- |
-- Module    : Foreign.CUDA.Runtime.Exec
-- Copyright : [2009..2014] Trevor L. McDonell
-- License   : BSD
--
-- Kernel execution control for C-for-CUDA runtime interface
--
--------------------------------------------------------------------------------

module Foreign.CUDA.Runtime.Exec (

  -- * Kernel Execution
  Fun, FunAttributes(..), FunParam(..), CacheConfig(..),
  attributes, setConfig, setParams, setCacheConfig, launch, launchKernel,

) where

#include "cbits/stubs.h"
{# context lib="cudart" #}

-- Friends
import Foreign.CUDA.Runtime.Stream                      ( Stream(..), defaultStream )
import Foreign.CUDA.Runtime.Error
import Foreign.CUDA.Internal.C2HS

-- System
import Foreign
import Foreign.C
import Control.Monad

#c
typedef struct cudaFuncAttributes cudaFuncAttributes;
#endc

--------------------------------------------------------------------------------
-- Data Types
--------------------------------------------------------------------------------

-- |
-- A @__global__@ device function.
--
-- Note that the use of a string naming a function was deprecated in CUDA 4.1
-- and removed in CUDA 5.0.
--
#if CUDART_VERSION >= 5000
type Fun = FunPtr ()
#else
type Fun = String
#endif

--
-- Function Attributes
--
{# pointer *cudaFuncAttributes as ^ foreign -> FunAttributes nocode #}

data FunAttributes = FunAttributes
  {
    constSizeBytes           :: !Int64,
    localSizeBytes           :: !Int64,
    sharedSizeBytes          :: !Int64,
    maxKernelThreadsPerBlock :: !Int,   -- ^ maximum block size that can be successively launched (based on register usage)
    numRegs                  :: !Int    -- ^ number of registers required for each thread
  }
  deriving (Show)

instance Storable FunAttributes where
  sizeOf _    = {# sizeof cudaFuncAttributes #}
  alignment _ = alignment (undefined :: Ptr ())

  poke _ _    = error "Can not poke Foreign.CUDA.Runtime.FunAttributes"
  peek p      = do
    cs <- cIntConv `fmap` {#get cudaFuncAttributes.constSizeBytes#} p
    ls <- cIntConv `fmap` {#get cudaFuncAttributes.localSizeBytes#} p
    ss <- cIntConv `fmap` {#get cudaFuncAttributes.sharedSizeBytes#} p
    tb <- cIntConv `fmap` {#get cudaFuncAttributes.maxThreadsPerBlock#} p
    nr <- cIntConv `fmap` {#get cudaFuncAttributes.numRegs#} p

    return FunAttributes
      {
        constSizeBytes           = cs,
        localSizeBytes           = ls,
        sharedSizeBytes          = ss,
        maxKernelThreadsPerBlock = tb,
        numRegs                  = nr
      }

#if CUDART_VERSION < 3000
data CacheConfig
#else
-- |
-- Cache configuration preference
--
{# enum cudaFuncCache as CacheConfig
    { }
    with prefix="cudaFuncCachePrefer" deriving (Eq, Show) #}
#endif

-- |
-- Kernel function parameters. Doubles will be converted to an internal float
-- representation on devices that do not support doubles natively.
--
data FunParam where
  IArg :: !Int             -> FunParam
  FArg :: !Float           -> FunParam
  DArg :: !Double          -> FunParam
  VArg :: Storable a => !a -> FunParam

--------------------------------------------------------------------------------
-- Execution Control
--------------------------------------------------------------------------------

-- |
-- Obtain the attributes of the named @__global__@ device function. This
-- itemises the requirements to successfully launch the given kernel.
--
{-# INLINEABLE attributes #-}
attributes :: Fun -> IO FunAttributes
attributes !fn = resultIfOk =<< cudaFuncGetAttributes fn

{-# INLINE cudaFuncGetAttributes #-}
{# fun unsafe cudaFuncGetAttributes
  { alloca-  `FunAttributes' peek*
  , withFun* `Fun'                 } -> `Status' cToEnum #}

-- |
-- Specify the grid and block dimensions for a device call. Used in conjunction
-- with 'setParams', this pushes data onto the execution stack that will be
-- popped when a function is 'launch'ed.
--
{-# INLINEABLE setConfig #-}
setConfig :: (Int,Int)          -- ^ grid dimensions
          -> (Int,Int,Int)      -- ^ block dimensions
          -> Int64              -- ^ shared memory per block (bytes)
          -> Maybe Stream       -- ^ associated processing stream
          -> IO ()
setConfig (!gx,!gy) (!bx,!by,!bz) !sharedMem !mst =
  nothingIfOk =<<
    cudaConfigureCallSimple gx gy bx by bz sharedMem (maybe defaultStream id mst)

--
-- The FFI does not support passing deferenced structures to C functions, as
-- this is highly platform/compiler dependent. Wrap our own function stub
-- accepting plain integers.
--
{-# INLINE cudaConfigureCallSimple #-}
{# fun unsafe cudaConfigureCallSimple
  {           `Int', `Int'
  ,           `Int', `Int', `Int'
  , cIntConv  `Int64'
  , useStream `Stream'            } -> `Status' cToEnum #}

-- |
-- Set the argument parameters that will be passed to the next kernel
-- invocation. This is used in conjunction with 'setConfig' to control kernel
-- execution.
--
{-# INLINEABLE setParams #-}
setParams :: [FunParam] -> IO ()
setParams = foldM_ k 0
  where
    k !offset !arg = do
      let s = size arg
      set arg s offset >>= nothingIfOk
      return (offset + s)

    size (IArg _) = sizeOf (undefined :: Int)
    size (FArg _) = sizeOf (undefined :: Float)
    size (DArg _) = sizeOf (undefined :: Double)
    size (VArg a) = sizeOf a

    set (IArg v) s o = cudaSetupArgument v s o
    set (FArg v) s o = cudaSetupArgument v s o
    set (VArg v) s o = cudaSetupArgument v s o
    set (DArg v) s o =
      cudaSetDoubleForDevice v >>= resultIfOk >>= \d ->
      cudaSetupArgument d s o

{-# INLINE cudaSetupArgument #-}
{# fun unsafe cudaSetupArgument
  `Storable a' =>
  { with'* `a'
  ,        `Int'
  ,        `Int'   } -> `Status' cToEnum #}
  where
    with' v a = with v $ \p -> a (castPtr p)

{-# INLINE cudaSetDoubleForDevice #-}
{# fun unsafe cudaSetDoubleForDevice
  { with'* `Double' peek'* } -> `Status' cToEnum #}
  where
    with' v a = with v $ \p -> a (castPtr p)
    peek'     = peek . castPtr

-- |
-- On devices where the L1 cache and shared memory use the same hardware
-- resources, this sets the preferred cache configuration for the given device
-- function. This is only a preference; the driver is free to choose a different
-- configuration as required to execute the function.
--
-- Switching between configuration modes may insert a device-side
-- synchronisation point for streamed kernel launches
--
{-# INLINEABLE setCacheConfig #-}
setCacheConfig :: Fun -> CacheConfig -> IO ()
#if CUDART_VERSION < 3000
setCacheConfig _ _       = requireSDK 'setCacheConfig 3.0
#else
setCacheConfig !fn !pref = nothingIfOk =<< cudaFuncSetCacheConfig fn pref

{-# INLINE cudaFuncSetCacheConfig #-}
{# fun unsafe cudaFuncSetCacheConfig
  { withFun*  `Fun'
  , cFromEnum `CacheConfig' } -> `Status' cToEnum #}
#endif

-- |
-- Invoke the @__global__@ kernel function on the device. This must be preceded
-- by a call to 'setConfig' and (if appropriate) 'setParams'.
--
{-# INLINEABLE launch #-}
launch :: Fun -> IO ()
launch !fn = nothingIfOk =<< cudaLaunch fn

{-# INLINE cudaLaunch #-}
{# fun unsafe cudaLaunch
  { withFun* `Fun' } -> `Status' cToEnum #}

-- |
-- Invoke a kernel on a @(gx * gy)@ grid of blocks, where each block contains
-- @(tx * ty * tz)@ threads and has access to a given number of bytes of shared
-- memory. The launch may also be associated with a specific 'Stream'.
--
{-# INLINEABLE launchKernel #-}
launchKernel
    :: Fun              -- ^ Device function symbol
    -> (Int,Int)        -- ^ grid dimensions
    -> (Int,Int,Int)    -- ^ thread block shape
    -> Int64            -- ^ shared memory per block (bytes)
    -> Maybe Stream     -- ^ (optional) execution stream
    -> [FunParam]
    -> IO ()
launchKernel !fn !grid !block !sm !mst !args = do
  setConfig grid block sm mst
  setParams args
  launch fn

--------------------------------------------------------------------------------
-- Internals
--------------------------------------------------------------------------------

-- CUDA 5.0 changed the type of a kernel function from char* to void*
--
withFun :: Fun -> (Ptr a -> IO b) -> IO b
#if CUDART_VERSION >= 5000
withFun fn action = action (castFunPtrToPtr fn)
#else
withFun           = withCString
#endif