-- Copyright 2006-2008, Galois, Inc.
-- This software is distributed under a standard, three-clause BSD license.
-- Please see the file LICENSE, distributed with this software, for specific
-- terms and conditions.
-- |Support for inter-domain communication through typed
-- communication channels.  The module provides types that represent
-- the input and ouput ends of open unidirectional communication
-- channels.  These channels are parameterized over the types of
-- messages that can be sent over them, so that domains can exchange
-- messages in a type-safe manner.
--
-- There are also bidirectional channels, parameterized over the types
-- of messages in each direction.
--
module Communication.IVC(
         InChannel, OutChannel, InOutChannel
       , ReadableChan, WriteableChan
       , makeNewInChannel, acceptNewInChannel
       , makeNewOutChannel, acceptNewOutChannel
       , makeNewInOutChannel, acceptNewInOutChannel
       -- , closeInChannel, closeOutChannel, closeInOutChannel
       , get, put, peer
       )
 where

import Communication.Rendezvous
import Control.Concurrent.MVar
import Control.Exception
import Control.Monad
import Data.Binary hiding (get,put)
import Data.Binary.Get(runGet, getWordhost)
import Data.Binary.Put(runPut, putWordhost, putLazyByteString)
import qualified Data.ByteString as BSS
import Data.ByteString.Lazy(ByteString)
import qualified Data.ByteString.Lazy as BS
import Foreign.Marshal.Alloc
import Foreign.Ptr
import Foreign.Storable
import Hypervisor.DomainInfo
import Hypervisor.ErrorCodes
import Hypervisor.Memory
import Hypervisor.Port

data IVCSetup
  = Lend DomId [GrantRef] Port
  | Borrow GrantHandle

data InChannel a = InChannel {
    ichSetupData :: IVCSetup
  , ichInChannel :: InChan
  , ichPeer      :: DomId
  }

-- |Make a new input channel, targetting the given domain. The second argument
-- is the number of pages to use for the channel. (Note: the actual size of the
-- transfer buffer in memory will be slightly smaller than n * pageSize, because
-- some bookkeeping space is required)
makeNewInChannel :: Binary a => DomId -> Word -> IO (InChannel a)
makeNewInChannel target npages = do
  (grefs, port, ichn) <- makeNewChan target npages buildRawInChan
  return (InChannel (Lend target grefs port) ichn target)

-- |Accept a new input channel, given the input data.
acceptNewInChannel :: Binary a =>
                      DomId -> [GrantRef] -> Port ->
                      IO (InChannel a)
acceptNewInChannel target grants port = do
  (ichn, gh) <- acceptNewChan target grants port buildRawInChan
  return (InChannel (Borrow gh) ichn target)

{- These close functions seem correct, but cause occasional access violations.
 -
-- |Ungracefully revoke or unmap the channel's grants.
-- This may cause a page fault if used without coordination.
-- throws @ErrorCode@ in case of failure.
closeInChannel :: Binary a => InChannel a -> IO ()
closeInChannel ich = case ichSetupData ich of
  Lend _ grants _ -> mapM_ endAccess grants
  Borrow gh       -> unmapGrant gh 
-}

data OutChannel a = OutChannel {
    ochSetupData  :: IVCSetup
  , ochOutChannel :: OutChan
  , ochPeer       :: DomId
  }

-- |Make a new output channel, targetting the given domain. The second argument
-- is the number of pages to use for the channel. (Note: the actual size of the
-- transfer buffer in memory will be slightly smaller than n * pageSize, because
-- some bookkeeping space is required)
makeNewOutChannel :: Binary a =>
                     DomId -> Word ->
                     IO (OutChannel a)
makeNewOutChannel target npages = do
  (grefs, port, ochn) <- makeNewChan target npages buildRawOutChan
  return (OutChannel (Lend target grefs port) ochn target)

-- |Accept a new output channel, given the input data
acceptNewOutChannel :: Binary a =>
                       DomId -> [GrantRef] -> Port ->
                       IO (OutChannel a)
