Add withNewClient.

Add withNewClient.

withNewClient :: RpcReply (Client req resp) -> Rpc req resp a -> IO (RpcReply a)

-- 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
  putStrLn "clientWaitForInitialMetadata(..)"
  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
          putStrLn ("metadata: " ++ show md)
          writeIORef initialMDRef (Just md)
          return (RpcOk md)
        RpcError err -> do
          putStrLn "clientWaitForInitialMetadata: callBatch failed"
          return (RpcError err)

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

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

recvMessage :: ClientReaderWriter -> IO (RpcReply (Maybe L.ByteString))
recvMessage crw = do
  putStrLn "recvMessage(..)"
  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 ]

clientSendClose :: ClientReaderWriter -> IO (RpcReply ())
clientSendClose 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)

sendClose :: Rpc req resp ()
sendClose = do
  crw <- askCrw
  joinReply =<< liftIO (clientSendClose 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

-- 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
  putStrLn "clientWaitForInitialMetadata(..)"
  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
          putStrLn ("metadata: " ++ show md)
          writeIORef initialMDRef (Just md)
          return (RpcOk md)
        RpcError err -> do
          putStrLn "clientWaitForInitialMetadata: callBatch failed"
          return (RpcError err)

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

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

recvMessage :: ClientReaderWriter -> IO (RpcReply (Maybe L.ByteString))
recvMessage crw = do
  putStrLn "recvMessage(..)"
  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 ]

clientSendClose :: ClientReaderWriter -> IO (RpcReply ())
clientSendClose 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)

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)

sendClose :: Rpc req resp ()
sendClose = do
  crw <- askCrw
  joinReply =<< liftIO (clientSendClose 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


Haddock wibble

Haddock wibble

{-# LANGUAGE ForeignFunctionInterface #-}
--------------------------------------------------------------------------------
-- |
-- Module    : Foreign.CUDA.Driver.Utils
-- Copyright : [2009..2014] Trevor L. McDonell
-- License   : BSD
--
-- Utility functions
--
--------------------------------------------------------------------------------

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

#include "cbits/stubs.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.
--
{-# INLINEABLE driverVersion #-}
driverVersion :: IO Int
driverVersion =  resultIfOk =<< cuDriverGetVersion

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

{-# LANGUAGE ForeignFunctionInterface #-}
--------------------------------------------------------------------------------
-- |
-- Module    : Foreign.CUDA.Driver.Utils
-- Copyright : [2009..2014] Trevor L. McDonell
-- License   : BSD
--
-- Utility functions
--
--------------------------------------------------------------------------------

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

#include "cbits/stubs.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
--
{-# INLINEABLE driverVersion #-}
driverVersion :: IO Int
driverVersion =  resultIfOk =<< cuDriverGetVersion

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


slightly better error messages

slightly better error messages

Ignore-this: 3f84db1930247040095d1aba6715fe3b

darcs-hash:20090928090016-115f9-f16170f2165b30f45b405a30fd9d6db78fb9d4bb.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,
    withDevicePtrByteOff,

    -- ** 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,
    peekArrayElemOff,
    pokeArray,
    pokeArrayElemOff,

    -- ** 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

-- |
-- Similar to `withDevicePtr', but with the address advanced by the given offset
-- in bytes.
--
withDevicePtrByteOff :: DevicePtr a -> Int -> (Ptr a -> IO b) -> IO b
withDevicePtrByteOff dptr offset action =
    withForeignPtr dptr $ \p -> action (p `plusPtr` offset)

--
-- 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 (moduleForceEither `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 =  peekArrayElemOff 0

-- |
-- Retrieve the specified number of elements from device memory, beginning at
-- the specified index (from zero), and return as a Haskell list.
--
peekArrayElemOff :: Storable a => Int -> Int -> DevicePtr a -> IO (Either String [a])
peekArrayElemOff o 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 `plusPtr` (o * F.sizeOf x)) 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 = pokeArrayElemOff 0

-- |
-- Store the list elements consecutively in device memory, beginning at the
-- specified index (from zero)
--
pokeArrayElemOff :: Storable a => Int -> DevicePtr a -> [a] -> IO (Maybe String)
pokeArrayElemOff o 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 `plusPtr` (o * F.sizeOf x)) 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
    moduleForceEither `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
    moduleForceEither `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)

moduleForceEither :: Either String a -> a
moduleForceEither (Left  s) = moduleErr "Marshal" s
moduleForceEither (Right r) = r

{-# 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,
    withDevicePtrByteOff,

    -- ** 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,
    peekArrayElemOff,
    pokeArray,
    pokeArrayElemOff,

    -- ** 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

-- |
-- Similar to `withDevicePtr', but with the address advanced by the given offset
-- in bytes.
--
withDevicePtrByteOff :: DevicePtr a -> Int -> (Ptr a -> IO b) -> IO b
withDevicePtrByteOff dptr offset action =
    withForeignPtr dptr $ \p -> action (p `plusPtr` offset)

--
-- 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 =  peekArrayElemOff 0

-- |
-- Retrieve the specified number of elements from device memory, beginning at
-- the specified index (from zero), and return as a Haskell list.
--
peekArrayElemOff :: Storable a => Int -> Int -> DevicePtr a -> IO (Either String [a])
peekArrayElemOff o 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 `plusPtr` (o * F.sizeOf x)) 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 = pokeArrayElemOff 0

-- |
-- Store the list elements consecutively in device memory, beginning at the
-- specified index (from zero)
--
pokeArrayElemOff :: Storable a => Int -> DevicePtr a -> [a] -> IO (Maybe String)
pokeArrayElemOff o 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 `plusPtr` (o * F.sizeOf x)) 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)