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 LLVM.Compile
    ( compileLLVM, runPasses, standardCompilationPipeline
    , LLVMOptLevel (..)) where
  
  #ifdef DEX_DEBUG
  import qualified LLVM.Analysis as L
  #endif
  import qualified LLVM.AST as L
  import qualified LLVM.AST.Global as GL
  import qualified LLVM.Module as Mod
  import qualified LLVM.Internal.Module as Mod
  #if MIN_VERSION_llvm_hs(15,0,0)
  import qualified LLVM.Passes as P
  #else
  import qualified LLVM.PassManager as P
  import qualified LLVM.Transforms as P
  #endif
  import qualified LLVM.Target as T
  import qualified LLVM.Shims
  import LLVM.Context
  import qualified Data.ByteString       as BS
  import qualified Data.ByteString.Char8 as B
  import System.IO.Unsafe
  
  import Control.Monad
  
  import PPrint ()
  import Paths_dex  (getDataFileName)
  import Types.Source
  import MonadUtil
  
  
  data LLVMOptLevel = OptALittle       -- -O1
                    | OptAggressively  -- -O3, incl. auto vectorization
    deriving Show
  
  compileLLVM :: PassLogger -> LLVMOptLevel -> L.Module -> String -> IO BS.ByteString
  compileLLVM logger opt ast exportName = do
    tm <- LLVM.Shims.newDefaultHostTargetMachine
    withContext \c -> do
      {-# SCC "LLVM.Internal.Module.withModuleFromAST" #-} Mod.withModuleFromAST c ast \m -> do
        standardCompilationPipeline opt
          logger
          [exportName] tm m
        {-# SCC "LLVM.Internal.Module.moduleObject" #-} Mod.moduleObject tm m
  {-# SCC compileLLVM #-}
  
  -- === LLVM passes ===
  
  runDefaultPasses :: LLVMOptLevel -> T.TargetMachine -> Mod.Module -> IO ()
  runDefaultPasses opt t m = do
  #if MIN_VERSION_llvm_hs(15,0,0)
    -- TODO: Modify llvm-hs to request vectorization
    P.runPasses (P.PassSetSpec [P.CuratedPassSet lvl] (Just t)) m
    where lvl = case opt of OptALittle -> 1; OptAggressively -> 3
  #else
    P.withPassManager defaultPasses \pm -> void $ P.runPassManager pm m
    case extraPasses of
      [] -> return ()
      _  -> runPasses extraPasses (Just t) m
    where
      defaultPasses = case opt of
        OptALittle ->
          P.defaultCuratedPassSetSpec {P.optLevel = Just 1, P.targetMachine = Just t}
        OptAggressively ->
          P.defaultCuratedPassSetSpec
            { P.optLevel = Just 3
            , P.loopVectorize = Just True
            , P.superwordLevelParallelismVectorize = Just True
            , P.targetMachine = Just t }
      extraPasses = []
  #endif
  
  #if MIN_VERSION_llvm_hs(15,0,0)
  runPasses :: [P.ModulePass] -> Maybe T.TargetMachine -> Mod.Module -> IO ()
  runPasses passes mt m = P.runPasses (P.PassSetSpec passes mt) m
  #else
  runPasses :: [P.Pass] -> Maybe T.TargetMachine -> Mod.Module -> IO ()
  runPasses passes mt m = do
    dl <- case mt of
           Just t  -> Just <$> T.getTargetMachineDataLayout t
           Nothing -> return Nothing
    let passSpec = P.PassSetSpec passes dl Nothing mt
    P.withPassManager passSpec \pm -> void $ P.runPassManager pm m
  #endif
  
  standardCompilationPipeline :: LLVMOptLevel -> PassLogger -> [String] -> T.TargetMachine -> Mod.Module -> IO ()
  standardCompilationPipeline opt logger exports tm m = do
    -- TODO: this means we build a copy of dexrt for every function. Is that
    -- really necessary?
    linkDexrt m
    internalize exports m
    {-# SCC showLLVM   #-} logPass JitPass $ showModule m
  #ifdef DEX_DEBUG
    {-# SCC verifyLLVM #-} L.verify m
  #endif
    {-# SCC runPasses  #-} runDefaultPasses opt tm m
    {-# SCC showOptimizedLLVM #-} logPass LLVMOpt $ showModule m
    {-# SCC showAssembly      #-} logPass AsmPass $ showAsm tm m
    where
      logPass :: PassName -> IO String -> IO ()
      logPass passName showIt = do
        s <- case ioLogLevel logger of
          DebugLogLevel -> Just <$> showIt
          NormalLogLevel -> return Nothing
        ioLogAction logger $ Outputs [PassResult passName s]
  {-# SCC standardCompilationPipeline #-}
  
  internalize :: [String] -> Mod.Module -> IO ()
  internalize names m = runPasses [P.InternalizeFunctions names, P.GlobalDeadCodeElimination] Nothing m
  
  -- === dex runtime ===
  
  {-# NOINLINE dexrtAST #-}
  dexrtAST :: L.Module
  dexrtAST = unsafePerformIO $ do
    dexrtBC <- B.readFile =<< getDataFileName "src/lib/dexrt.bc"
    withContext \ctx -> do
      Mod.withModuleFromBitcode ctx (("dexrt.c" :: String), dexrtBC) \m ->
        stripFunctionAnnotations <$> Mod.moduleAST m
    where
      -- We strip the function annotations for dexrt functions, because clang
      -- usually assumes that it's compiling for some generic CPU instead of the
      -- target that we select here. That creates a lot of confusion, making
      -- optimizations much less effective.
      stripFunctionAnnotations ast =
        ast { L.moduleDefinitions = stripDef <$> L.moduleDefinitions ast }
      stripDef def@(L.GlobalDefinition f@L.Function{..}) = case basicBlocks of
        [] -> def
        _  -> L.GlobalDefinition $ f { GL.functionAttributes = [] }
      stripDef def = def
  
  linkDexrt :: Mod.Module -> IO ()
  linkDexrt m = do
    ctx <- Mod.moduleContext m
    dataLayout <- Mod.getDataLayout =<< Mod.readModule m
    targetTriple <- Mod.getTargetTriple =<< Mod.readModule m
    let dexrtTargetAST = dexrtAST { L.moduleDataLayout = dataLayout
                                  , L.moduleTargetTriple = targetTriple }
    Mod.withModuleFromAST ctx dexrtTargetAST \dexrtm -> do
      Mod.linkModules m dexrtm
      runPasses [P.AlwaysInline True] Nothing m
  
  -- === printing ===
  
  showModule :: Mod.Module -> IO String
  showModule m = B.unpack <$> Mod.moduleLLVMAssembly m
  
  showAsm :: T.TargetMachine -> Mod.Module -> IO String
  showAsm t m' = do
    ctx <- Mod.moduleContext m'
    -- Uncomment this to dump assembly to a file that can be linked to a C benchmark suite:
    -- withModuleClone ctx m' \m -> Mod.writeObjectToFile t (Mod.File "asm.o") m
    withModuleClone ctx m' \m -> B.unpack <$> Mod.moduleTargetAssembly t m
  
  withModuleClone :: Context -> Mod.Module -> (Mod.Module -> IO a) -> IO a
  withModuleClone ctx m f = do
    -- It would be better to go through bitcode, but apparently that doesn't work...
    bc <- Mod.moduleLLVMAssembly m
    Mod.withModuleFromLLVMAssembly ctx (("<string>" :: String), bc) f
  