content: {-# LANGUAGE OverloadedStrings #-}
  {-# LANGUAGE RankNTypes #-}
  {-# LANGUAGE ScopedTypeVariables #-}
  {-# Language ImplicitParams #-}
  {-# Language MultiWayIf #-}
  {-# Language PatternSynonyms #-}
  {-# Language ScopedTypeVariables #-}
  {-# Language TypeFamilies #-}
  {-# Language TypeOperators #-}
  {-# Language ViewPatterns #-}
  
  module Crux.Goal where
  
  import Control.Concurrent.Async (async, asyncThreadId, waitAnyCatch)
  import Control.Exception (throwTo, SomeException, displayException)
  import Control.Lens ((^.), view)
  
  import Control.Monad (forM, forM_, unless, when)
  import Data.Either (partitionEithers)
  import Data.IORef
  import Data.Maybe (fromMaybe)
  import qualified Data.Map as Map
  import qualified Data.Parameterized.Map as MapF
  import qualified Data.Text as Text
  import           Data.Void
  import           Prettyprinter
  import           System.Exit (ExitCode(ExitSuccess))
  import           System.FilePath ((<.>), splitExtension)
  import           System.IO (Handle, IOMode(..), withFile)
  import qualified System.Timeout as ST
  
  import What4.Interface (notPred, getConfiguration, IsExprBuilder)
  import What4.Config (setOpt, getOptionSetting, Opt, ConfigOption)
  import What4.ProgramLoc (ProgramLoc)
  import What4.SatResult(SatResult(..))
  import What4.Expr (ExprBuilder, GroundEvalFn(..), BoolExpr, GroundValueWrapper(..))
  import What4.Protocol.Online( OnlineSolver, SolverProcess, inNewFrame, inNewFrame2Open
                              , inNewFrame2Close, solverEvalFuns, solverConn
                              , check, getUnsatCore, getAbducts )
  import What4.Protocol.SMTWriter( SMTReadWriter, mkFormula, assumeFormulaWithFreshName
                                 , assumeFormula, smtExprGroundEvalFn )
  import qualified What4.Solver as WS
  import Lang.Crucible.Backend
  import Lang.Crucible.Backend.Online
          ( OnlineBackend, withSolverProcess, enableOnlineBackend, solverInteractionFile )
  import Lang.Crucible.Simulator.SimError
          ( SimError(..), SimErrorReason(..) )
  import Lang.Crucible.Simulator.ExecutionTree
          (ctxSymInterface)
  import Lang.Crucible.Panic (panic)
  
  import Crux.Types
  import Crux.Model
  import Crux.Log as Log
  import Crux.Config.Common
  import Crux.ProgressBar
  
  
  symCfg :: (IsExprBuilder sym, Opt t a) => sym -> ConfigOption t -> a -> IO ()
  symCfg sym x y =
    do opt <- getOptionSetting x (getConfiguration sym)
       _   <- setOpt opt y
       pure ()
  
  
  -- | Simplify the proved goals.
  provedGoalsTree :: forall sym.
    ( IsSymInterface sym
    ) =>
    sym ->
    Maybe (Goals (Assumptions sym) (Assertion sym, [ProgramLoc], ProofResult sym)) ->
    IO (Maybe ProvedGoals)
  provedGoalsTree sym = traverse (go mempty)
    where
    go :: Assumptions sym ->
          Goals (Assumptions sym) (Assertion sym, [ProgramLoc], ProofResult sym) ->
          IO ProvedGoals
    go asmps gs =
      case gs of
        Assuming as gs1 -> go (asmps <> as) gs1
        Prove (p,locs,r) -> proveToGoal sym asmps p locs r
        ProveConj g1 g2 -> Branch <$> go asmps g1 <*> go asmps g2
  
  proveToGoal ::
    (IsSymInterface sym) =>
    sym ->
    Assumptions sym ->
    Assertion sym ->
    [ProgramLoc] ->
    ProofResult sym ->
    IO ProvedGoals
  proveToGoal sym allAsmps p locs pr =
    case pr of
      NotProved ex cex s ->
        do as <- flattenAssumptions sym allAsmps
           return (NotProvedGoal (map showAsmp as) (showGoal p) ex locs cex s)
      Proved xs ->
        case partitionEithers xs of
          (asmps, [])  -> return (ProvedGoal (map showAsmp asmps) (showGoal p) locs True)
          (asmps, _:_) -> return (ProvedGoal (map showAsmp asmps) (showGoal p) locs False)
  
   where
   showAsmp x = forgetAssumption x
   showGoal x = x^.labeledPredMsg
  
  
  countGoals :: Goals a b -> Int
  countGoals gs =
    case gs of
      Assuming _ gs1  -> countGoals gs1
      Prove _         -> 1
      ProveConj g1 g2 -> countGoals g1 + countGoals g2
  
  isResourceExhausted :: LabeledPred p SimError -> Bool
  isResourceExhausted (view labeledPredMsg -> SimError _ (ResourceExhausted _)) = True
  isResourceExhausted _ = False
  
  updateProcessedGoals ::
    LabeledPred p SimError ->
    ProofResult a ->
    ProcessedGoals ->
    ProcessedGoals
  updateProcessedGoals _ (Proved _) pgs =
    pgs{ totalProcessedGoals = 1 + totalProcessedGoals pgs
       , provedGoals = 1 + provedGoals pgs
       }
  
  updateProcessedGoals res (NotProved _ _ _) pgs | isResourceExhausted res =
    pgs{ totalProcessedGoals = 1 + totalProcessedGoals pgs
       , incompleteGoals = 1 + incompleteGoals pgs
       }
  
  updateProcessedGoals _ (NotProved _ (Just _) _) pgs =
    pgs{ totalProcessedGoals = 1 + totalProcessedGoals pgs
       , disprovedGoals = 1 + disprovedGoals pgs
       }
  
  updateProcessedGoals _ (NotProved _ Nothing _) pgs =
    pgs{ totalProcessedGoals = 1 + totalProcessedGoals pgs }
  
  -- | A function that can be used to generate a pretty explanation of a
  -- simulation error.
  
  type Explainer sym t ann = Maybe (GroundEvalFn t)
                             -> LPred sym SimError
                             -> IO (Doc ann)
  
  type ProverCallback sym =
    forall ext personality t st fs.
      (sym ~ ExprBuilder t st fs) =>
      CruxOptions ->
      SimCtxt personality sym ext ->
      Explainer sym t Void ->
      Maybe (Goals (Assumptions sym) (Assertion sym)) ->
      IO (ProcessedGoals, Maybe (Goals (Assumptions sym) (Assertion sym, [ProgramLoc], ProofResult sym)))
  
  -- | Discharge a tree of proof obligations ('Goals') by using a non-online solver
  --
  -- This function traverses the 'Goals' tree while keeping track of a collection
  -- of assumptions in scope for each goal.  For each proof goal encountered in
  -- the tree, it creates a fresh solver connection using the provided solver
  -- adapter.
  --
  -- This is in contrast to 'proveGoalsOnline', which uses an online solver
  -- connection with scoped assumption frames.  This function allows using a wider
  -- variety of solvers (i.e., ones that don't have support for online solving)
  -- and would in principle enable parallel goal evaluation (though the tree
  -- structure makes that a bit trickier, it isn't too hard).
  --
  -- Note that this function uses the same symbolic backend ('ExprBuilder') as the
  -- symbolic execution phase, which should not be a problem.
  proveGoalsOffline :: forall st sym p t fs personality msgs.
    ( sym ~ ExprBuilder t st fs
    , Logs msgs
    , SupportsCruxLogMessage msgs
    ) =>
    [WS.SolverAdapter st] ->
    CruxOptions ->
    SimCtxt personality sym p ->
    (Maybe (GroundEvalFn t) -> Assertion sym -> IO (Doc Void)) ->
    Maybe (Goals (Assumptions sym) (Assertion sym)) ->
    IO (ProcessedGoals, Maybe (Goals (Assumptions sym) (Assertion sym, [ProgramLoc], ProofResult sym)))
  proveGoalsOffline _adapter _opts _ctx _explainFailure Nothing = return (ProcessedGoals 0 0 0 0, Nothing)
  proveGoalsOffline adapters opts ctx explainFailure (Just gs0) = do
    goalNum <- newIORef (ProcessedGoals 0 0 0 0)
    (start,end,finish) <- proverMilestoneCallbacks gs0
    gs <- go (start,end) goalNum mempty gs0
    nms <- readIORef goalNum
    finish
    return (nms, Just gs)
  
    where
      sym = ctx^.ctxSymInterface
  
      failfast = proofGoalsFailFast opts
  
      go :: SupportsCruxLogMessage msgs
         => (ProverMilestoneStartGoal, ProverMilestoneEndGoal)
         -> IORef ProcessedGoals
         -> Assumptions sym
         -> Goals (Assumptions sym) (Assertion sym)
         -> IO (Goals (Assumptions sym) (Assertion sym, [ProgramLoc], ProofResult sym))
      go (start,end) goalNum assumptionsInScope gs =
        case gs of
          Assuming ps gs1 -> do
            res <- go (start,end) goalNum (assumptionsInScope <> ps) gs1
            return (Assuming ps res)
  
          ProveConj g1 g2 -> do
            g1' <- go (start,end) goalNum assumptionsInScope g1
            numDisproved <- disprovedGoals <$> readIORef goalNum
            if failfast && numDisproved > 0
              then return g1'
              else ProveConj g1' <$> go (start,end) goalNum assumptionsInScope g2
  
          Prove p -> do
            goalNumber <- totalProcessedGoals <$> readIORef goalNum
            start goalNumber
  
            -- Conjoin all of the in-scope assumptions, the goal, then negate and
            -- check sat with the adapter
            assumptions <- assumptionsPred sym assumptionsInScope
            goal <- notPred sym (p ^. labeledPred)
  
            res <- dispatchSolversOnGoalAsync (goalTimeout opts) adapters
                             (runOneSolver p assumptionsInScope assumptions goal goalNumber)
            end goalNumber
            case res of
              Right Nothing -> do
                let details = NotProved "(timeout)" Nothing []
                let locs = assumptionsTopLevelLocs assumptionsInScope
                modifyIORef' goalNum (updateProcessedGoals p details)
                return (Prove (p, locs, details))
  
              Right (Just (locs,details)) -> do
                modifyIORef' goalNum (updateProcessedGoals p details)
                case details of
                  NotProved _ (Just _) _ ->
                    when (failfast && not (isResourceExhausted p)) $
                      sayCrux Log.FoundCounterExample
                  _ -> return ()
                return (Prove (p, locs, details))
              Left es -> do
                modifyIORef' goalNum (updateProcessedGoals p (NotProved mempty Nothing []))
                let allExceptions = unlines (displayException <$> es)
                fail allExceptions
  
      runOneSolver :: Assertion sym
                   -> Assumptions sym
                   -> BoolExpr t
                   -> BoolExpr t
                   -> Integer
                   -> WS.SolverAdapter st
                   -> IO ([ProgramLoc], ProofResult sym)
      runOneSolver p assumptionsInScope assumptions goal goalNumber adapter =
        -- Create a file to a single offline solver interaction, assuming the user
        -- has selected this option. A single Crux session might invoke an offline
        -- solver multiple times, so each file has unique suffixes to indicate the
        -- goal number and which solver in particular was used to solve the goal.
        -- (Recall that a Crux user can choose multiple offline solvers, so it is
        -- important to indicate which solver was picked for each goal.)
        withMaybeOfflineSolverOutputHandle (offlineSolverOutput opts) $ \mbLogHandle ->
        let logData = WS.defaultLogData { WS.logHandle = mbLogHandle } in
        WS.solver_adapter_check_sat adapter (ctx ^. ctxSymInterface) logData [assumptions, goal] $ \satRes ->
          case satRes of
            Unsat _ -> do
              -- NOTE: We don't have an easy way to get an unsat core here
              -- because we don't have a solver connection.
              as <- flattenAssumptions sym assumptionsInScope
              let core = fmap Left as ++ [ Right p ]
              let locs = assumptionsTopLevelLocs assumptionsInScope
              return (locs, Proved core)
            Sat (evalFn, _) -> do
              evs  <- concretizeEvents (groundEval evalFn) assumptionsInScope
              let vals = evalModelFromEvents evs
              explain <- explainFailure (Just evalFn) p
              let locs = map eventLoc evs
              return (locs, NotProved explain (Just (vals,evs)) [])
            Unknown -> do
              explain <- explainFailure Nothing p
              let locs = assumptionsTopLevelLocs assumptionsInScope
              return (locs, NotProved explain Nothing [])
        where
          -- Create a handle for a file based on the template. For instance, if
          -- the template is @solver-output.smt2@, then create a file named
          -- @solver-output-<goal number>-<solver name>.smt2@.
          withOfflineSolverOutputHandle :: FilePath -> (Handle -> IO r) -> IO r
          withOfflineSolverOutputHandle template k =
            let (templateName, templateExt) = splitExtension template
                fileName = templateName ++ "-" ++ show goalNumber
                                        ++ "-" ++ WS.solver_adapter_name adapter
                                       <.> templateExt in
            withFile fileName WriteMode k
  
          -- Lift @withOfflineSolverOutputHandle@ to a @Maybe FilePath@ argument.
          withMaybeOfflineSolverOutputHandle ::
            Maybe FilePath -> (Maybe Handle -> IO r) -> IO r
          withMaybeOfflineSolverOutputHandle mbTemplate k =
            case mbTemplate of
              Just template -> withOfflineSolverOutputHandle template (k . Just)
              Nothing       -> k Nothing
  
  evalModelFromEvents :: [CrucibleEvent GroundValueWrapper] -> ModelView
  evalModelFromEvents evs = ModelView (foldl f (modelVals emptyModelView) evs)
   where
     f m (CreateVariableEvent loc nm tpr (GVW v)) = MapF.insertWith jn tpr (Vals [Entry nm loc v]) m
     f m _ = m
  
     jn (Vals new) (Vals old) = Vals (old++new)
  
  dispatchSolversOnGoalAsync :: forall a b s time.
                                (RealFrac time)
                             => Maybe time
                             -> [WS.SolverAdapter s]
                             -> (WS.SolverAdapter s -> IO (b,ProofResult a))
                             -> IO (Either [SomeException] (Maybe (b,ProofResult a)))
  dispatchSolversOnGoalAsync mtimeoutSeconds adapters withAdapter = do
    asyncs <- forM adapters (async . withAdapter)
    await asyncs []
    where
      await [] es = return $ Left es
      await as es = do
        mresult <- withTimeout $ waitAnyCatch as
        case mresult of
          Just (a, result) -> do
            let as' = filter (/= a) as
            case result of
              Left  exc ->
                await as' (exc : es)
              Right (x, r@(Proved _)) -> do
                mapM_ kill as'
                return $ Right (Just (x,r))
              Right (x,r@(NotProved _ (Just _) _)) -> do
                mapM_ kill as'
                return $ Right (Just (x,r))
              Right _ ->
                await as' es
          Nothing -> do
            mapM_ kill as
            return $ Right $ Nothing
  
      withTimeout action
        | Just seconds <- mtimeoutSeconds = ST.timeout (round seconds * 1000000) action
        | otherwise = Just <$> action
  
      -- `cancel` from async blocks until the canceled thread has terminated.
      kill a = throwTo (asyncThreadId a) ExitSuccess
  
  
  -- | Returns three actions, called respectively when the proving process for a
  -- goal is started, when it is ended, and when the final goal is solved.  These
  -- handlers should handle all necessary output / notifications to external
  -- observers, based on the run options.
  proverMilestoneCallbacks ::
    Log.Logs msgs =>
    Log.SupportsCruxLogMessage msgs =>
    Goals asmp ast -> IO ProverMilestoneCallbacks
  proverMilestoneCallbacks goals = do
    (start, end, finish) <-
      if view quiet ?outputConfig then
        return silentProverMilestoneCallbacks
      else
        prepStatus "Checking: " (countGoals goals)
    return
      ( start <> sayCrux . Log.StartedGoal
      , end <> sayCrux . Log.EndedGoal
      , finish
      )
  
  
  -- | Prove a collection of goals.  The result is a goal tree, where
  -- each goal is annotated with the outcome of the proof.
  --
  -- NOTE: This function takes an explicit symbolic backend as an argument, even
  -- though the symbolic backend used for symbolic execution is available in the
  -- 'SimCtxt'.  We do that so that we can use separate solvers for path
  -- satisfiability checking and goal discharge.
  proveGoalsOnline ::
    forall sym personality p msgs goalSolver s st fs.
    ( sym ~ ExprBuilder s st fs
    , OnlineSolver goalSolver
    , Logs msgs
    , SupportsCruxLogMessage msgs
    ) =>
    OnlineBackend goalSolver s st fs ->
    CruxOptions ->
    SimCtxt personality sym p ->
    (Maybe (GroundEvalFn s) -> Assertion sym -> IO (Doc Void)) ->
    Maybe (Goals (Assumptions sym) (Assertion sym)) ->
    IO (ProcessedGoals, Maybe (Goals (Assumptions sym) (Assertion sym, [ProgramLoc], ProofResult sym)))
  proveGoalsOnline _ _opts _ctxt _explainFailure Nothing =
       return (ProcessedGoals 0 0 0 0, Nothing)
  
  proveGoalsOnline bak opts _ctxt explainFailure (Just gs0) =
    do
       -- send solver interactions to the correct file
       mapM_ (symCfg sym solverInteractionFile) (fmap Text.pack (onlineSolverOutput opts))
       -- initial goal count
       goalNum <- newIORef (ProcessedGoals 0 0 0 0)
       -- nameMap is a mutable ref to a map from Text to Either (Assumption sym) (Assertion sym)
       nameMap <- newIORef Map.empty
       when (unsatCores opts && yicesMCSat opts) $
         sayCrux Log.SkippingUnsatCoresBecauseMCSatEnabled
       -- callbacks for starting a goal, ending a goal, and finishing all goals
       (start,end,finish) <- proverMilestoneCallbacks gs0
       -- make sure online features are enabled
       enableOpt <- getOptionSetting enableOnlineBackend (getConfiguration sym)
       _ <- setOpt enableOpt True
       -- @go@ traverses a proof tree, processing/solving each goal as it traverses it.
       -- It also updates goal count and nameMap
       res <- withSolverProcess bak (panic "proveGoalsOnline" ["Online solving not enabled!"]) $ \sp ->
                inNewFrame sp (go (start,end) sp mempty goalNum gs0 nameMap)
       nms <- readIORef goalNum
       finish
       return (nms, Just res)
  
    where
    sym = backendGetSym bak
  
    bindName nm p nameMap = modifyIORef nameMap (Map.insert nm p)
  
    hasUnsatCores = unsatCores opts && not (yicesMCSat opts)
  
    howManyAbducts = fromMaybe 0 (getNAbducts opts)
  
    usingAbducts = howManyAbducts > 0
  
    failfast = proofGoalsFailFast opts
  
    go (start,end) sp assumptionsInScope gn gs nameMap = do
      -- traverse goal tree
      case gs of
        -- case: assumption in context for all the contained goals
        Assuming asms gs1 ->
          do ps <- flattenAssumptions sym asms
             forM_ ps $ \asm ->
               unless (trivialAssumption asm) $
                 -- extract predicate from assumption
                 do let p = assumptionPred asm
                    -- create formula, assert to SMT solver, create new name and add to nameMap
                    nm <- doAssume =<< mkFormula conn p
                    bindName nm (Left asm) nameMap
             -- recursive call
             res <- go (start,end) sp (assumptionsInScope <> asms) gn gs1 nameMap
             return (Assuming (mconcat (map singleAssumption ps)) res)
        -- case: proof obligation in the context of all previously-made assumptions
        Prove p ->
          -- number of processed goals gives goal number to prove
          do goalNumber <- totalProcessedGoals <$> readIORef gn
             start goalNumber
             -- negate goal, create formula
             t <- mkFormula conn =<< notPred sym (p ^. labeledPred)
             -- assert formula to SMT solver, create new name and add to nameMap.
             -- This is done in a new assertion frame if abduction is turned on since
             -- this assertion would need to be removed before asking for abducts
             let inNewFrame2ForAbducts =
                   if usingAbducts then inNewFrame2 sp else id
             ret <- inNewFrame2ForAbducts $ do
               nm <- doAssume t
               bindName nm (Right p) nameMap
               -- check-sat with SMT solver, pattern match on result
               res <- check sp "proof"
               case res of
                 Unsat () ->
                   -- build unsat core, which is the entire assertion set by default
                   do namemap <- readIORef nameMap
                      core <- if hasUnsatCores then
                                 map (lookupnm namemap) <$> getUnsatCore sp
                              -- default unsat core: entire assertion set
                              else return (Map.elems namemap)
                      let locs = assumptionsTopLevelLocs assumptionsInScope
                      return $ UnsatResult core locs
                 Sat ()  ->
                   do -- evaluate counter-example
                      f <- smtExprGroundEvalFn conn (solverEvalFuns sp)
                      evs <- concretizeEvents (groundEval f) assumptionsInScope
                      explain <- explainFailure (Just f) p
                      return $ SatResult explain evs
                 Unknown ->
                   do explain <- explainFailure Nothing p
                      let locs = assumptionsTopLevelLocs assumptionsInScope
                      return $ UnknownResult explain locs
             end goalNumber
             smtResultToGoals p ret
        -- case: conjunction of goals
        ProveConj g1 g2 ->
          do g1' <- inNewFrame sp (go (start,end) sp assumptionsInScope gn g1 nameMap)
             -- NB, we don't need 'inNewFrame' here because
             --  we don't need to back up to this point again.
             if failfast then
               do numDisproved <- disprovedGoals <$> readIORef gn
                  if numDisproved > 0 then
                    return g1'
                  else
                    ProveConj g1' <$> go (start,end) sp assumptionsInScope gn g2 nameMap
             else
               ProveConj g1' <$> go (start,end) sp assumptionsInScope gn g2 nameMap
  
      where
      conn = solverConn sp
  
      lookupnm namemap x =
        fromMaybe (error $ "Named predicate " ++ show x ++ " not found!")
                  (Map.lookup x namemap)
  
      doAssume formula = do
        namemap <- readIORef nameMap
        if hasUnsatCores
        then assumeFormulaWithFreshName conn formula
        else assumeFormula conn formula >> return (Text.pack ("x" ++ show (Map.size namemap)))
  
      -- Convert an 'SMTResult' to a 'ProofResult'. This function should be
      -- called /after/ 'inNewFrame2' so that the abducts can be queried properly
      -- in the SatResult case.
      smtResultToGoals :: LabeledPred (BoolExpr s) SimError
                       -> SMTResult sym
                       -> IO ( Goals asmp (LabeledPred (BoolExpr s) SimError
                             , [ProgramLoc]
                             , ProofResult sym)
                             )
      smtResultToGoals p smtRes = do
        (locs, gt) <- case smtRes of
          UnsatResult core locs -> do
            let pr = Proved core
            return (locs, pr)
          SatResult explain evs -> do
            let vals = evalModelFromEvents evs
            abds <- if usingAbducts then
                      getAbducts sp (fromIntegral howManyAbducts) "abd" (p ^. labeledPred)
                    else
                      return []
            let gt = NotProved explain (Just (vals,evs)) abds
            when (failfast && not (isResourceExhausted p)) $
              sayCrux Log.FoundCounterExample
            let locs = map eventLoc evs
            return (locs, gt)
          UnknownResult explain locs -> do
            let gt = NotProved explain Nothing []
            return (locs, gt)
  
        -- update goal count
        modifyIORef' gn (updateProcessedGoals p gt)
        return (Prove (p, locs, gt))
  
  -- | Like 'inNewFrame', but specifically for frame @2@. This is used for the
  -- purpose of generating abducts.
  
  -- TODO: Upstream this to @what4@ (Issue what4#218).
  inNewFrame2 :: SMTReadWriter solver => SolverProcess scope solver -> IO a -> IO a
  inNewFrame2 sp action = do
    inNewFrame2Open sp
    val <- action
    inNewFrame2Close sp
    return val
  
  -- | An intermediate data structure used in 'proveGoalsOnline'. This can be
  -- thought of as a halfway point between a 'SatResult' and a 'ProofResult'.
  data SMTResult sym
    = UnsatResult [Either (Assumption sym) (Assertion sym)]
                  [ProgramLoc]
    | SatResult (Doc Void)
                [CrucibleEvent GroundValueWrapper]
    | UnknownResult (Doc Void)
                    [ProgramLoc]
  