acceptNewOutChannel target grants port = do
  (ochn, gh) <- acceptNewChan target grants port buildRawOutChan
  return (OutChannel (Borrow gh) ochn target)

data InOutChannel a b = InOutChannel {
    bchSetupData  :: IVCSetup
  , bchInPercent  :: Float
  , bchInChannel  :: InChan
  , bchOutChannel :: OutChan
  , bchPeer       :: DomId
  }

{-
closeOutChannel :: Binary a => OutChannel a -> IO ()
closeOutChannel och = case ochSetupData och of
  Lend _ grants _ -> mapM_ endAccess grants
  Borrow gh       -> unmapGrant gh 
-}

-- |Make a new input / output channel targetting the given domain. The second
-- argument is the number of pages to use, while the third argument tells the
-- system what percentage to use for the input channel. This third argument
-- must be between 0 and 1, inclusive.
makeNewInOutChannel :: (Binary a, Binary b) =>
                       DomId -> Word -> Float ->
                       IO (InOutChannel a b)
makeNewInOutChannel target npages perc
  | (perc < 0) || (perc > 1.0) = throwIO EINVAL
  | otherwise                  = do
     (grs, p, (ich,och)) <- makeNewChan target npages (buildIOChan perc npages)
     return (InOutChannel (Lend target grs p) perc ich och target)

-- |Accept a new input / out channel, given the input data
acceptNewInOutChannel :: (Binary a, Binary b) =>
                         DomId -> [GrantRef] -> Port -> Float ->
                         IO (InOutChannel a b)
acceptNewInOutChannel target grants port perc
  | (perc < 0) || (perc > 1.0) = throwIO EINVAL
  | otherwise                  = do
     let npages = fromIntegral (length grants)
     ((i, o), gh) <- acceptNewChan target grants port (buildIOChan perc npages)
     return (InOutChannel (Borrow gh) perc i o target)

{-
closeInOutChannel :: (Binary a, Binary b) => InOutChannel a b -> IO ()
closeInOutChannel bch = case bchSetupData bch of
  Lend _ grants _ -> mapM_ endAccess grants
  Borrow gh       -> unmapGrant gh 
-}

buildIOChan :: Float -> Word ->
               Bool -> Ptr Word8 -> Word -> Port ->
               IO (InChan, OutChan)
buildIOChan perc npages doClear ptr _ port = do
  let p1Size   = floor ((fromIntegral (npages * 4096)) * perc)
      p2Size   = (npages * 4096) - p1Size
      b1Size   = p1Size - bookkeepingOverhead
      b2Size   = p2Size - bookkeepingOverhead
  let (inPtr, inSize, outPtr, outSize) =
        if doClear
          then (ptr, b1Size, ptr `plusPtrW` p1Size, b2Size)
          else (ptr `plusPtrW` p1Size, b2Size, ptr, b1Size)
  ichn <- buildRawInChan  doClear inPtr  inSize  port
  ochn <- buildRawOutChan doClear outPtr outSize port
  setPortHandler port $ tryWriteData ochn >> tryReadData  ichn
  return (ichn, ochn)

makeNewChan :: DomId -> Word ->
               (Bool -> Ptr Word8 -> Word -> Port -> IO a) ->
               IO ([GrantRef], Port, a)
makeNewChan target npages buildChan = do
  let len = npages * 4096
  ptr  <- mallocBytes (fromIntegral len)
  refs <- grantAccess target ptr len True
  port <- allocPort target
  ichn <- buildChan True ptr (len - bookkeepingOverhead) port
  return (refs, port, ichn)

acceptNewChan :: DomId -> [GrantRef] -> Port ->
                 (Bool -> Ptr Word8 -> Word -> Port -> IO a) ->
                 IO (a, GrantHandle)
acceptNewChan target grefs port buildChan = do
  myport <- bindRemotePort target port
  (ptr, gh) <- mapGrants target grefs True
  let size = (length grefs * 4096) - bookkeepingOverhead
  chan <- buildChan False ptr (fromIntegral size) myport
  return (chan, gh)

-- -----------------------------------------------------------------------------

