content: {-# LANGUAGE ConstraintKinds #-}
  {-# LANGUAGE DataKinds #-}
  {-# LANGUAGE FlexibleContexts #-}
  {-# LANGUAGE GADTs #-}
  {-# LANGUAGE ImplicitParams #-}
  {-# LANGUAGE PolyKinds #-}
  {-# LANGUAGE RankNTypes #-}
  {-# LANGUAGE ScopedTypeVariables #-}
  {-# LANGUAGE TypeApplications #-}
  {-# LANGUAGE TypeOperators #-}
  module Main ( main ) where
  
  import           Control.Lens ( (^.) )
  import qualified Control.Lens as L
  import           Control.Monad ( unless, forM_ )
  import qualified Data.Macaw.CFG as MC
  import qualified Data.Macaw.Discovery as MD
  import qualified Data.Macaw.Refinement as MR
  import           Data.Macaw.Symbolic ( SymArchConstraints )
  import qualified Data.Map as M
  import           Data.Maybe ( isNothing )
  import           Data.Monoid
  import           Data.Parameterized.Some
  import qualified Data.Text.IO as TIO
  import qualified Lang.Crucible.LLVM.MemModel as LLVM
  import qualified Options.Applicative as O
  import           Prettyprinter as PP
  
  import           Prelude
  
  import           Initialization ( withElf, withRefinedDiscovery )
  import           Options( Options(..) )
  import           Summary
  
  optionsParser :: O.Parser Options
  optionsParser = Options
                  <$> O.strArgument ( O.metavar "FILE"
                                      <> O.help "The binary ELF file to perform discovery on"
                                    )
                  <*> O.switch ( O.long "unrefined"
                               <> O.short 'u'
                               <> O.help "No refinement of discovery results"
                               )
                  <*> O.switch ( O.long "summary"
                               <> O.short 's'
                               <> O.help "Only show summary of discovery/refinement.\n\
                                         \Without this flag a full list of all discovered\n\
                                         \functions and blocks is output."
                               )
                  <*> O.switch ( O.long "verbose"
                               <> O.short 'v'
                               <> O.help "Show verbose information about each discovered\n\
                                         \function and block."
                                         )
                  <*> O.option O.auto ( O.metavar "SOLVER"
                                     <> O.help "The SMT solver to use"
                                     <> O.short 'S'
                                     <> O.long "solver"
                                     <> O.value MR.Yices
                                      )
                  <*> O.optional ( O.strOption ( O.metavar "FILE"
                                              <> O.long "solver-interaction-file"
                                              <> O.help "File to save solver interactions to"
                                               ))
                  <*> O.option O.auto ( O.metavar "INT"
                                      <> O.help "The maximum number of models to consider valid for a given indirect call"
                                      <> O.value (MR.maximumModelCount MR.defaultRefinementConfig)
                                      <> O.long "maximum-model-count"
                                      <> O.short 'm'
                                      )
                  <*> O.option O.auto ( O.metavar "NAT"
                                      <> O.help "The number of solver processes to run concurrently (minimum 1)"
                                      <> O.value 1
                                      <> O.short 'N'
                                      <> O.long "solver-processes"
                                      )
                  <*> O.option O.auto ( O.metavar "NAT"
                                      <> O.help "The number of seconds to wait for the solver before timing out"
                                      <> O.value (MR.timeoutSeconds MR.defaultRefinementConfig)
                                      <> O.long "timeout"
                                      <> O.short 't'
                                      )
  
  main :: IO ()
  main = O.execParser optParser >>= doRefinement
    where optParser = O.info ( optionsParser O.<**> O.helper )
                      ( O.fullDesc
                      <> O.progDesc "A tool to show refined code discovery for ELF binaries"
                      <> O.header "run-refinement - code discovery output"
                      )
  
  doRefinement :: Options -> IO ()
  doRefinement opts = do
    withElf opts $ \_ archInfo bin unrefinedDI -> do
      case unrefined opts of
        True -> showDiscoveryInfo opts unrefinedDI Nothing
        False ->
          let ?memOpts = LLVM.defaultMemOptions in
          withRefinedDiscovery opts archInfo bin $ \refinedDI refinedInfo -> do
            showDiscoveryInfo opts unrefinedDI (Just (refinedDI, refinedInfo))
  
  showDiscoveryInfo :: (SymArchConstraints arch)
                    => Options
                    -> MD.DiscoveryState arch
                    -> Maybe (MD.DiscoveryState arch, MR.RefinementInfo arch)
                    -> IO ()
  showDiscoveryInfo opts unrefinedDI mrefinedDI = do
    unless (summaryOnly opts) $ do
      if verbose opts
        then showDetails unrefinedDI mrefinedDI
        else showOverview unrefinedDI mrefinedDI
  
    showSummary unrefinedDI mrefinedDI
  
  
  
  showSummary :: forall arch . (MC.MemWidth (MC.ArchAddrWidth arch)) => MD.DiscoveryState arch -> Maybe (MD.DiscoveryState arch, MR.RefinementInfo arch) -> IO ()
  showSummary unrefinedDI mdirefined =
    let summarize di = vcat [ pretty ":: ==== SUMMARY ===="
                            , pretty (summarizeInfo di)
                            ]
    in case mdirefined of
      Nothing -> putStrLn (show (summarize unrefinedDI))
      Just (refinedDI, rinfo) -> do
        let lhs = PP.vcat [ PP.pretty "Unrefined"
                          , summarize unrefinedDI
                          ]
        let rhs = PP.vcat [ PP.pretty "Refined"
                          , PP.nest 4 (summarizeRefinementInfo rinfo)
                          , summarize refinedDI
                          ]
        putStrLn (show (PP.vsep [ lhs, rhs ]))
  
  showOverview :: (MC.MemWidth (MC.ArchAddrWidth arch))
               => MD.DiscoveryState arch
               -> Maybe (MD.DiscoveryState arch, MR.RefinementInfo arch)
               -> IO ()
  showOverview unrefinedDI mrefinedDI =
    let getIssue dfi (blkAddr, pblk) =
          let issue = case MD.pblockTermStmt pblk of
                MD.ParsedTranslateError r -> pretty "Translation failure:" <+> pretty (show r)
                MD.ClassifyFailure _ rsns
                  | isNothing (lookup (MD.pblockAddr pblk) (MD.discoveredClassifyFailureResolutions dfi)) ->
                    PP.vcat [ PP.pretty "Classify failure: "
                            , PP.nest 4 (PP.vcat (map PP.pretty rsns))
                            ]
                _ -> emptyDoc
          in hsep [ pretty "Block @", pretty $ show blkAddr, issue ]
        funcSummary (funAddr, (Some dfi)) =
          let blkSummary = map (getIssue dfi) (dfi ^. MD.parsedBlocks . L.to M.toList)
          in vcat [ pretty "Function @" <+> pretty (show funAddr)
                  , indent 2 $ vcat blkSummary
                  ]
        summaries di = map funcSummary (di ^. MD.funInfo . L.to M.toList)
    in case mrefinedDI of
      Nothing -> putStrLn (show (PP.vcat (summaries unrefinedDI)))
      Just (refinedDI, rinfo) -> do
        let lhs = PP.vcat ( PP.pretty "Unrefined"
                          : PP.pretty "========="
                          : summaries unrefinedDI
                          )
        let rhs = PP.vcat ( PP.pretty "Refined"
                          : PP.pretty "======="
                          : PP.nest 4 (summarizeRefinementInfo rinfo)
                          : summaries refinedDI
                          )
        putStrLn (show (PP.vsep [ lhs, rhs ]))
  
  summarizeRefinementInfo :: (MC.MemWidth (MC.ArchAddrWidth arch)) => MR.RefinementInfo arch -> PP.Doc a
  summarizeRefinementInfo rinfo =
    PP.vcat [ PP.pretty "Timeouts: " <> PP.list (fmap (PP.pretty . show) (MR.refinementTimeouts rinfo))
            , PP.pretty "Spurious: " <> PP.list (fmap (PP.pretty . show) (MR.refinementSpurious rinfo))
            , PP.pretty "No Models: " <> PP.list (fmap (PP.pretty . show) (MR.refinementNoModels rinfo))
            , PP.pretty "Errors: " <> PP.list (fmap (PP.pretty . show) (MR.refinementErrors rinfo))
            ]
  
  showDetails :: (SymArchConstraints arch)
              => MD.DiscoveryState arch
              -> Maybe (MD.DiscoveryState arch, MR.RefinementInfo arch)
              -> IO ()
  showDetails di _ =
    forM_ (M.toList (di ^. MD.funInfo)) $ \(funAddr, Some dfi) -> do
      putStrLn $ "===== BEGIN FUNCTION " ++ show funAddr ++ " ====="
      forM_ (M.toList (dfi ^. MD.parsedBlocks)) $ \(blockAddr, pb) -> do
        putStrLn $ "== begin block " ++ show blockAddr ++ " =="
        putStrLn . show $ MD.pblockStmts pb
        putStrLn ""
        case MD.pblockTermStmt pb of
          MD.ParsedTranslateError r -> do
            putStr "*** "
            putStr "TRANSLATE ERROR: "
            TIO.putStrLn r
          e@(MD.ClassifyFailure {}) -> do
            putStr "*** "
            putStr "CLASSIFY FAILURE: "
            putStrLn $ show e
          r -> do
            putStr "### block terminates with: "
            putStrLn $ show r
        putStrLn ""
      putStrLn ""
      putStrLn ""
  