-- 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 LLVM.Link
  ( createLinker, destroyLinker
  , addExplicitLinkMap, addObjectFile, getFunctionPointer
  , ExplicitLinkMap
  ) where

import Data.String (fromString)
import Foreign.Ptr
import qualified Data.ByteString as BS

import System.IO
import System.IO.Temp

import qualified LLVM.OrcJIT as OrcJIT
import qualified LLVM.Internal.OrcJIT as OrcJIT
import qualified LLVM.Internal.Target as Target
import qualified LLVM.Internal.FFI.Target as FFI

import qualified LLVM.Shims

data Linker = Linker
  { linkerExecutionSession :: OrcJIT.ExecutionSession
#ifdef darwin_HOST_OS
  , linkerLinkLayer        :: OrcJIT.ObjectLinkingLayer
#else
  , linkerLinkLayer        :: OrcJIT.RTDyldObjectLinkingLayer
#endif
  , _linkerTargetMachine   :: Target.TargetMachine
   -- We ought to just need the link layer and the mangler but but llvm-hs
   -- requires a full `IRCompileLayer` for incidental reasons. TODO: fix.
  , linkerIRLayer          :: OrcJIT.IRCompileLayer
  , linkerDylib            :: OrcJIT.JITDylib }

instance OrcJIT.IRLayer Linker where
  -- llvm-hs requires an compile/IR layer but don't actually need it for the
  -- linking functions we call. TODO: update llvm-hs to expose more precise
  -- requirements for its linking functions.
  getIRLayer    l = OrcJIT.getIRLayer    $ linkerIRLayer l
  getDataLayout l = OrcJIT.getDataLayout $ linkerIRLayer l
  getMangler    l = OrcJIT.getMangler    $ linkerIRLayer l

type CName = String

type ExplicitLinkMap = [(CName, Ptr ())]

createLinker :: IO Linker
createLinker = do
  -- TODO: should this be a parameter to `createLinker` instead?
  tm <- LLVM.Shims.newDefaultHostTargetMachine
  s         <- OrcJIT.createExecutionSession
#ifdef darwin_HOST_OS
  linkLayer <- OrcJIT.createObjectLinkingLayer s
#else
  linkLayer <- OrcJIT.createRTDyldObjectLinkingLayer s
#endif
  dylib     <- OrcJIT.createJITDylib s "main_dylib"
  compileLayer <- OrcJIT.createIRCompileLayer s linkLayer tm
  OrcJIT.addDynamicLibrarySearchGeneratorForCurrentProcess compileLayer dylib
  return $ Linker s linkLayer tm compileLayer dylib

destroyLinker :: Linker  -> IO ()
destroyLinker (Linker session _ (Target.TargetMachine tm) _ _) = do
  -- dylib, link layer and IRLayer should get cleaned up automatically
  OrcJIT.disposeExecutionSession session
  FFI.disposeTargetMachine tm

addExplicitLinkMap :: Linker -> ExplicitLinkMap -> IO ()
addExplicitLinkMap l linkMap = do
  let (linkedNames, linkedPtrs) = unzip linkMap
  let flags = OrcJIT.defaultJITSymbolFlags { OrcJIT.jitSymbolAbsolute = True }
  let ptrSymbols = [OrcJIT.JITSymbol (ptrToWordPtr ptr) flags | ptr <- linkedPtrs]
  mangledNames <- mapM (OrcJIT.mangleSymbol l . fromString) linkedNames
  OrcJIT.defineAbsoluteSymbols (linkerDylib l) $ zip mangledNames ptrSymbols
  mapM_ OrcJIT.disposeMangledSymbol mangledNames

addObjectFile :: Linker -> BS.ByteString -> IO ()
addObjectFile l objFileContents = do
  withSystemTempFile "objfile.o" \path h -> do
    BS.hPut h objFileContents
    hFlush h
    OrcJIT.addObjectFile (linkerLinkLayer l) (linkerDylib l) path

getFunctionPointer :: Linker -> CName -> IO (FunPtr a)
getFunctionPointer l name = do
  OrcJIT.lookupSymbol (linkerExecutionSession l) (linkerIRLayer l)
    (linkerDylib l) (fromString name) >>= \case
      Right (OrcJIT.JITSymbol funcAddr _) ->
        return $ castPtrToFunPtr $ wordPtrToPtr funcAddr
      Left s -> error $ "Couldn't find function: " ++ name ++ "\n" ++ show s