instance Binary a => RendezvousCapable Word (InChannel a) (OutChannel a) where
  makeConnection other size = do
    res <- makeNewOutChannel other size
    let Lend _ grs ps = ochSetupData res
    return (grs, [ps], return res)
  acceptConnection other refs [port] _ = acceptNewInChannel other refs port
  acceptConnection _ _ _ _ = fail "Should only have received one port!"

instance Binary a => RendezvousCapable Word (OutChannel a) (InChannel a) where
  makeConnection other size = do
    res <- makeNewInChannel other size
    let Lend _ grs ps = ichSetupData res
    return (grs, [ps], return res)
  acceptConnection other refs [port] _ = acceptNewOutChannel other refs port
  acceptConnection _ _ _ _ = fail "Should only have received one port!"

instance (Binary a, Binary b) =>
           RendezvousCapable (Float, Word) (InOutChannel a b) (InOutChannel b a)
 where
  makeConnection other (perc, size) = do
    res <- makeNewInOutChannel other size perc
    let Lend _ grs ps = bchSetupData res
    return (grs, [ps], return res)
  acceptConnection other refs [port] (perc, _) =
    acceptNewInOutChannel other refs port perc
  acceptConnection _ _ _ _ =
    fail "Should only have received one port!"

-- -----------------------------------------------------------------------------

class WriteableChan c a | c -> a where
  put :: c -> a -> IO ()

instance Binary a => WriteableChan (OutChannel a) a where
  put c = putBinary (ochOutChannel c)

instance Binary b => WriteableChan (InOutChannel a b) b where
  put c = putBinary (bchOutChannel c)

class ReadableChan c a | c -> a where
  get :: c -> IO a

instance Binary a => ReadableChan (InChannel a) a where
  get c = getBinary (ichInChannel c)

instance Binary a => ReadableChan (InOutChannel a b) a where
  get c = getBinary (bchInChannel c)

putBinary :: Binary a => OutChan -> a -> IO ()
putBinary oc x = runWriteRequest oc (encode x)

getBinary :: Binary a => InChan -> IO a
getBinary ic = decode `fmap` runReadRequest ic

class CommChan c where
  peer :: c -> DomId

instance CommChan (InChannel a) where
  peer = ichPeer

instance CommChan (OutChannel a) where
  peer = ochPeer

instance CommChan (InOutChannel a b) where
  peer = bchPeer


-- -----------------------------------------------------------------------------

--
-- A communications channel is composed of something of a pair of a pointer
-- and a size, where:
--
--  +-----------------------+ ptr + 0
--  +         ...           |
--  +         ...           |
--  +     buffer space      |
--  +         ...           |
--  +         ...           |
--  +-----------------------+ ptr + size
--  +    bytes consumed     |
--  +-----------------------+ ptr + size + 4
--  +    bytes produced     |
--  +-----------------------+ ptr + size + 8
--

bytesConsumed :: Ptr Word8 -> Word -> IO Word32
bytesConsumed p s = peekByteOff (castPtr p) (fromIntegral s)

bytesProduced :: Ptr Word8 -> Word -> IO Word32
bytesProduced p s = peekByteOff (castPtr p) (fromIntegral s + 4)

setBytesConsumed :: Ptr Word8 -> Word -> Word32 -> IO ()
setBytesConsumed p s v = pokeByteOff (castPtr p) (fromIntegral s) v

setBytesProduced :: Ptr Word8 -> Word -> Word32 -> IO ()
setBytesProduced p s v = pokeByteOff (castPtr p) (fromIntegral s + 4) v

bookkeepingOverhead :: Integral a => a
bookkeepingOverhead = 8

-- Internal-only data structure
data OutChan = OutChan {
    ocBuffer  :: Ptr Word8
  , ocSize    :: Word
  , ocModulus :: Word32
  , ocPort    :: Port
  , ocWaiting :: MVar [(ByteString, MVar ())]
  }

buildRawOutChan :: Bool -> Ptr Word8 -> Word -> Port -> IO OutChan
buildRawOutChan doClear buf size port = do
  when doClear $ bzero buf size
  waiters <- newMVar []
  let res = OutChan buf size (computeModulus size) port waiters
  setPortHandler port $ tryWriteData res
  return res

