Reorder Call.chs

Reorder Call.chs

-- 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.IORef

import           Foreign.C.Types              as C
import qualified Foreign.ForeignPtr           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.ByteBuffer#}
{#import Network.Grpc.Lib.Metadata#}
{#import Network.Grpc.Lib.TimeSpec#}
{#import Network.Grpc.Lib.Types#}
{#import Network.Grpc.Lib.Grpc#}

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

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

type Deadline = TimeSpec

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

emptyChannelArgs :: ChannelArgs
emptyChannelArgs = ChannelArgs C.nullPtr

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

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

withTimeout :: Deadline -> ClientContext -> ClientContext
withTimeout ts (ClientContext chan cq cqt _) = ClientContext chan cq cqt ts

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 -> MethodName -> Arg -> [Metadata] -> IO (RpcReply (UnaryResult L.ByteString))
callUnary ctx@(ClientContext chan cq _ deadline) method arg mds =
  C.withForeignPtr chan $ \chanPtr ->
    bracket (grpcChannelCreateCall chanPtr C.nullPtr defaultPropagationMask cq method "localhost" deadline) grpcCallDestroy $ \call0 -> newMVar call0 >>= \mcall -> do
      crw <- newClientReaderWriter ctx mcall

      sendInitOp <- opSendInitialMetadata mds
      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 -> MethodName -> Arg -> IO (RpcReply (Client B.ByteString L.ByteString))
callDownstream ctx@(ClientContext chan cq _ deadline) method arg =
  C.withForeignPtr chan $ \chanPtr -> do
    mcall <- grpcChannelCreateCall chanPtr C.nullPtr defaultPropagationMask cq method "localhost" deadline >>= newMVar
    crw <- newClientReaderWriter ctx mcall

    sendInitOp <- opSendInitialMetadata []
    sendCloseOp <- opSendCloseFromClient
    sendMessageOp <- opSendMessage arg

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

callUpstream :: ClientContext -> MethodName -> IO (RpcReply (Client B.ByteString L.ByteString))
callUpstream ctx@(ClientContext chan cq _ deadline) method =
  C.withForeignPtr chan $ \chanPtr -> do
    mcall <- grpcChannelCreateCall chanPtr C.nullPtr defaultPropagationMask cq method "localhost" deadline >>= newMVar
    crw <- newClientReaderWriter ctx mcall

    sendInitOp <- opSendInitialMetadata []

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

callBidi :: ClientContext -> MethodName -> [Metadata] -> IO (RpcReply (Client B.ByteString L.ByteString))
callBidi ctx@(ClientContext chan cq _ deadline) method mds = do
  C.withForeignPtr chan $ \chanPtr -> do
    mcall <- grpcChannelCreateCall chanPtr C.nullPtr defaultPropagationMask cq method "localhost" deadline >>= newMVar

    crw <- newClientReaderWriter ctx mcall
    sendInitOp <- opSendInitialMetadata mds

    res <- callBatch crw [ OpX sendInitOp ]
    case res of
      RpcOk _ -> return (RpcOk (Client crw defaultEncoder defaultDecoder))
      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 :: !(IORef (Maybe RpcStatus))
}

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

data OpX = forall t. OpX (OpT t)

data OpArray = OpArray { opArrPtr :: !GrpcOpPtr
                       , opArrLen :: !CULong
                       , opArrFree :: !(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
  CQ.withEvent (ccWorker ctx) $ \eDesc -> do
    arr <- toArray ops
    callStatus <- withMVar (callMVar_ crw) $ \call ->
      grpcCallStartBatch call (opArrPtr arr) (opArrLen arr) (CQ.eventTag eDesc) reservedPtr
    -- print ("callStatus: " ++ show callStatus)
    case callStatus of
      CallOk -> do
        e <- CQ.interruptibleWaitEvent eDesc
        -- print ("event: " ++ show e)
        opArrFree arr -- TODO: We leak if we're interrupted.
        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
         -- print callStatus
         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
      writeIORef statusFromServer (Just 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 crw@(ClientReaderWriter{..}) = do
  status <- readIORef statusFromServer
  case status of
    Nothing -> do
      recvStatusOp <- opRecvStatusOnClient crw
      res <- callBatch crw [ OpX recvStatusOp ]
      case res of
        RpcOk _ -> do
          st <- opRead recvStatusOp
          return (RpcOk st)
        RpcError err -> do
          return (RpcError err)
    Just st -> return (RpcOk st)

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 { callMVar_ = callMVar } =
  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)

initialMetadata :: Rpc req resp [Metadata]
initialMetadata = do
  crw <- askCrw
  joinReply =<< liftIO (clientWaitForInitialMetadata crw)

waitForStatus :: Rpc req resp RpcStatus
waitForStatus = do
  crw <- askCrw
  joinReply =<< liftIO (clientWaitForStatus crw)

receiveMessage :: Rpc req resp (Maybe resp)
receiveMessage = do
  crw <- askCrw
  msg <- joinReply =<< liftIO (clientRead crw)
  case msg of
    Nothing -> return Nothing
    Just x -> do
      decoder <- askDecoder
      liftM Just (joinReply =<< liftIO (decoder x))

receiveAllMessages :: Rpc req resp [resp]
receiveAllMessages = do
  crw <- askCrw
  decoder <- askDecoder
  let go acc = do
        value <- joinReply =<< liftIO (clientRead crw)
        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 o = do
  crw <- askCrw
  encoder <- askEncoder
  x <- joinReply =<< liftIO (encoder o)
  joinReply =<< liftIO (clientWrite crw x)

sendHalfClose :: Rpc req resp ()
sendHalfClose = do
  crw <- askCrw
  joinReply =<< liftIO (clientSendHalfClose crw)

closeCall :: Rpc req resp ()
closeCall = do
  crw <- askCrw
  liftIO (clientClose crw)

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.IORef

import           Foreign.C.Types              as C
import qualified Foreign.ForeignPtr           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.ByteBuffer#}
{#import Network.Grpc.Lib.Metadata#}
{#import Network.Grpc.Lib.TimeSpec#}
{#import Network.Grpc.Lib.Types#}
{#import Network.Grpc.Lib.Grpc#}

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

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

type Deadline = TimeSpec

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

emptyChannelArgs :: ChannelArgs
emptyChannelArgs = ChannelArgs C.nullPtr

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

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

withTimeout :: Deadline -> ClientContext -> ClientContext
withTimeout ts (ClientContext chan cq cqt _) = ClientContext chan cq cqt ts

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

callUnary :: ClientContext -> MethodName -> Arg -> [Metadata] -> IO (RpcReply (UnaryResult L.ByteString))
callUnary ctx@(ClientContext chan cq _ deadline) method arg mds =
  C.withForeignPtr chan $ \chanPtr ->
    bracket (grpcChannelCreateCall chanPtr C.nullPtr defaultPropagationMask cq method "localhost" deadline) grpcCallDestroy $ \call0 -> newMVar call0 >>= \mcall -> do
      crw <- newClientReaderWriter ctx mcall

      sendInitOp <- opSendInitialMetadata mds
      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)

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)

callDownstream :: ClientContext -> MethodName -> Arg -> IO (RpcReply (Client B.ByteString L.ByteString))
callDownstream ctx@(ClientContext chan cq _ deadline) method arg =
  C.withForeignPtr chan $ \chanPtr -> do
    mcall <- grpcChannelCreateCall chanPtr C.nullPtr defaultPropagationMask cq method "localhost" deadline >>= newMVar
    crw <- newClientReaderWriter ctx mcall

    sendInitOp <- opSendInitialMetadata []
    sendCloseOp <- opSendCloseFromClient
    sendMessageOp <- opSendMessage arg

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

callUpstream :: ClientContext -> MethodName -> IO (RpcReply (Client B.ByteString L.ByteString))
callUpstream ctx@(ClientContext chan cq _ deadline) method =
  C.withForeignPtr chan $ \chanPtr -> do
    mcall <- grpcChannelCreateCall chanPtr C.nullPtr defaultPropagationMask cq method "localhost" deadline >>= newMVar
    crw <- newClientReaderWriter ctx mcall

    sendInitOp <- opSendInitialMetadata []

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

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

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

data OpX = forall t. OpX (OpT t)

data OpArray = OpArray { opArrPtr :: !GrpcOpPtr
                       , opArrLen :: !CULong
                       , opArrFree :: !(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
  CQ.withEvent (ccWorker ctx) $ \eDesc -> do
    arr <- toArray ops
    callStatus <- withMVar (callMVar_ crw) $ \call ->
      grpcCallStartBatch call (opArrPtr arr) (opArrLen arr) (CQ.eventTag eDesc) reservedPtr
    -- print ("callStatus: " ++ show callStatus)
    case callStatus of
      CallOk -> do
        e <- CQ.interruptibleWaitEvent eDesc
        -- print ("event: " ++ show e)
        opArrFree arr -- TODO: We leak if we're interrupted.
        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
         -- print callStatus
         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
      writeIORef statusFromServer (Just 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 crw@(ClientReaderWriter{..}) = do
  status <- readIORef statusFromServer
  case status of
    Nothing -> do
      recvStatusOp <- opRecvStatusOnClient crw
      res <- callBatch crw [ OpX recvStatusOp ]
      case res of
        RpcOk _ -> do
          st <- opRead recvStatusOp
          return (RpcOk st)
        RpcError err -> do
          return (RpcError err)
    Just st -> return (RpcOk st)

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 { callMVar_ = callMVar } =
  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)

initialMetadata :: Rpc req resp [Metadata]
initialMetadata = do
  crw <- askCrw
  joinReply =<< liftIO (clientWaitForInitialMetadata crw)

waitForStatus :: Rpc req resp RpcStatus
waitForStatus = do
  crw <- askCrw
  joinReply =<< liftIO (clientWaitForStatus crw)

receiveMessage :: Rpc req resp (Maybe resp)
receiveMessage = do
  crw <- askCrw
  msg <- joinReply =<< liftIO (clientRead crw)
  case msg of
    Nothing -> return Nothing
    Just x -> do
      decoder <- askDecoder
      liftM Just (joinReply =<< liftIO (decoder x))

receiveAllMessages :: Rpc req resp [resp]
receiveAllMessages = do
  crw <- askCrw
  decoder <- askDecoder
  let go acc = do
        value <- joinReply =<< liftIO (clientRead crw)
        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 o = do
  crw <- askCrw
  encoder <- askEncoder
  x <- joinReply =<< liftIO (encoder o)
  joinReply =<< liftIO (clientWrite crw x)

sendHalfClose :: Rpc req resp ()
sendHalfClose = do
  crw <- askCrw
  joinReply =<< liftIO (clientSendHalfClose crw)

closeCall :: Rpc req resp ()
closeCall = do
  crw <- askCrw
  liftIO (clientClose crw)

callBidi :: ClientContext -> MethodName -> [Metadata] -> IO (RpcReply (Client B.ByteString L.ByteString))
callBidi ctx@(ClientContext chan cq _ deadline) method mds = do
  C.withForeignPtr chan $ \chanPtr -> do
    mcall <- grpcChannelCreateCall chanPtr C.nullPtr defaultPropagationMask cq method "localhost" deadline >>= newMVar

    crw <- newClientReaderWriter ctx mcall
    sendInitOp <- opSendInitialMetadata mds

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

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

data RpcReply a
  = RpcOk a
  | RpcError RpcError
  deriving Show

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


haskell-bindings-improve-cmake: fix reference counting issue with keyListMeta

module Elektra.Key (Key (..), Namespace (..), withKey,
    keyNew, keyNewWithValue, keyDup, keyCopy, keyClear, keyIncRef, keyDecRef, keyGetRef,
    keyName, keyGetNameSize, keyUnescapedName, keyGetUnescapedNameSize, keySetName, keyGetFullNameSize, keyGetFullName,
    keyAddName,  keyBaseName, keyGetBaseName, keyGetBaseNameSize, keyAddBaseName, keySetBaseName, keyGetNamespace,
    keyString, keyGetValueSize, keySetString, keySet,
    keyRewindMeta, keyNextMeta, keyCurrentMeta, keyCopyMeta, keyCopyAllMeta, keyGetMeta, keySetMeta, keyListMeta,
    keyCmp, keyNeedSync, keyIsBelow, keyIsDirectBelow, keyRel, keyIsInactive, keyIsBinary, keyIsString, keyPtrNull, ifKey) where

#include <kdb.h>
import Foreign.Marshal.Alloc (allocaBytes)
import Foreign.Ptr (castPtr, nullPtr)
import Foreign.ForeignPtr (withForeignPtr)
import System.IO.Unsafe (unsafePerformIO)
import Control.Monad (liftM)

{#context lib="libelektra" #}

-- ***
-- TYPE DEFINITIONS
-- ***

{#pointer *Key foreign finalizer keyDel newtype #}

{#typedef size_t Int#}
{#typedef ssize_t Int #}

{#enum KEY_NS_NONE as Namespace { underscoreToCase } deriving (Show, Eq) #}
{#enum KEY_NAME as ElektraKeyVarargs { underscoreToCase } deriving (Show, Eq) #}

-- ***
-- KEY CREATION / DELETION / COPPY METHODS
-- ***

keyPtrNull :: Key -> IO Bool
keyPtrNull (Key ptr) = withForeignPtr ptr (return . (== nullPtr))

-- as we use haskell's reference counting here, increase the number by one
-- so it gets deleted properly when haskell calls the finalizer
keyNew :: String -> IO Key
keyNew name = do
    key <- keyNewRaw name KeyEnd
    keyIncRef key
    return key
keyNewWithValue :: String -> String -> IO Key
keyNewWithValue name value = do
    key <- keyNewRawWithValue name KeyValue value KeyEnd
    keyIncRef key
    return key
{#fun unsafe variadic keyNew[keyswitch_t] as keyNewRaw {`String', `ElektraKeyVarargs'} -> `Key' #}
{#fun unsafe variadic keyNew[keyswitch_t, const char *, keyswitch_t]
    as keyNewRawWithValue {`String', `ElektraKeyVarargs', `String', `ElektraKeyVarargs'} -> `Key' #}
keyDup :: Key -> IO Key
keyDup key = do
    dup <- keyDupRaw key
    keyIncRef dup
    return dup
{#fun unsafe keyDup as keyDupRaw {`Key'} -> `Key' #}
{#fun unsafe keyCopy {`Key', `Key'} -> `Int' #}
{#fun unsafe keyClear {`Key'} -> `Int' #}
{#fun unsafe keyIncRef {`Key'} -> `Int' #}
{#fun unsafe keyDecRef {`Key'} -> `Int' #}
{#fun unsafe keyGetRef {`Key'} -> `Int' #}

-- ***
-- KEY NAME MANIPULATION METHODS
-- ***

{#fun unsafe keyName {`Key'} -> `String' #}
keyGetName = keyName
{#fun unsafe keyGetNameSize {`Key'} -> `Int' #}
keyUnescapedName :: Key -> IO String
keyUnescapedName key = do
    size <- keyGetUnescapedNameSize key
    withKey key $ (\cKey -> do
        result <- {#call unsafe keyUnescapedName as keyUnescapedNameRaw #} cKey
        C2HSImp.peekCStringLen (castPtr result, size))
{#fun unsafe keyGetUnescapedNameSize {`Key'} -> `Int' #}
{#fun unsafe keySetName {`Key', `String'} -> `Int' #}
keyGetFullName :: (Key) -> IO String
keyGetFullName key = do
    size <- keyGetFullNameSize key
    withKey key $ \cKey ->
        allocaBytes size (\result -> do
            {#call unsafe keyGetFullName as keyGetFullNameRaw #} cKey result size
            C2HSImp.peekCString result)
{#fun unsafe keyGetFullNameSize {`Key'} -> `Int' #}
{#fun unsafe keyBaseName {`Key'} -> `String' #}
keyGetBaseName :: Key -> IO String
keyGetBaseName = keyBaseName
{#fun unsafe keyGetBaseNameSize {`Key'} -> `Int' #}
{#fun unsafe keyAddBaseName {`Key', `String'} -> `Int' #}
{#fun unsafe keyAddName {`Key', `String'} -> `Int' #}
{#fun unsafe keySetBaseName {`Key', `String'} -> `Int' #}
{#fun unsafe keyGetNamespace {`Key'} -> `Namespace' #}

-- ***
-- KEY VALUE MANIPULATION METHODS
-- ***

-- we don't handle binary data currently
-- {#fun unsafe keyValue {`Key'} -> `PPtr' #}
{#fun unsafe keyString {`Key'} -> `String' #}
{#fun unsafe keyGetValueSize {`Key'} -> `Int' #}
keySet :: Show a => Key -> a -> IO Int
keySet key = keySetString key . show
{#fun unsafe keySetString {`Key', `String'} -> `Int' #}

-- ***
-- KEY META MANIPULATION METHODS
-- ***

{#fun unsafe keyRewindMeta {`Key'} -> `Int' #}
{#fun unsafe keyNextMeta {`Key'} -> `Key' #}
{#fun unsafe keyCurrentMeta {`Key'} -> `Key' #}
{#fun unsafe keyCopyMeta {`Key', `Key', `String'} -> `Int' #}
{#fun unsafe keyCopyAllMeta {`Key', `Key'} -> `Int' #}
{#fun unsafe keyGetMeta {`Key', `String'} -> `Key' #}
{#fun unsafe keySetMeta {`Key', `String', `String'} -> `Int' #}
keyListMeta :: Key -> IO [Key]
keyListMeta key = keyRewindMeta key >> listMeta []
    where
        listMeta res = do
            cur <- keyNextMeta key
            isNull <- keyPtrNull cur
            if isNull then return res else keyIncRef cur >> liftM (cur :) (listMeta res)

-- ***
-- KEY TESTING METHODS
-- ***

{#fun unsafe keyCmp {`Key', `Key'} -> `Int' #}
{#fun unsafe keyNeedSync {`Key'} -> `Int' #}
{#fun unsafe keyIsBelow {`Key', `Key'} -> `Int' #}
{#fun unsafe keyIsDirectBelow {`Key', `Key'} -> `Int' #}
{#fun unsafe keyRel {`Key', `Key'} -> `Int' #}
{#fun unsafe keyIsInactive {`Key'} -> `Int' #}
{#fun unsafe keyIsBinary {`Key'} -> `Int' #}
{#fun unsafe keyIsString {`Key'} -> `Int' #}

-- ***
-- COMMON HASKELL TYPE CLASSES
-- unsafePerformIO should be ok here, as those functions don't alter the key's state
-- ***

instance Show Key where
    show key = unsafePerformIO $ do
        name <- keyName key
        value <- keyString key
        ref <- keyGetRef key
        return $ name ++ " " ++ value ++ " " ++ (show ref)

instance Eq Key where
    key1 == key2 = unsafePerformIO $ fmap (== 0) $ keyCmp key1 key2

-- ***
-- ADDITIONAL HELPERS USEFUL IN HASKELL
-- ***

ifKey :: IO Key -> (Key -> IO a) -> IO a -> IO a
ifKey k t f = do
    null <- k >>= keyPtrNull
    if null then f else k >>= t

module Elektra.Key (Key (..), Namespace (..), withKey,
    keyNew, keyNewWithValue, keyDup, keyCopy, keyClear, keyIncRef, keyDecRef, keyGetRef,
    keyName, keyGetNameSize, keyUnescapedName, keyGetUnescapedNameSize, keySetName, keyGetFullNameSize, keyGetFullName,
    keyAddName,  keyBaseName, keyGetBaseName, keyGetBaseNameSize, keyAddBaseName, keySetBaseName, keyGetNamespace,
    keyString, keyGetValueSize, keySetString, keySet,
    keyRewindMeta, keyNextMeta, keyCurrentMeta, keyCopyMeta, keyCopyAllMeta, keyGetMeta, keySetMeta, keyListMeta,
    keyCmp, keyNeedSync, keyIsBelow, keyIsDirectBelow, keyRel, keyIsInactive, keyIsBinary, keyIsString, keyPtrNull, ifKey) where

#include <kdb.h>
import Foreign.Marshal.Alloc (allocaBytes)
import Foreign.Ptr (castPtr, nullPtr)
import Foreign.ForeignPtr (withForeignPtr)
import System.IO.Unsafe (unsafePerformIO)
import Control.Monad (liftM)

{#context lib="libelektra" #}

-- ***
-- TYPE DEFINITIONS
-- ***

{#pointer *Key foreign finalizer keyDel newtype #}

{#typedef size_t Int#}
{#typedef ssize_t Int #}

{#enum KEY_NS_NONE as Namespace { underscoreToCase } deriving (Show, Eq) #}
{#enum KEY_NAME as ElektraKeyVarargs { underscoreToCase } deriving (Show, Eq) #}

-- ***
-- KEY CREATION / DELETION / COPPY METHODS
-- ***

keyPtrNull :: Key -> IO Bool
keyPtrNull (Key ptr) = withForeignPtr ptr (return . (== nullPtr))

-- as we use haskell's reference counting here, increase the number by one
-- so it gets deleted properly when haskell calls the finalizer
keyNew :: String -> IO Key
keyNew name = do
    key <- keyNewRaw name KeyEnd
    keyIncRef key
    return key
keyNewWithValue :: String -> String -> IO Key
keyNewWithValue name value = do
    key <- keyNewRawWithValue name KeyValue value KeyEnd
    keyIncRef key
    return key
{#fun unsafe variadic keyNew[keyswitch_t] as keyNewRaw {`String', `ElektraKeyVarargs'} -> `Key' #}
{#fun unsafe variadic keyNew[keyswitch_t, const char *, keyswitch_t]
    as keyNewRawWithValue {`String', `ElektraKeyVarargs', `String', `ElektraKeyVarargs'} -> `Key' #}
keyDup :: Key -> IO Key
keyDup key = do
    dup <- keyDupRaw key
    keyIncRef dup
    return dup
{#fun unsafe keyDup as keyDupRaw {`Key'} -> `Key' #}
{#fun unsafe keyCopy {`Key', `Key'} -> `Int' #}
{#fun unsafe keyClear {`Key'} -> `Int' #}
{#fun unsafe keyIncRef {`Key'} -> `Int' #}
{#fun unsafe keyDecRef {`Key'} -> `Int' #}
{#fun unsafe keyGetRef {`Key'} -> `Int' #}

-- ***
-- KEY NAME MANIPULATION METHODS
-- ***

{#fun unsafe keyName {`Key'} -> `String' #}
keyGetName = keyName
{#fun unsafe keyGetNameSize {`Key'} -> `Int' #}
keyUnescapedName :: Key -> IO String
keyUnescapedName key = do
    size <- keyGetUnescapedNameSize key
    withKey key $ (\cKey -> do
        result <- {#call unsafe keyUnescapedName as keyUnescapedNameRaw #} cKey
        C2HSImp.peekCStringLen (castPtr result, size))
{#fun unsafe keyGetUnescapedNameSize {`Key'} -> `Int' #}
{#fun unsafe keySetName {`Key', `String'} -> `Int' #}
keyGetFullName :: (Key) -> IO String
keyGetFullName key = do
    size <- keyGetFullNameSize key
    withKey key $ \cKey ->
        allocaBytes size (\result -> do
            {#call unsafe keyGetFullName as keyGetFullNameRaw #} cKey result size
            C2HSImp.peekCString result)
{#fun unsafe keyGetFullNameSize {`Key'} -> `Int' #}
{#fun unsafe keyBaseName {`Key'} -> `String' #}
keyGetBaseName :: Key -> IO String
keyGetBaseName = keyBaseName
{#fun unsafe keyGetBaseNameSize {`Key'} -> `Int' #}
{#fun unsafe keyAddBaseName {`Key', `String'} -> `Int' #}
{#fun unsafe keyAddName {`Key', `String'} -> `Int' #}
{#fun unsafe keySetBaseName {`Key', `String'} -> `Int' #}
{#fun unsafe keyGetNamespace {`Key'} -> `Namespace' #}

-- ***
-- KEY VALUE MANIPULATION METHODS
-- ***

-- we don't handle binary data currently
-- {#fun unsafe keyValue {`Key'} -> `PPtr' #}
{#fun unsafe keyString {`Key'} -> `String' #}
{#fun unsafe keyGetValueSize {`Key'} -> `Int' #}
keySet :: Show a => Key -> a -> IO Int
keySet key = keySetString key . show
{#fun unsafe keySetString {`Key', `String'} -> `Int' #}

-- ***
-- KEY META MANIPULATION METHODS
-- ***

{#fun unsafe keyRewindMeta {`Key'} -> `Int' #}
{#fun unsafe keyNextMeta {`Key'} -> `Key' #}
{#fun unsafe keyCurrentMeta {`Key'} -> `Key' #}
{#fun unsafe keyCopyMeta {`Key', `Key', `String'} -> `Int' #}
{#fun unsafe keyCopyAllMeta {`Key', `Key'} -> `Int' #}
{#fun unsafe keyGetMeta {`Key', `String'} -> `Key' #}
{#fun unsafe keySetMeta {`Key', `String', `String'} -> `Int' #}
keyListMeta :: Key -> IO [Key]
keyListMeta key = keyRewindMeta key >> listMeta []
    where
        listMeta res = do
            cur <- keyNextMeta key
            isNull <- keyPtrNull cur
            if isNull then return res else liftM (cur :) (listMeta res)

-- ***
-- KEY TESTING METHODS
-- ***

{#fun unsafe keyCmp {`Key', `Key'} -> `Int' #}
{#fun unsafe keyNeedSync {`Key'} -> `Int' #}
{#fun unsafe keyIsBelow {`Key', `Key'} -> `Int' #}
{#fun unsafe keyIsDirectBelow {`Key', `Key'} -> `Int' #}
{#fun unsafe keyRel {`Key', `Key'} -> `Int' #}
{#fun unsafe keyIsInactive {`Key'} -> `Int' #}
{#fun unsafe keyIsBinary {`Key'} -> `Int' #}
{#fun unsafe keyIsString {`Key'} -> `Int' #}

-- ***
-- COMMON HASKELL TYPE CLASSES
-- unsafePerformIO should be ok here, as those functions don't alter the key's state
-- ***

instance Show Key where
    show key = unsafePerformIO $ do
        name <- keyName key
        value <- keyString key
        ref <- keyGetRef key
        return $ name ++ " " ++ value ++ " " ++ (show ref)

instance Eq Key where
    key1 == key2 = unsafePerformIO $ fmap (== 0) $ keyCmp key1 key2

-- ***
-- ADDITIONAL HELPERS USEFUL IN HASKELL
-- ***

ifKey :: IO Key -> (Key -> IO a) -> IO a -> IO a
ifKey k t f = do
    null <- k >>= keyPtrNull
    if null then f else k >>= t


automatically release device memory

automatically release device memory

Ignore-this: 2d19d6510749575a05bfd806bd59e7ac

Not yet convinced this is the right way to do it, but the functionality (when
implemented correctly, at least) would be useful.

darcs-hash:20090916021904-115f9-8ba8dbbd6e152e612e912a8558a96e6f89b24aa7.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,
    withDevicePtr,

    -- ** Dynamic allocation
    --
    -- Basic methods to allocate a block of memory on the device. The memory
    -- should be deallocated using 'free' when no longer necessary. The
    -- advantage is that failed allocations will not raise an exception.
    --
    free,
    malloc,
    malloc2D,
    malloc3D,
    memset,
    memset2D,
    memset3D,

    -- ** Local allocation
    --
    -- Execute a computation, passing as argument a pointer to a newly allocated
    -- block of memory. The memory is freed when the computation terminates.
    --
    alloca,
    allocaBytes,
    allocaBytesMemset,

    -- ** Marshalling
    --
    -- Store and retrieve Haskell lists as arrays on the device
    --
    peek,
    poke,
    peekArray,
    pokeArray,

    -- ** Combined allocation and marshalling
    --
    -- Write Haskell values (or lists) to newly allocated device memory. This
    -- may be a temporary store.
    --
    new,
    with,
    newArray,
    withArray,
    withArrayLen,

    -- ** Copying
    --
    -- Copy data between the host and device
    --
  )
  where

import Data.Int
import Control.Exception

import Foreign.C
import Foreign.Ptr
import Foreign.Concurrent
import Foreign.ForeignPtr hiding (newForeignPtr, addForeignPtrFinalizer)
import Foreign.Storable (Storable)

import qualified Foreign.Storable as F
import qualified Foreign.Marshal  as F

import Foreign.CUDA.Utils
import Foreign.CUDA.Error
import Foreign.CUDA.Stream
import Foreign.CUDA.Internal.C2HS

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

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

-- |
-- A reference to data stored on the device
--
type DevicePtr a = ForeignPtr a

-- |
-- Unwrap the device pointer, yielding a device memory address that can be
-- passed to a kernel function callable from the host code. This allows the type
-- system to strictly separate host and device memory references, a non-obvious
-- distinction in pure CUDA.
--
withDevicePtr :: DevicePtr a -> (Ptr a -> IO b) -> IO b
withDevicePtr =  withForeignPtr

--
-- Internal, unwrap and cast. Required as some functions such as memset work
-- with untyped `void' data.
--
withDevicePtr' :: DevicePtr a -> (Ptr b -> IO c) -> IO c
withDevicePtr' =  withForeignPtr . castForeignPtr

--
-- Wrap a device memory pointer. The memory will be freed once the last
-- reference to the data is dropped, although there is no guarantee as to
-- exactly when this will occur.
--
newDevicePtr   :: Ptr a -> IO (DevicePtr a)
newDevicePtr p =  newForeignPtr p (free_ 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 a))
malloc bytes =  do
    (rv,ptr) <- cudaMalloc bytes
    case rv of
        Success -> Right `fmap` newDevicePtr (castPtr ptr)
        _       -> return . Left . describe $ rv

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

-- |
-- 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 a, Int64))
malloc2D (width,height) =  do
    (rv,ptr,pitch) <- cudaMallocPitch width height
    case rv of
        Success -> (\p -> Right (p,pitch)) `fmap` newDevicePtr (castPtr ptr)
        _       -> return . Left . describe $ rv

{# fun unsafe cudaMallocPitch
    { alloca'-  `Ptr ()' peek'*       ,
      alloca'-  `Int64'  peekIntConv* ,
      cIntConv  `Int64'               ,
      cIntConv  `Int64'               } -> `Status' cToEnum #}
    where
        -- C->Haskell doesn't like qualified imports
        --
        alloca' :: Storable a => (Ptr a -> IO b) -> IO b
        alloca' =  F.alloca
        peek'   =  F.peek

-- |
-- 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 a, Int64))
malloc3D = moduleErr "malloc3D" "not implemented yet"

-- |
-- Free previously allocated memory on the device
--

--free   :: DevicePtr a -> IO (Maybe String)
--free p =  nothingIfOk `fmap` (withDevicePtr p cudaFree)

free :: DevicePtr a -> IO ()
free =  finalizeForeignPtr

free_   :: Ptr a -> IO ()
free_ p =  cudaFree p >>= \rv ->
    case rv of
        Success -> return ()
        _       -> moduleErr "free" (describe rv)

{# fun unsafe cudaFree
    { castPtr `Ptr a' } -> `Status' cToEnum #}

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

-- |
-- Initialise device memory to a given value
--
memset :: DevicePtr a                   -- ^ The device memory
       -> Int64                         -- ^ Number of bytes
       -> Int                           -- ^ Value to set for each byte
       -> IO (Maybe String)
memset ptr bytes symbol =
    nothingIfOk `fmap` cudaMemset ptr symbol bytes

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

-- |
-- Initialise a matrix to a given value
--
memset2D :: DevicePtr a                 -- ^ 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 a' ,
      cIntConv        `Int64'       ,
                      `Int'         ,
      cIntConv        `Int64'       ,
      cIntConv        `Int64'       } -> `Status' cToEnum #}

-- |
-- Initialise the elements of a 3D array to a given value
--
memset3D :: DevicePtr a                 -- ^ 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"

--------------------------------------------------------------------------------
-- Local allocation
--------------------------------------------------------------------------------

-- |
-- Execute a computation, passing a pointer to a temporarily allocated block of
-- memory
--
alloca :: Storable a => (DevicePtr a -> IO b) -> IO b
alloca =  doAlloca undefined
  where
    doAlloca   :: Storable a' => a' -> (DevicePtr a' -> IO b') -> IO b'
    doAlloca x =  allocaBytes (fromIntegral (F.sizeOf x))

-- |
-- Execute a computation, passing a pointer to a block of 'n' bytes of memory
--
allocaBytes :: Int64 -> (DevicePtr a -> IO b) -> IO b
allocaBytes bytes =
    bracket (forceEither `fmap` malloc bytes) free

-- |
-- Execute a computation, passing a pointer to a  block of memory initialised to
-- a given value
--
allocaBytesMemset :: Int64 -> Int -> (DevicePtr a -> IO b) -> IO b
allocaBytesMemset bytes symbol f =
    allocaBytes bytes $ \dptr ->
    memset dptr bytes symbol >>= \rv ->
      case rv of
        Nothing -> f dptr
        Just e  -> error e

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

-- |
-- Retrieve the specified value from device memory
--
peek :: Storable a => DevicePtr a -> IO (Either String a)
peek d = doPeek undefined
  where
    doPeek   :: Storable a' => a' -> IO (Either String a')
    doPeek x =  F.alloca         $ \hptr ->
                withDevicePtr' d $ \dptr ->
                memcpy hptr dptr (fromIntegral (F.sizeOf x)) DeviceToHost >>= \rv ->
                case rv of
                    Nothing -> Right `fmap` F.peek hptr
                    Just s  -> return (Left s)

-- |
-- Retrieve the specified number of elements from device memory and return as a
-- Haskell list
--
peekArray :: Storable a => Int -> DevicePtr a -> IO (Either String [a])
peekArray n d = doPeek undefined
  where
    doPeek   :: Storable a' => a' -> IO (Either String [a'])
    doPeek x =  let bytes = fromIntegral n * (fromIntegral (F.sizeOf x)) in
                F.allocaArray  n $ \hptr ->
                withDevicePtr' d $ \dptr ->
                memcpy hptr dptr bytes DeviceToHost >>= \rv ->
                case rv of
                   Nothing -> Right `fmap` F.peekArray n hptr
                   Just s  -> return (Left s)

-- |
-- Store the given value to device memory
--
poke :: Storable a => DevicePtr a -> a -> IO (Maybe String)
poke d v =
    F.with v        $ \hptr ->
    withDevicePtr d $ \dptr ->
    memcpy dptr hptr (fromIntegral (F.sizeOf v)) HostToDevice

-- |
-- Store the list elements consecutively in device memory
--
pokeArray :: Storable a => DevicePtr a -> [a] -> IO (Maybe String)
pokeArray d = doPoke undefined
  where
    doPoke     :: Storable a' => a' -> [a'] -> IO (Maybe String)
    doPoke x v =  F.withArrayLen v $ \n hptr ->
                  withDevicePtr' d $ \dptr ->
                  let bytes = (fromIntegral n) * (fromIntegral (F.sizeOf x)) in
                  memcpy dptr hptr bytes HostToDevice

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

-- |
-- Allocate a block of memory on the device and transfer a value into it. The
-- memory may be deallocated using 'free' when no longer required.
--
new :: Storable a => a -> IO (DevicePtr a)
new v =
    let bytes = fromIntegral (F.sizeOf v) in
    forceEither `fmap` malloc bytes >>= \dptr ->
    poke dptr v >>= \rv ->
    case rv of
        Nothing -> return dptr
        Just s  -> moduleErr "new" s

-- |
-- Write the list of storable elements into a newly allocated linear array on
-- the device. The memory may be deallocated when no longer required.
--
newArray   :: Storable a => [a] -> IO (DevicePtr a)
newArray v =
    let bytes = fromIntegral (length v) * fromIntegral (F.sizeOf (head v)) in
    forceEither  `fmap` malloc bytes >>= \dptr ->
    pokeArray dptr v >>= \rv ->
    case rv of
        Nothing -> return dptr
        Just s  -> moduleErr "newArray" s

-- |
-- Execute a computation passing a pointer to a temporarily allocated block of
-- device memory containing the value
--
with :: Storable a => a -> (DevicePtr a -> IO b) -> IO b
with v f =
    alloca $ \dptr ->
    poke dptr v >>= \rv ->
    case rv of
        Nothing -> f dptr
        Just s  -> moduleErr "with" s

-- |
-- Temporarily store a list of variable in device memory and operate on them
--
withArray :: Storable a => [a] -> (DevicePtr a -> IO b) -> IO b
withArray vs f = withArrayLen vs (\_ -> f)

-- |
-- Like 'withArray', but the action gets the number of values as an additional
-- parameter
--
withArrayLen :: Storable a => [a] -> (Int -> DevicePtr a -> IO b) -> IO b
withArrayLen vs f =
    let l = length vs
        b = fromIntegral l * fromIntegral (F.sizeOf (head vs))
    in
        allocaBytes b $ \dptr ->
        pokeArray dptr vs >>= \rv ->
        case rv of
            Nothing -> f l dptr
            Just s  -> moduleErr "withArray" s

--------------------------------------------------------------------------------
-- 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 :: Stream           -- ^ asynchronous stream to operate over
            -> Ptr a            -- ^ destination
            -> Ptr a            -- ^ source
            -> Int64            -- ^ number of bytes
            -> CopyDirection
            -> IO (Maybe String)
memcpyAsync stream dst src bytes dir =
    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,
    withDevicePtr,

    -- ** Dynamic allocation
    --
    -- Basic methods to allocate a block of memory on the device. The memory
    -- should be deallocated using 'free' when no longer necessary. The
    -- advantage is that failed allocations will not raise an exception.
    --
    free,
    malloc,
    malloc2D,
    malloc3D,
    memset,
    memset2D,
    memset3D,

    -- ** Local allocation
    --
    -- Execute a computation, passing as argument a pointer to a newly allocated
    -- block of memory. The memory is freed when the computation terminates.
    --
    alloca,
    allocaBytes,
    allocaBytesMemset,

    -- ** Marshalling
    --
    -- Store and retrieve Haskell lists as arrays on the device
    --
    peek,
    poke,
    peekArray,
    pokeArray,

    -- ** Combined allocation and marshalling
    --
    -- Write Haskell values (or lists) to newly allocated device memory. This
    -- may be a temporary store.
    --
    new,
    with,
    newArray,
    withArray,
    withArrayLen,

    -- ** Copying
    --
    -- Copy data between the host and device
    --
  )
  where

import Data.Int
import Control.Exception

import Foreign.C
import Foreign.Ptr
import Foreign.ForeignPtr
import Foreign.Storable (Storable)

import qualified Foreign.Storable as F
import qualified Foreign.Marshal  as F

import Foreign.CUDA.Utils
import Foreign.CUDA.Error
import Foreign.CUDA.Stream
import Foreign.CUDA.Internal.C2HS

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

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

-- |
-- A reference to data stored on the device
--
type DevicePtr a = ForeignPtr a

-- |
-- Unwrap the device pointer, yielding a device memory address that can be
-- passed to a kernel function callable from the host code. This allows the type
-- system to strictly separate host and device memory references, a non-obvious
-- distinction in pure CUDA.
--
withDevicePtr :: DevicePtr a -> (Ptr a -> IO b) -> IO b
withDevicePtr =  withForeignPtr

--
-- Internal, unwrap and cast. Required as some functions such as memset work
-- with untyped `void' data.
--
withDevicePtr' :: DevicePtr a -> (Ptr b -> IO c) -> IO c
withDevicePtr' =  withForeignPtr . castForeignPtr

newDevicePtr :: Ptr a -> IO (DevicePtr a)
newDevicePtr =  newForeignPtr_

--
-- 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 a))
malloc bytes =  do
    (rv,ptr) <- cudaMalloc bytes
    case rv of
        Success -> Right `fmap` newDevicePtr (castPtr ptr)
        _       -> return . Left . describe $ rv

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

-- |
-- 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 a, Int64))
malloc2D (width,height) =  do
    (rv,ptr,pitch) <- cudaMallocPitch width height
    case rv of
        Success -> (\p -> Right (p,pitch)) `fmap` newDevicePtr (castPtr ptr)
        _       -> return . Left . describe $ rv

{# fun unsafe cudaMallocPitch
    { alloca'-  `Ptr ()' peek'*       ,
      alloca'-  `Int64'  peekIntConv* ,
      cIntConv  `Int64'               ,
      cIntConv  `Int64'               } -> `Status' cToEnum #}
    where
        -- C->Haskell doesn't like qualified imports
        --
        alloca' :: Storable a => (Ptr a -> IO b) -> IO b
        alloca' =  F.alloca
        peek'   =  F.peek

-- |
-- 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 a, Int64))
malloc3D = moduleErr "malloc3D" "not implemented yet"

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

free_   :: Ptr a -> IO ()
free_ p =  cudaFree p >>= \rv ->
    case rv of
        Success -> return ()
        _       -> moduleErr "free" (describe rv)

{# fun unsafe cudaFree
    { castPtr `Ptr a' } -> `Status' cToEnum #}

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

-- |
-- Initialise device memory to a given value
--
memset :: DevicePtr a                   -- ^ The device memory
       -> Int64                         -- ^ Number of bytes
       -> Int                           -- ^ Value to set for each byte
       -> IO (Maybe String)
memset ptr bytes symbol =
    nothingIfOk `fmap` cudaMemset ptr symbol bytes

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

-- |
-- Initialise a matrix to a given value
--
memset2D :: DevicePtr a                 -- ^ 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 a' ,
      cIntConv        `Int64'       ,
                      `Int'         ,
      cIntConv        `Int64'       ,
      cIntConv        `Int64'       } -> `Status' cToEnum #}

-- |
-- Initialise the elements of a 3D array to a given value
--
memset3D :: DevicePtr a                 -- ^ 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"

--------------------------------------------------------------------------------
-- Local allocation
--------------------------------------------------------------------------------

-- |
-- Execute a computation, passing a pointer to a temporarily allocated block of
-- memory
--
alloca :: Storable a => (DevicePtr a -> IO b) -> IO b
alloca =  doAlloca undefined
  where
    doAlloca   :: Storable a' => a' -> (DevicePtr a' -> IO b') -> IO b'
    doAlloca x =  allocaBytes (fromIntegral (F.sizeOf x))

-- |
-- Execute a computation, passing a pointer to a block of 'n' bytes of memory
--
allocaBytes :: Int64 -> (DevicePtr a -> IO b) -> IO b
allocaBytes bytes =
    bracket (forceEither `fmap` malloc bytes) free

-- |
-- Execute a computation, passing a pointer to a  block of memory initialised to
-- a given value
--
allocaBytesMemset :: Int64 -> Int -> (DevicePtr a -> IO b) -> IO b
allocaBytesMemset bytes symbol f =
    allocaBytes bytes $ \dptr ->
    memset dptr bytes symbol >>= \rv ->
      case rv of
        Nothing -> f dptr
        Just e  -> error e

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

-- |
-- Retrieve the specified value from device memory
--
peek :: Storable a => DevicePtr a -> IO (Either String a)
peek d = doPeek undefined
  where
    doPeek   :: Storable a' => a' -> IO (Either String a')
    doPeek x =  F.alloca         $ \hptr ->
                withDevicePtr' d $ \dptr ->
                memcpy hptr dptr (fromIntegral (F.sizeOf x)) DeviceToHost >>= \rv ->
                case rv of
                    Nothing -> Right `fmap` F.peek hptr
                    Just s  -> return (Left s)

-- |
-- Retrieve the specified number of elements from device memory and return as a
-- Haskell list
--
peekArray :: Storable a => Int -> DevicePtr a -> IO (Either String [a])
peekArray n d = doPeek undefined
  where
    doPeek   :: Storable a' => a' -> IO (Either String [a'])
    doPeek x =  let bytes = fromIntegral n * (fromIntegral (F.sizeOf x)) in
                F.allocaArray  n $ \hptr ->
                withDevicePtr' d $ \dptr ->
                memcpy hptr dptr bytes DeviceToHost >>= \rv ->
                case rv of
                   Nothing -> Right `fmap` F.peekArray n hptr
                   Just s  -> return (Left s)

-- |
-- Store the given value to device memory
--
poke :: Storable a => DevicePtr a -> a -> IO (Maybe String)
poke d v =
    F.with v        $ \hptr ->
    withDevicePtr d $ \dptr ->
    memcpy dptr hptr (fromIntegral (F.sizeOf v)) HostToDevice

-- |
-- Store the list elements consecutively in device memory
--
pokeArray :: Storable a => DevicePtr a -> [a] -> IO (Maybe String)
pokeArray d = doPoke undefined
  where
    doPoke     :: Storable a' => a' -> [a'] -> IO (Maybe String)
    doPoke x v =  F.withArrayLen v $ \n hptr ->
                  withDevicePtr' d $ \dptr ->
                  let bytes = (fromIntegral n) * (fromIntegral (F.sizeOf x)) in
                  memcpy dptr hptr bytes HostToDevice

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

-- |
-- Allocate a block of memory on the device and transfer a value into it. The
-- memory may be deallocated using 'free' when no longer required.
--
new :: Storable a => a -> IO (DevicePtr a)
new v =
    let bytes = fromIntegral (F.sizeOf v) in
    forceEither `fmap` malloc bytes >>= \dptr ->
    poke dptr v >>= \rv ->
    case rv of
        Nothing -> return dptr
        Just s  -> moduleErr "new" s

-- |
-- Write the list of storable elements into a newly allocated linear array on
-- the device. The memory may be deallocated when no longer required.
--
newArray   :: Storable a => [a] -> IO (DevicePtr a)
newArray v =
    let bytes = fromIntegral (length v) * fromIntegral (F.sizeOf (head v)) in
    forceEither  `fmap` malloc bytes >>= \dptr ->
    pokeArray dptr v >>= \rv ->
    case rv of
        Nothing -> return dptr
        Just s  -> moduleErr "newArray" s

-- |
-- Execute a computation passing a pointer to a temporarily allocated block of
-- device memory containing the value
--
with :: Storable a => a -> (DevicePtr a -> IO b) -> IO b
with v f =
    alloca $ \dptr ->
    poke dptr v >>= \rv ->
    case rv of
        Nothing -> f dptr
        Just s  -> moduleErr "with" s

-- |
-- Temporarily store a list of variable in device memory and operate on them
--
withArray :: Storable a => [a] -> (DevicePtr a -> IO b) -> IO b
withArray vs f = withArrayLen vs (\_ -> f)

-- |
-- Like 'withArray', but the action gets the number of values as an additional
-- parameter
--
withArrayLen :: Storable a => [a] -> (Int -> DevicePtr a -> IO b) -> IO b
withArrayLen vs f =
    let l = length vs
        b = fromIntegral l * fromIntegral (F.sizeOf (head vs))
    in
        allocaBytes b $ \dptr ->
        pokeArray dptr vs >>= \rv ->
        case rv of
            Nothing -> f l dptr
            Just s  -> moduleErr "withArray" s

--------------------------------------------------------------------------------
-- 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 :: Stream           -- ^ asynchronous stream to operate over
            -> Ptr a            -- ^ destination
            -> Ptr a            -- ^ source
            -> Int64            -- ^ number of bytes
            -> CopyDirection
            -> IO (Maybe String)
memcpyAsync stream dst src bytes dir =
    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)