content: -- Copyright 2022 Google LLC
  --
  -- Use of this source code is governed by a BSD-style
  -- license that can be found in the LICENSE file or at
  -- https://developers.google.com/open-source/licenses/bsd
  
  module Runtime
    ( loadLitVal
    , LLVMCallable (..), callEntryFun, callDtor, BenchRequirement (..)
    , createTLSKey, PThreadKey) where
  
  import Data.Int
  import GHC.IO.FD
  import GHC.IO.Handle.FD
  import System.IO
  import System.Process
  
  import Foreign.Marshal.Alloc
  import Foreign.Ptr
  import Foreign.C.Types (CInt (..))
  import Foreign.Storable hiding (alignment)
  import Control.Monad.IO.Class
  import Control.Monad
  import Control.Concurrent
  import Control.Exception hiding (throw)
  import qualified Control.Exception as E
  
  import Err
  import MonadUtil
  import PPrint ()
  
  import Types.Top hiding (DexDestructor)
  import Types.Source  hiding (CInt)
  import Types.Primitives
  
  -- === One-shot evaluation ===
  
  foreign import ccall "dynamic"
    callFunPtr :: DexExecutable -> Int32 -> Ptr () -> Ptr () -> IO DexExitCode
  
  foreign import ccall "dynamic"
    callDtor :: DexDestructor -> IO ()
  
  type DexDestructor = FunPtr (IO ())
  type DexExecutable = FunPtr (Int32 -> Ptr () -> Ptr () -> IO DexExitCode)
  type DexExitCode = Int
  
  
  data BenchRequirement =
      NoBench
    | DoBench Bool -- True means "must sync CUDA"
  
  data LLVMCallable = LLVMCallable
    { nativeFun :: NativeFunction
    , logger :: PassLogger
    , resultTypes :: [BaseType] }
  
  -- The NativeFunction needs to have been compiled with EntryFunCC.
  callEntryFun :: LLVMCallable -> [LitVal] -> IO [LitVal]
  callEntryFun LLVMCallable{nativeFun, logger, resultTypes} args = do
    withPipeToLogger logger \fd ->
      allocaCells (length args) \argsPtr ->
        allocaCells (length resultTypes) \resultPtr -> do
          storeLitVals argsPtr args
          let fPtr = castFunPtr $ nativeFunPtr nativeFun
          checkedCallFunPtr fd argsPtr resultPtr fPtr
          results <- loadLitVals resultPtr resultTypes
          return results
  {-# SCC callEntryFun #-}
  
  checkedCallFunPtr :: FD -> Ptr () -> Ptr () -> DexExecutable -> IO ()
  checkedCallFunPtr fd argsPtr resultPtr fPtr = do
    let (CInt fd') = fdFD fd
    exitCode <- callFunPtr fPtr fd' argsPtr resultPtr
    unless (exitCode == 0) $ throwErr RuntimeErr
  
  withPipeToLogger :: PassLogger -> (FD -> IO a) -> IO a
  withPipeToLogger logger writeAction = do
    result <- snd <$> withPipe
      (\h -> readStream h \s -> ioLogAction logger (Outputs [TextOut s]))
      (\h -> handleToFd h >>= writeAction)
    case result of
      Left e -> E.throw e
      Right ans -> return ans
  
  -- === serializing scalars ===
  
  loadLitVals :: MonadIO m => Ptr () -> [BaseType] -> m [LitVal]
  loadLitVals p types = zipWithM loadLitVal (ptrArray p) types
  
  _freeLitVals :: MonadIO m => Ptr () -> [BaseType] -> m ()
  _freeLitVals p types = zipWithM_ freeLitVal (ptrArray p) types
  
  storeLitVals :: MonadIO m => Ptr () -> [LitVal] -> m ()
  storeLitVals p xs = zipWithM_ storeLitVal (ptrArray p) xs
  
  loadLitVal :: MonadIO m => Ptr () -> BaseType -> m LitVal
  loadLitVal ptr (Scalar ty) = liftIO case ty of
    Int64Type   -> Int64Lit   <$> peek (castPtr ptr)
    Int32Type   -> Int32Lit   <$> peek (castPtr ptr)
    Word8Type   -> Word8Lit   <$> peek (castPtr ptr)
    Word32Type  -> Word32Lit  <$> peek (castPtr ptr)
    Word64Type  -> Word64Lit  <$> peek (castPtr ptr)
    Float64Type -> Float64Lit <$> peek (castPtr ptr)
    Float32Type -> Float32Lit <$> peek (castPtr ptr)
  loadLitVal ptrPtr (PtrType t) = do
    ptr <- liftIO $ peek $ castPtr ptrPtr
    PtrLit t <$> if ptr == nullPtr
      then return NullPtr
      else return $ PtrLitVal ptr
  loadLitVal _ (Vector _ _) = error "Vector loads not implemented"
  
  storeLitVal :: MonadIO m => Ptr () -> LitVal -> m ()
  storeLitVal ptr val = liftIO case val of
    Int64Lit   x -> poke (castPtr ptr) x
    Int32Lit   x -> poke (castPtr ptr) x
    Word8Lit   x -> poke (castPtr ptr) x
    Word32Lit  x -> poke (castPtr ptr) x
    Word64Lit  x -> poke (castPtr ptr) x
    Float64Lit x -> poke (castPtr ptr) x
    Float32Lit x -> poke (castPtr ptr) x
    PtrLit _ (PtrLitVal x) -> poke (castPtr ptr) x
    _ -> error "not implemented"
  
  foreign import ccall "free_dex"
    free_cpu :: Ptr () -> IO ()
  #ifdef DEX_CUDA
  foreign import ccall "dex_cuMemFree"
    free_gpu :: Ptr () -> IO ()
  #else
  free_gpu :: Ptr () -> IO ()
  free_gpu = error "Compiled without GPU support!"
  #endif
  
  freeLitVal :: MonadIO m => Ptr () -> BaseType -> m ()
  freeLitVal litValPtr ty = case ty of
    Scalar  _ -> return ()
    PtrType (CPU, Scalar _) -> liftIO $ free_cpu =<< loadPtr
    PtrType (GPU, Scalar _) -> liftIO $ free_gpu =<< loadPtr
    -- TODO: Handle pointers to pointers
    _ -> error "not implemented"
    where loadPtr = peek (castPtr litValPtr)
  
  cellSize :: Int
  cellSize = 8
  
  ptrArray :: Ptr () -> [Ptr ()]
  ptrArray p = map (\i -> p `plusPtr` (i * cellSize)) [0..]
  
  allocaCells :: MonadIO m => Int -> (Ptr () -> IO a) -> m a
  allocaCells n cont = liftIO $ allocaBytes (n * cellSize) cont
  
  -- ==== unix pipe utilities ===
  
  type IOExcept a = Either SomeException a
  
  withPipe :: (Handle -> IO a) -> (Handle -> IO b) -> IO (IOExcept a, IOExcept b)
  withPipe readAction writeAction = do
    (readHandle, writeHandle) <- createPipe
    waitForReader <- forkWithResult $ readAction  readHandle
    waitForWriter <- forkWithResult $ writeAction writeHandle
    y <- waitForWriter `finally` hClose writeHandle
    x <- waitForReader `finally` hClose readHandle
    return (x, y)
  
  forkWithResult :: IO a -> IO (IO (IOExcept a))
  forkWithResult action = do
    resultMVar <- newEmptyMVar
    void $ forkIO $ catch (do result <- action
                              putMVar resultMVar $ Right result)
                          (\e -> putMVar resultMVar $ Left (e::SomeException))
    return $ takeMVar resultMVar
  
  readStream :: Handle -> (String -> IO ()) -> IO ()
  readStream h action = go
    where
      go :: IO ()
      go = do
        eof <- hIsEOF h
        unless eof $ hGetLine h >>= action >> go
  
  -- ==== thread-local storage ===
  
  -- We use TLS to implement some dynamically scoped variables, like the file
  -- descriptor for the current output stream. We create the key here and we
  -- generate code for `pthread_setspecific` at entry functions and
  -- `pthread_getspecific` at use sites.
  
  data PThreadKey
  
  foreign import ccall "dex_pthread_key_create"
    pthreadKeyCreate :: IO (Ptr PThreadKey)
  
  createTLSKey :: IO (Ptr (Ptr ()))
  createTLSKey = castPtr <$> pthreadKeyCreate
  