{-# LANGUAGE OverloadedStrings #-}
module LLVM.Test.Instrumentation where

import Test.Tasty
import Test.Tasty.HUnit

import LLVM.Test.Support

import Control.Monad.Trans.Except
import Control.Monad.Except (catchError)
import Control.Monad.IO.Class

import Data.Functor hiding (void)
import qualified Data.List as List
import qualified Data.Set as Set
import qualified Data.Map as Map

import LLVM.Module
import LLVM.Context
import LLVM.PassManager
import LLVM.Transforms
import LLVM.Target

import LLVM.AST as A
import LLVM.AST.Type
import LLVM.AST.Name
import LLVM.AST.AddrSpace
import LLVM.AST.DataLayout
import LLVM.Triple
import qualified LLVM.AST.IntegerPredicate as IPred
import qualified LLVM.AST.Linkage as L
import qualified LLVM.AST.Visibility as V
import qualified LLVM.AST.CallingConvention as CC
import qualified LLVM.AST.Attribute as A
import qualified LLVM.AST.Global as G
import qualified LLVM.AST.Constant as C

instrument :: PassSetSpec -> A.Module -> IO A.Module
instrument s m = withContext $ \context -> withModuleFromAST context m $ \mIn' -> do
  withPassManager s $ \pm -> runPassManager pm mIn'
  moduleAST mIn'

ast = do
 dl <- withHostTargetMachineDefault getTargetMachineDataLayout
 triple <- liftIO getDefaultTargetTriple
 return $ Module "<string>" "<string>" (Just dl) (Just triple) [
  -- This function is needed for AddressSanitizerModule
  GlobalDefinition $ functionDefaults {
    G.returnType = void,
    G.name = Name "asan.module_ctor",
    G.basicBlocks = [BasicBlock (UnName 0) [] (Do (Ret Nothing []))]
  },
  GlobalDefinition $ functionDefaults {
    G.returnType = i32,
    G.name = Name "foo",
    G.parameters = ([Parameter i128 (Name "x") []],False),
    G.basicBlocks = [
      BasicBlock (UnName 0) [] (Do $ Br (Name "checkDone") []),
      BasicBlock (Name "checkDone") [
        UnName 1 := Phi {
         type' = i128,
         incomingValues = [
          (LocalReference i128 (Name "x"), UnName 0),
          (LocalReference i128 (Name "x'"), Name "even"),
          (LocalReference i128 (Name "x''"), Name "odd")
         ],
         metadata = []
        },
        Name "count" := Phi {
         type' = i32,
         incomingValues = [
          (ConstantOperand (C.Int 32 1), UnName 0),
          (LocalReference i32 (Name "count'"), Name "even"),
          (LocalReference i32 (Name "count'"), Name "odd")
         ],
         metadata = []
        },
        Name "count'" := Add {
         nsw = False,
         nuw = False,
         operand0 = LocalReference i32 (Name "count"),
         operand1 = ConstantOperand (C.Int 32 1),
         metadata = []
        },
        Name "is one" := ICmp {
         iPredicate = IPred.EQ,
         operand0 = LocalReference i128 (UnName 1),
         operand1 = ConstantOperand (C.Int 128 1),
         metadata = []
        }
      ] (
        Do $ CondBr (LocalReference i1 (Name "is one")) (Name "done") (Name "checkOdd") []
      ),
      BasicBlock (Name "checkOdd") [
        Name "is odd" := Trunc (LocalReference i128 (UnName 1)) i1 []
      ] (
       Do $ CondBr (LocalReference i1 (Name "is odd")) (Name "odd") (Name "even") []
      ),
      BasicBlock (Name "even") [
        Name "x'" := UDiv True (LocalReference i128 (UnName 1)) (ConstantOperand (C.Int 128 2)) []
      ] (
        Do $ Br (Name "checkDone") []
      ),
      BasicBlock (Name "odd") [
        UnName 2 := Mul False False (LocalReference i128 (UnName 1)) (ConstantOperand (C.Int 128 3)) [],
        Name "x''" := Add False False (LocalReference i128 (UnName 2)) (ConstantOperand (C.Int 128 1)) []
      ] (
        Do $ Br (Name "checkDone") []
      ),
      BasicBlock (Name "done") [
      ] (
        Do $ Ret (Just (LocalReference i32 (Name "count'"))) []
      )
     ]
   },
  GlobalDefinition $ functionDefaults {
    G.returnType = i32,
    G.name = Name "main",
    G.parameters = ([
      Parameter i32 (Name "argc") [],
      Parameter (ptr (ptr i8)) (Name "argv") []
     ],False),
    G.basicBlocks = [
      BasicBlock (UnName 0) [
        UnName 1 := Call {
          tailCallKind = Nothing,
          callingConvention = CC.C,
          returnAttributes = [],
          function = Right
            (ConstantOperand
              (C.GlobalReference
                 (PointerType
                    { pointerReferent = FunctionType i32 [i128] False
                    , pointerAddrSpace = AddrSpace 0})
                 (Name "foo"))),
          arguments = [
           (ConstantOperand (C.Int 128 9491828328), [])
          ],
          functionAttributes = [],
          metadata = []
        }
      ] (
        Do $ Ret (Just (LocalReference i32 (UnName 1))) []
      )
     ]
   }
  ]

isMemorySanitizerSupported :: IO Bool
isMemorySanitizerSupported = do
  triple <- getProcessTargetTriple
  let ~(Right triple') = runExcept (parseTriple triple)
  let os' = os triple'
  return $ Set.member os' (Set.fromList [FreeBSD, NetBSD, Linux])

instrumentationPasses :: [(TestName, Pass, IO Bool)]
instrumentationPasses = [
    ("GCOVProfiler", defaultGCOVProfiler, return True),
    ("AddressSanitizer", defaultAddressSanitizer, return True),
    ("AddressSanitizerModule", defaultAddressSanitizerModule, return True),
    ("ThreadSanitizer", defaultThreadSanitizer, return True),
    ("BoundsChecking", BoundsChecking, return True),
    ("MemorySanitizer", defaultMemorySanitizer, isMemorySanitizerSupported)
  ]

tests =
  testGroup "Instrumentation" [
  testGroup "basic" [
    testCase n $ do
      shouldTest <- checkIfShouldTest
      if not shouldTest
        then return ()
        else do
          triple <- getProcessTargetTriple
          withTargetLibraryInfo triple $ \tli -> do
            dl <- withHostTargetMachineDefault getTargetMachineDataLayout
            ast <- ast
            ast' <- instrument (defaultPassSetSpec { transforms = [p], dataLayout = Just dl, targetLibraryInfo = Just tli }) ast
            let names ast = [ n | GlobalDefinition d <- moduleDefinitions ast, Name n <- return (G.name d) ]
            names ast' `List.intersect` names ast @?= names ast
    |
    (n, p, checkIfShouldTest) <- instrumentationPasses
    ]
  ]