runWriteRequest :: OutChan -> ByteString -> IO ()
runWriteRequest och !bs = do
  resMV <- newEmptyMVar
  waiters <- takeMVar (ocWaiting och)
  putMVar (ocWaiting och) $! (msg, resMV) : waiters
  tryWriteData och
  takeMVar resMV
 where
  !msg = runPut $ do
     putWordhost (fromIntegral (BS.length bs))
     putLazyByteString bs

tryWriteData :: OutChan -> IO ()
tryWriteData och = do
  waiters           <- takeMVar (ocWaiting och)
  cons              <- bytesConsumed (ocBuffer och) (ocSize och)
  prod              <- bytesProduced (ocBuffer och) (ocSize och)
  (waiters', prod') <- doPossibleWrites prod cons waiters
  setBytesProduced (ocBuffer och) (ocSize och) prod'
  when (prod /= prod') $ sendOnPort (ocPort och)
  putMVar (ocWaiting och) $! waiters'
 where
  bufferSize = fromIntegral (ocSize och)
  --
  doPossibleWrites :: Word32 -> Word32 ->
                      [(ByteString, MVar())] ->
                      IO ([(ByteString, MVar())], Word32)
  doPossibleWrites prod _ [] = return ([], prod)
  doPossibleWrites prod cons ls@((bstr, resMV):rest) = do
    -- this is an awkward way to deal with rollver, but it should work.
    let unread = if prod >= cons then prod - cons else overflow
        overflow = prod + (ocModulus och - cons)
        avail = bufferSize - unread
        bstrLn  = fromIntegral (BS.length bstr)
      -- In this case, the buffer is full.
    if | avail == 0 -> return (ls, prod)
      -- In this case, we have enough space to write the full bytestring.
       | avail > bstrLn -> do
           writeBS (ocBuffer och) (ocSize och) prod bstr
           putMVar resMV ()
           let prod' = (prod + fromIntegral bstrLn) `mod` ocModulus och
           doPossibleWrites prod' cons rest
      -- In this case, we have space to do a write, but not the whole
      -- bytestring
       | otherwise -> do
           let (h,t)   = BS.splitAt (fromIntegral avail) bstr
           writeBS (ocBuffer och) (ocSize och) prod h
           let prod' = fromIntegral (prod + avail) `mod` ocModulus och
           return ((t, resMV) : rest, prod')

writeBS :: Ptr Word8 -> Word -> Word32 -> ByteString -> IO ()
writeBS buffer size logical_off lbstr =
  foldM_ doWrite logical_off (BS.toChunks lbstr)
 where
  doWrite :: Word32 -> BSS.ByteString -> IO Word32
  doWrite loff bstr = BSS.useAsCStringLen bstr $ \ (dptr, dlenI) -> do
    let real_off = fromIntegral (loff `mod` fromIntegral size)
        destPtr  = buffer `plusPtrW` real_off
        dlen     = fromIntegral dlenI
    if real_off + dlen > size
      then do let part1s = size - real_off
                  part2s = dlen - part1s
              memcpy destPtr dptr                     part1s
              memcpy buffer  (dptr `plusPtrW` part1s) part2s
      else    memcpy destPtr dptr dlen
    return (loff + fromIntegral dlen)

-- Internal-only data structure
data InChan = InChan {
    icBuffer    :: Ptr Word8
  , icSize      :: Word
  , icModulus   :: Word32
  , icPort      :: Port
  , icStateMV   :: MVar InChanState
  }

data InChanState = NeedSize [MVar ByteString]
                 | GotSize !Word32 ByteString [MVar ByteString]

buildRawInChan :: Bool -> Ptr Word8 -> Word -> Port -> IO InChan
buildRawInChan doClear buf size port = do
  when doClear $ bzero buf size
  stateMV <- newMVar (NeedSize [])
  let res = InChan buf size (computeModulus size) port stateMV
  setPortHandler port $ tryReadData res
  return res

runReadRequest :: InChan -> IO ByteString
runReadRequest ich = do
  resMV <- newEmptyMVar
  istate <- takeMVar (icStateMV ich)
  case istate of
    NeedSize waiters ->
      putMVar (icStateMV ich) $! NeedSize (waiters ++ [resMV])
    GotSize n acc waiters ->
      putMVar (icStateMV ich) $! GotSize n acc (waiters ++ [resMV])
  tryReadData ich
  takeMVar resMV

tryReadData :: InChan -> IO ()
tryReadData ich = modifyMVar_ (icStateMV ich) $ \ istate -> do
  prod             <- bytesProduced (icBuffer ich) (icSize ich)
  cons             <- bytesConsumed (icBuffer ich) (icSize ich)
  (istate', cons') <- doPossibleReads prod cons istate
  setBytesConsumed (icBuffer ich) (icSize ich) cons'
  when (cons /= cons') $ sendOnPort (icPort ich)
  return istate'
 where
  doPossibleReads :: Word32 -> Word32 -> InChanState -> IO (InChanState, Word32)
  doPossibleReads prod cons istate = do
    let avail = if prod >= cons then prod - cons else overflow
        overflow = prod + (icModulus ich - cons)
    case istate of
      -- If we need to get a size, we have waiters, and there's at least
      -- four bytes to read, then we should read off the size.
      NeedSize ws@(_:_) | avail >= sizeSize -> do
        sizeBS <- readBS (icBuffer ich) (icSize ich) cons sizeSize
        let size = runGet getWordhost sizeBS
        let istate' = GotSize (fromIntegral size) BS.empty ws
            cons'   = (cons + sizeSize) `mod` icModulus ich
        doPossibleReads prod cons' istate'
      -- If we have some data, but not enough, update ourselves with the
      -- new data and the lesser requirement.
      GotSize n acc ws | (avail > 0) && (n > avail) -> do
        part <- readBS (icBuffer ich) (icSize ich) cons avail
        let istate' = GotSize (n - avail) (acc `BS.append` part) ws
            cons'   = (cons + avail) `mod` icModulus ich
        doPossibleReads prod cons' istate'
      -- If we can read everything, do it!
      GotSize n acc (f:rest) | (avail > 0) && (n <= avail) -> do
        endp <- readBS (icBuffer ich) (icSize ich) cons n
        putMVar f (acc `BS.append` endp)
        let cons' = (cons + n) `mod` icModulus ich
        doPossibleReads prod cons' (NeedSize rest)
      -- Otherwise, we can't do anything
      _ ->
        return (istate, cons)

readBS :: Ptr Word8 -> Word -> Word32 -> Word32 -> IO ByteString
readBS !buffer !sizeW !logical_off !amt = do
  let real_off = logical_off `mod` size
      readPtr  = buffer `plusPtrW` real_off
      part1sz  = size - real_off
      part2sz  = amt - part1sz
  if real_off + amt > size
    then do part1 <- packCStringLen readPtr part1sz
            part2 <- packCStringLen buffer  part2sz
            return $! BS.fromStrict part1 `BS.append` BS.fromStrict part2
    else    BS.fromStrict `fmap` packCStringLen readPtr amt
 where
  size = fromIntegral sizeW
  packCStringLen p s = BSS.packCStringLen (castPtr p, fromIntegral s)

plusPtrW :: Integral b => Ptr a -> b -> Ptr a
plusPtrW p x = p `plusPtr` (fromIntegral x)

sizeSize :: Integral a => a
sizeSize = fromIntegral (BS.length (runPut (putWordhost 0)))

computeModulus :: Word -> Word32
computeModulus size
  | base == 0 = fromIntegral q * (fromIntegral size - 1)
  | otherwise = base
  where
   base  = fromIntegral q * fromIntegral size
   size' = fromIntegral size :: Word64
   q     = 0x100000000 `div` size'

foreign import ccall unsafe "strings.h bzero"
  bzero :: Ptr a -> Word -> IO ()

foreign import ccall unsafe "string.h memcpy"
  memcpy :: Ptr a -> Ptr b -> Word -> IO ()
