content: {-# LANGUAGE DataKinds #-}
  {-# LANGUAGE FlexibleContexts #-}
  {-# LANGUAGE GADTs #-}
  {-# LANGUAGE ImplicitParams #-}
  {-# LANGUAGE LambdaCase #-}
  {-# LANGUAGE MagicHash #-}
  {-# LANGUAGE OverloadedStrings #-}
  {-# LANGUAGE PatternSynonyms #-}
  {-# LANGUAGE RankNTypes #-}
  {-# LANGUAGE ScopedTypeVariables #-}
  {-# LANGUAGE TypeApplications #-}
  {-# LANGUAGE TypeOperators #-}
  
  module Crux.LLVM.Simulate where
  
  import Control.Monad (unless)
  import Data.String (fromString)
  import qualified Data.Map.Strict as Map
  import Data.IORef
  import qualified Data.List as List
  import Data.Maybe ( fromMaybe )
  import qualified Data.Parameterized.Map as MapF
  import Data.Sequence (Seq)
  import qualified Data.Traversable as T
  import Control.Lens ((&), (%~), (%=), (^.), use, view)
  import Control.Monad.IO.Class (liftIO)
  import Data.Text as Text (Text, pack)
  import GHC.Exts ( proxy# )
  
  import System.IO (stdout)
  
  import Data.Parameterized.Some (Some(..))
  import Data.Parameterized.Context (pattern Empty, pattern (:>))
  
  import Data.LLVM.BitCode (ParseWarning(..), parseBitCodeFromFileWithWarnings,
                            ppParseWarnings)
  import qualified Text.LLVM as LLVM
  import Prettyprinter
  
  -- what4
  import qualified What4.Expr.Builder as WEB
  import What4.Interface (bvLit, bvZero)
  
  -- crucible
  import Lang.Crucible.Backend
  import Lang.Crucible.Types
  import Lang.Crucible.CFG.Core(AnyCFG(..), CFG, cfgArgTypes)
  import Lang.Crucible.FunctionHandle(newHandleAllocator,HandleAllocator)
  import Lang.Crucible.Simulator
    ( RegMap(..), assignReg, emptyRegMap
    , fnBindingsFromList, runOverrideSim, callCFG
    , initSimContext, profilingMetrics
    , ExecState( InitialState ), GlobalVar
    , SimState, defaultAbortHandler, printHandle
    , ppSimError, ovrWithBackend
    )
  import Lang.Crucible.Simulator.ExecutionTree ( actFrame, gpGlobals, stateGlobals, stateTree )
  import Lang.Crucible.Simulator.GlobalState ( insertGlobal, lookupGlobal )
  import Lang.Crucible.Simulator.Profiling ( Metric(Metric) )
  
  
  -- crucible-llvm
  import Lang.Crucible.LLVM(llvmExtensionImpl, llvmGlobals, registerLazyModule )
  import Lang.Crucible.LLVM.Bytes ( bytesToBV )
  import Lang.Crucible.LLVM.Globals
          ( initializeAllMemory, populateAllGlobals )
  import Lang.Crucible.LLVM.MemModel
          ( Mem, MemImpl, mkMemVar, withPtrWidth, memAllocCount, memWriteCount
          , MemOptions(..), HasLLVMAnn, LLVMAnnMap
          , explainCex, CexExplanation(..), doAlloca
          , pattern LLVMPointerRepr, LLVMPointerType
          , pattern PtrRepr, pattern PtrWidth, llvmPointer_bv
          )
  import Lang.Crucible.LLVM.MemModel.CallStack (ppCallStack)
  import Lang.Crucible.LLVM.MemType (MemType(..), SymType(..), i8, memTypeAlign, memTypeSize)
  import Lang.Crucible.LLVM.PrettyPrint (ppSymbol)
  import Lang.Crucible.LLVM.Translation
          ( translateModule, ModuleTranslation, globalInitMap
          , transContext, getTranslatedCFG, llvmPtrWidth, llvmTypeCtx
          , LLVMTranslationWarning(..)
          )
  import Lang.Crucible.LLVM.Intrinsics
          (llvmIntrinsicTypes, register_llvm_overrides )
  
  import Lang.Crucible.LLVM.Errors( ppBB )
  import Lang.Crucible.LLVM.Extension( LLVM, ArchWidth )
  import Lang.Crucible.LLVM.SymIO
         ( LLVMFileSystem, llvmSymIOIntrinsicTypes, symio_overrides, initialLLVMFileSystem, SomeOverrideSim(..) )
  import Lang.Crucible.LLVM.TypeContext (llvmDataLayout)
  import qualified Lang.Crucible.SymIO as SymIO
  import qualified Lang.Crucible.SymIO.Loader as SymIO.Loader
  
  -- crux
  import qualified Crux
  
  import Crux.Types
  import Crux.Log ( outputHandle )
  
  import Crux.LLVM.Config
  import qualified Crux.LLVM.Log as Log
  import Crux.LLVM.Overrides
  
  -- | Create a simulator context for the given architecture.
  setupSimCtxt ::
    (IsSymBackend sym bak, HasLLVMAnn sym) =>
    HandleAllocator ->
    bak ->
    MemOptions ->
    GlobalVar Mem ->
    SimCtxt Crux sym LLVM
  setupSimCtxt halloc bak mo memVar =
    initSimContext bak
                   (MapF.union llvmIntrinsicTypes llvmSymIOIntrinsicTypes)
                   halloc
                   stdout
                   (fnBindingsFromList [])
                   (llvmExtensionImpl mo)
                   CruxPersonality
      & profilingMetrics %~ Map.union (memMetrics memVar)
  
  -- | Parse an LLVM bit-code file.
  parseLLVM ::
    Crux.Logs msgs =>
    Log.SupportsCruxLLVMLogMessage msgs =>
    FilePath ->
    IO LLVM.Module
  parseLLVM file =
    do ok <- parseBitCodeFromFileWithWarnings file
       case ok of
         Left err -> throwCError (LLVMParseError err)
         Right (m, warnings) -> do
           sayParseWarnings warnings
           return m
  
  registerFunctions ::
    Crux.Logs msgs =>
    Log.SupportsCruxLLVMLogMessage msgs =>
    (ArchOk arch, IsSymInterface sym, HasLLVMAnn sym, ptrW ~ ArchWidth arch) =>
    LLVMOptions ->
    LLVM.Module ->
    ModuleTranslation arch ->
    Maybe (LLVMFileSystem ptrW) ->
    OverM Crux sym LLVM ()
  registerFunctions llvmOpts llvm_module mtrans fs0 =
    do let llvm_ctx = mtrans ^. transContext
       let ?lc = llvm_ctx ^. llvmTypeCtx
           ?intrinsicsOpts = intrinsicsOpts llvmOpts
           ?memOpts = memOpts llvmOpts
  
       -- register the callable override functions
       _ <- register_llvm_overrides llvm_module []
              (concat [ cruxLLVMOverrides proxy#
                      , svCompOverrides
                      , cbmcOverrides proxy#
                      , maybe [] symio_overrides fs0
                      ])
              llvm_ctx
  
       -- register all the functions defined in the LLVM module
       registerLazyModule sayTranslationWarning mtrans
  
  simulateLLVMFile ::
    Crux.Logs msgs =>
    Log.SupportsCruxLLVMLogMessage msgs =>
    -- | Path to the LLVM module
    FilePath ->
    LLVMOptions ->
    Crux.SimulatorCallbacks msgs st Crux.Types.CruxSimulationResult
  simulateLLVMFile llvm_file llvmOpts =
    Crux.SimulatorCallbacks $
      do bbMapRef <- newIORef (Map.empty :: LLVMAnnMap sym)
         let ?recordLLVMAnnotation =
               \callStack an bb ->
                 modifyIORef bbMapRef (Map.insert an (callStack, bb))
         return $
           Crux.SimulatorHooks
             { Crux.setupHook =
                 \bak maybeOnline ->
                   do halloc <- newHandleAllocator
                      setupFileSim halloc llvm_file llvmOpts bak maybeOnline
             , Crux.onErrorHook =
                 \bak -> return (explainFailure (backendGetSym bak) bbMapRef)
             , Crux.resultHook = \_sym result -> return result
             }
  
  setupFileSim :: forall sym bak t st fs msgs
                . Crux.Logs msgs
               => Log.SupportsCruxLLVMLogMessage msgs
               => IsSymBackend sym bak
               => sym ~ WEB.ExprBuilder t st fs
               => HasLLVMAnn sym
               => HandleAllocator
               -> FilePath
               -> LLVMOptions
               -> bak
               -> Maybe (Crux.SomeOnlineSolver sym bak)
               -> IO (Crux.RunnableState sym)
  setupFileSim halloc llvm_file llvmOpts bak _maybeOnline =
    do let sym = backendGetSym bak
       memVar <- mkMemVar "crux:llvm_memory" halloc
  
       let simctx = (setupSimCtxt halloc bak (memOpts llvmOpts) memVar)
                    { printHandle = view outputHandle ?outputConfig }
  
       prepped <- prepLLVMModule llvmOpts halloc bak llvm_file memVar
  
       let globSt = llvmGlobals memVar (prepMem prepped)
  
       Some trans <- return $ prepSomeTrans prepped
       llvmPtrWidth (trans ^. transContext) $ \ptrW -> withPtrWidth ptrW $ do
         mContents <- T.traverse (SymIO.Loader.loadInitialFiles sym) (symFSRoot llvmOpts)
         -- We modify the defaults, which are extremely conservative.  Unless the
         -- user directs otherwise, we default to connecting stdin, stdout, and stderr.
         let defaultFileContents = (SymIO.emptyInitialFileSystemContents @sym)
                                   { SymIO.useStderr = True
                                   , SymIO.useStdout = True
                                   , SymIO.concreteFiles = Map.fromList [(SymIO.StdinTarget, mempty)]
                                   }
         let fsContents = fromMaybe defaultFileContents mContents
         let mirroredOutputs = [ (SymIO.StdoutTarget, ?outputConfig ^. Crux.outputHandle)
                               , (SymIO.StderrTarget, ?outputConfig ^. Crux.outputHandle)
                               ]
         (fs0, globSt', SomeOverrideSim initFSOverride) <- initialLLVMFileSystem halloc sym ptrW fsContents mirroredOutputs globSt
         return $ Crux.RunnableState $
           InitialState simctx globSt' defaultAbortHandler UnitRepr $
             runOverrideSim UnitRepr $
               withPtrWidth ptrW $
                  do registerFunctions llvmOpts (prepLLVMMod prepped) trans (Just fs0)
                     initFSOverride
                     checkFun llvmOpts trans memVar
  
  
  data PreppedLLVM sym = PreppedLLVM { prepLLVMMod :: LLVM.Module
                                     , prepSomeTrans :: Some ModuleTranslation
                                     , prepMemVar :: GlobalVar Mem
                                     , prepMem :: MemImpl sym
                                     }
  
  -- | Given an LLVM Bitcode file, and a GlobalVar memory, translate the
  -- file into the Crucible representation and add the globals and
  -- definitions from the file to the GlobalVar memory.
  
  prepLLVMModule :: IsSymBackend sym bak
                 => HasLLVMAnn sym
                 => Crux.Logs msgs
                 => Log.SupportsCruxLLVMLogMessage msgs
                 => LLVMOptions
                 -> HandleAllocator
                 -> bak
                 -> FilePath  -- for the LLVM bitcode file
                 -> GlobalVar Mem
                 -> IO (PreppedLLVM sym)
  prepLLVMModule llvmOpts halloc bak bcFile memVar = do
      llvmMod <- parseLLVM bcFile
      Some trans <-
          let ?transOpts = transOpts llvmOpts
           in translateModule halloc memVar llvmMod
      mem <- let llvmCtxt = trans ^. transContext in
               let ?lc = llvmCtxt ^. llvmTypeCtx
                   ?memOpts = memOpts llvmOpts
               in llvmPtrWidth llvmCtxt $ \ptrW ->
                 withPtrWidth ptrW $ do
                 Log.sayCruxLLVM (Log.UsingPointerWidthForFile (intValue ptrW) (Text.pack bcFile))
                 populateAllGlobals bak (trans ^. globalInitMap)
                   =<< initializeAllMemory bak llvmCtxt llvmMod
      return $ PreppedLLVM llvmMod (Some trans) memVar mem
  
  sayParseWarnings ::
    Log.SupportsCruxLLVMLogMessage msgs =>
    Crux.Logs msgs =>
    Seq ParseWarning -> IO ()
  sayParseWarnings warnings =
    unless (null warnings) $
      Log.sayCruxLLVM $ Log.ParseWarning $ Text.pack $ show $
      ppParseWarnings warnings
  
  sayTranslationWarning ::
    Log.SupportsCruxLLVMLogMessage msgs =>
    Crux.Logs msgs =>
    LLVMTranslationWarning -> IO ()
  sayTranslationWarning = Log.sayCruxLLVM . f
    where
      f (LLVMTranslationWarning s p msg) =
          Log.TranslationWarning (Text.pack (show (ppSymbol s))) (Text.pack (show p)) msg
  
  checkFun ::
    forall arch msgs personality sym.
    IsSymInterface sym =>
    HasLLVMAnn sym =>
    ArchOk arch =>
    Crux.Logs msgs =>
    Log.SupportsCruxLLVMLogMessage msgs =>
    LLVMOptions ->
    ModuleTranslation arch ->
    GlobalVar Mem ->
    OverM personality sym LLVM ()
  checkFun llvmOpts trans memVar =
    liftIO (getTranslatedCFG trans (fromString nm)) >>= \case
      Just (_, AnyCFG anyCfg, warns) ->
        do liftIO (mapM_ sayTranslationWarning warns)
           case cfgArgTypes anyCfg of
             Empty -> simulateFun anyCfg emptyRegMap
  
             (Empty :> LLVMPointerRepr w :> PtrRepr)
               |  isMain, shouldSupplyMainArguments
               ,  Just Refl <- testEquality w (knownNat @32)
               -> checkMainWithArguments anyCfg
  
             _ -> throwCError (BadFun nm isMain)
  
      Nothing -> throwCError (MissingFun nm)
    where
      nm     = entryPoint llvmOpts
      isMain = nm == "main"
      shouldSupplyMainArguments =
        case supplyMainArguments llvmOpts of
          NoArguments    -> False
          EmptyArguments -> True
  
      simulateFun :: CFG LLVM blocks ctx ret ->
                     RegMap sym ctx ->
                     OverM personality sym LLVM ()
      simulateFun anyCfg args = do
        liftIO $ Log.sayCruxLLVM (Log.SimulatingFunction (Text.pack nm))
        callCFG anyCfg args >> return ()
  
      -- Simulate a function with the signature @int main(int argc, char* argv[])@.
      -- We do so by behaving as if we are starting from this entrypoint:
      --
      --   int crucible_main() {
      --     int argc = 0;
      --     char *argv[] = {};
      --     return main(argc, argv);
      --   }
      --
      -- @argc@ can easily be created with 'bvLit'. For @argv@, we have to choose
      -- what sort of allocation we want when creating its @LLVMPointer@. We've
      -- opted for 'doAlloca' since the LLVM bitcode to allocate @argv@ will use
      -- the @alloca@ instruction. This means that @argv@ will be
      -- stack-allocated, which is a reasonable choice, given that the @main@
      -- function's @argv@ argument actually lives on the stack in most binaries.
      checkMainWithArguments ::
        CFG LLVM blocks (EmptyCtx ::> LLVMPointerType 32 ::> TPtr arch) ret ->
        OverM personality sym LLVM ()
      checkMainWithArguments anyCfg = ovrWithBackend $ \bak -> do
        let ?memOpts  = memOpts llvmOpts
        let w         = knownNat @32
            dl        = llvmDataLayout (trans^.transContext.llvmTypeCtx)
            tp        = ArrayType 0 (PtrType (MemType i8))
            tp_sz     = memTypeSize  dl tp
            alignment = memTypeAlign dl tp
            loc       = "crux-llvm main(argc, argv) simulation"
            sym       = backendGetSym bak
  
        gs  <- use (stateTree.actFrame.gpGlobals)
        mem <- case lookupGlobal memVar gs of
                 Just mem -> pure mem
                 Nothing  -> fail "checkFun.checkMainWithArguments: Memory missing from global vars"
        argc <- liftIO $ llvmPointer_bv sym =<< bvZero sym w
        sz   <- liftIO $ bvLit sym PtrWidth $ bytesToBV PtrWidth tp_sz
        (argv, mem') <- liftIO $ doAlloca bak mem sz alignment loc
        stateTree.actFrame.gpGlobals %= insertGlobal memVar mem'
  
        let args = assignReg PtrRepr argv $
                   assignReg (LLVMPointerRepr w) argc emptyRegMap
        simulateFun anyCfg args
  
  ---------------------------------------------------------------------
  -- Profiling
  
  memMetrics :: forall p sym ext
               . GlobalVar Mem
              -> Map.Map Text (Metric p sym ext)
  memMetrics memVar = Map.fromList
                      -- Note: These map keys are used by profile.js in
                      -- https://github.com/GaloisInc/sympro-ui and must
                      -- match the names there.
                      [ ("LLVM.allocs", allocs)
                      , ("LLVM.writes", writes)
                      ]
    where
      allocs = Metric $ measureMemBy memAllocCount
      writes = Metric $ measureMemBy memWriteCount
  
      measureMemBy :: (MemImpl sym -> Int)
                   -> SimState p sym ext r f args
                   -> IO Integer
      measureMemBy f st = do
        let globals = st ^. stateGlobals
        case lookupGlobal memVar globals of
          Just mem -> return $ toInteger (f mem)
          Nothing -> fail "Memory missing from global vars"
  
  ----------------------------------------------------------------------
  
  -- arbitrary, we should probably make this limit configurable
  detailLimit :: Int
  detailLimit = 10
  
  explainFailure :: IsSymInterface sym
                 => sym ~ WEB.ExprBuilder t st fs
                 => sym
                 -> IORef (LLVMAnnMap sym)
                 -> Crux.Explainer sym t ann
  explainFailure sym bbMapRef evalFn gl =
    do bb <- readIORef bbMapRef
       ex <- explainCex sym bb evalFn >>= \f -> f (gl ^. labeledPred)
       let details =
             case ex of
               NoExplanation -> mempty
               DisjOfFailures xs ->
                 case ppBBPair <$> xs of
                   []  -> mempty
                   [x] -> indent 2 x
                   xs' ->
                     let xs'' = List.take detailLimit xs'
                         xs'l = length xs'
                         msg1 = "Failing conditions::"
                         msg2 = if xs'l > detailLimit
                                then "[only displaying the first"
                                     <+> pretty detailLimit
                                     <+> "of" <+> pretty xs'l
                                     <+> "failed conditions]"
                                else "Total failed conditions:" <+> pretty xs'l
                     in nest 2 $ vcat $ msg1 : xs'' <> [msg2]
       return $ vcat [ ppSimError (gl^. labeledPredMsg), details ]
    where ppBBPair (callStack, bb) =
            vsep [ ppBB bb
                 , "in context:"
                 , indent 2 (ppCallStack callStack)
                 ]
  