------------------------------------------------------------------------
-- |
-- Module      : What4.Solver.Z3
-- Description : Solver adapter code for Z3
-- Copyright   : (c) Galois, Inc 2015-2020
-- License     : BSD3
-- Maintainer  : Rob Dockins <rdockins@galois.com>
-- Stability   : provisional
--
-- Z3-specific tweaks to the basic SMTLib2 solver interface.
------------------------------------------------------------------------
{-# LANGUAGE FlexibleInstances #-}
{-# LANGUAGE LambdaCase #-}
{-# LANGUAGE MultiParamTypeClasses #-}
{-# LANGUAGE MultiWayIf #-}
{-# LANGUAGE OverloadedStrings #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE TypeApplications #-}
{-# LANGUAGE TypeOperators #-}
{-# LANGUAGE GADTs #-}
module What4.Solver.Z3
  ( Z3(..)
  , z3Adapter
  , z3Path
  , z3Timeout
  , z3Options
  , z3Tactic
  , z3TacticDefault
  , z3Features
  , runZ3InOverride
  , withZ3
  , writeZ3SMT2File
  , runZ3Horn
  , writeZ3HornSMT2File
  ) where

import           Control.Monad ( when )
import qualified Data.Bimap as Bimap
import           Data.Bits
import           Data.Foldable
import           Data.Map.Strict (Map)
import qualified Data.Map as Map
import           Data.String
import           Data.Text (Text)
import qualified Data.Text as T
import           System.IO

import           Data.Parameterized.Map (MapF)
import qualified Data.Parameterized.Map as MapF
import           Data.Parameterized.Some
import           What4.BaseTypes
import           What4.Concrete
import           What4.Config
import           What4.Expr.Builder
import           What4.Expr.GroundEval
import           What4.Interface
import           What4.ProblemFeatures
import           What4.Protocol.Online
import qualified What4.Protocol.SMTLib2 as SMT2
import           What4.Protocol.SMTLib2.Response ( strictSMTParseOpt )
import qualified What4.Protocol.SMTLib2.Response as RSP
import qualified What4.Protocol.SMTLib2.Syntax as Syntax
import           What4.Protocol.SMTWriter
import           What4.SatResult
import           What4.Solver.Adapter
import           What4.Utils.Process

data Z3 = Z3 deriving Show

-- | Path to Z3
z3Path :: ConfigOption (BaseStringType Unicode)
z3Path = configOption knownRepr "solver.z3.path"

z3PathOLD :: ConfigOption (BaseStringType Unicode)
z3PathOLD = configOption knownRepr "z3_path"

-- | Per-check timeout, in milliseconds (zero is none)
z3Timeout :: ConfigOption BaseIntegerType
z3Timeout = configOption knownRepr "solver.z3.timeout"

z3TimeoutOLD :: ConfigOption BaseIntegerType
z3TimeoutOLD = configOption knownRepr "z3_timeout"
-- | Strict parsing specifically for Z3 interaction?  If set,
-- overrides solver.strict_parsing, otherwise defaults to
-- solver.strict_parsing.
z3StrictParsing  :: ConfigOption BaseBoolType
z3StrictParsing = configOption knownRepr "solver.z3.strict_parsing"

-- | Z3 tactic
z3Tactic :: ConfigOption (BaseStringType Unicode)
z3Tactic = configOption knownRepr "solver.z3.tactic"

z3TacticDefault :: Text
z3TacticDefault = ""

z3Options :: [ConfigDesc]
z3Options =
  let mkPath co = mkOpt co
                  executablePathOptSty
                  (Just "Z3 executable path")
                  (Just (ConcreteString "z3"))
      mkTmo co = mkOpt co
                 integerOptSty
                 (Just "Per-check timeout in milliseconds (zero is none)")
                 (Just (ConcreteInteger 0))
      p = mkPath z3Path
      t = mkTmo z3Timeout
  in [ p, t
     , copyOpt (const $ configOptionText z3StrictParsing) strictSMTParseOpt
     , mkOpt z3Tactic stringOptSty (Just "Z3 tactic") (Just (ConcreteString (UnicodeLiteral z3TacticDefault)))
     , deprecatedOpt [p] $ mkPath z3PathOLD
     , deprecatedOpt [t] $ mkTmo z3TimeoutOLD
     ] <> SMT2.smtlib2Options

z3Adapter :: SolverAdapter st
z3Adapter =
  SolverAdapter
  { solver_adapter_name = "z3"
  , solver_adapter_config_options = z3Options
  , solver_adapter_check_sat = runZ3InOverride
  , solver_adapter_write_smt2 = writeZ3SMT2File
  }

indexType :: [SMT2.Sort] -> SMT2.Sort
indexType [i] = i
indexType il = SMT2.smtlib2StructSort @Z3 il

indexCtor :: [SMT2.Term] -> SMT2.Term
indexCtor [i] = i
indexCtor il = SMT2.smtlib2StructCtor @Z3 il

instance SMT2.SMTLib2Tweaks Z3 where
  smtlib2tweaks = Z3

  smtlib2arrayType il r = SMT2.arraySort (indexType il) r

  smtlib2arrayConstant = Just $ \idx rtp v ->
    SMT2.arrayConst (indexType idx) rtp v
  smtlib2arraySelect a i = SMT2.arraySelect a (indexCtor i)
  smtlib2arrayUpdate a i = SMT2.arrayStore a (indexCtor i)

  -- Z3 uses a datatype declaration command that differs from the
  -- SMTLib 2.6 standard
  smtlib2declareStructCmd n = Just $
      let type_name i = fromString ('T' : show (i-1))
          params = builder_list $ type_name  <$> [1..n]
          n_str = fromString (show n)
          tp = "Struct" <> n_str
          ctor = "mk-struct" <> n_str
          field_def i = app field_nm [type_name i]
            where field_nm = "struct" <> n_str <> "-proj" <> fromString (show (i-1))
          fields = field_def <$> [1..n]
          decl = app tp [app ctor fields]
          decls = "(" <> decl <> ")"
       in Syntax.Cmd $ app "declare-datatypes" [ params, decls ]

z3Features :: ProblemFeatures
z3Features = useNonlinearArithmetic
         .|. useIntegerArithmetic
         .|. useQuantifiers
         .|. useSymbolicArrays
         .|. useStructs
         .|. useStrings
         .|. useFloatingPoint
         .|. useBitvectors

writeZ3SMT2File
   :: ExprBuilder t st fs
   -> Handle
   -> [BoolExpr t]
   -> IO ()
writeZ3SMT2File = SMT2.writeDefaultSMT2 Z3 "Z3" z3Features (Just z3StrictParsing)

instance SMT2.SMTLib2GenericSolver Z3 where
  defaultSolverPath _ = findSolverPath z3Path . getConfiguration

  defaultSolverArgs _ sym = do
    let cfg = getConfiguration sym
    timeout <- getOption =<< getOptionSetting z3Timeout cfg
    let extraOpts = case timeout of
                      Just (ConcreteInteger n) | n > 0 -> ["-t:" ++ show n]
                      _ -> []
    tactic <- getOpt =<< getOptionSetting z3Tactic cfg
    let tacticOpt = if tactic /= z3TacticDefault then ["tactic.default_tactic=" ++ T.unpack tactic] else []
    return $ tacticOpt ++ ["-smt2", "-in"] ++ extraOpts

  getErrorBehavior _ = SMT2.queryErrorBehavior

  defaultFeatures _ = z3Features

  supportsResetAssertions _ = True

  setDefaultLogicAndOptions writer = do
    -- Tell Z3 to produce models.
    SMT2.setOption writer "produce-models" "true"
    -- Tell Z3 to round and print algebraic reals as decimal
    SMT2.setOption writer "pp.decimal" "true"
    -- Tell Z3 to compute UNSAT cores, if that feature is enabled
    when (supportedFeatures writer `hasProblemFeature` useUnsatCores) $
      SMT2.setOption writer "produce-unsat-cores" "true"

runZ3InOverride
  :: ExprBuilder t st fs
  -> LogData
  -> [BoolExpr t]
  -> (SatResult (GroundEvalFn t, Maybe (ExprRangeBindings t)) () -> IO a)
  -> IO a
runZ3InOverride = SMT2.runSolverInOverride Z3 nullAcknowledgementAction
                  z3Features (Just z3StrictParsing)

-- | Run Z3 in a session. Z3 will be configured to produce models, but
-- otherwise left with the default configuration.
withZ3
  :: ExprBuilder t st fs
  -> FilePath
    -- ^ Path to Z3 executable
  -> LogData
  -> (SMT2.Session t Z3 -> IO a)
    -- ^ Action to run
  -> IO a
withZ3 = SMT2.withSolver Z3 nullAcknowledgementAction
         z3Features (Just z3StrictParsing)


setInteractiveLogicAndOptions ::
  SMT2.SMTLib2Tweaks a =>
  WriterConn t (SMT2.Writer a) ->
  IO ()
setInteractiveLogicAndOptions writer = do
    -- Tell Z3 to acknowledge successful commands
    SMT2.setOption writer "print-success"  "true"
    -- Tell Z3 to produce models
    SMT2.setOption writer "produce-models" "true"
    -- Tell Z3 to round and print algebraic reals as decimal
    SMT2.setOption writer "pp.decimal" "true"
    -- Tell Z3 to make declarations global, so they are not removed by 'pop' commands
    SMT2.setOption writer "global-declarations" "true"
    -- Tell Z3 to compute UNSAT cores, if that feature is enabled
    when (supportedFeatures writer `hasProblemFeature` useUnsatCores) $ do
      SMT2.setOption writer "produce-unsat-cores" "true"

instance OnlineSolver (SMT2.Writer Z3) where

  startSolverProcess feat mbIOh sym = do
    timeout <- SolverGoalTimeout <$>
               (getOpt =<< getOptionSetting z3Timeout (getConfiguration sym))
    SMT2.startSolver Z3 SMT2.smtAckResult setInteractiveLogicAndOptions
      timeout feat (Just z3StrictParsing) mbIOh sym

  shutdownSolverProcess = SMT2.shutdownSolver Z3

-- | Check the satisfiability of a set of constrained Horn clauses (CHCs).
--
-- CHCs are represented as pure SMT-LIB2 implications. For more information, see
-- the [Z3 guide](https://microsoft.github.io/z3guide/docs/fixedpoints/intro/).
--
-- There are two ways to solve the CHCs: either by directly solving the problem
-- as is, or by transforming the problem into a set of linear integer arithmetic
-- (LIA) CHCs and solving that instead. The latter is done by replacing all
-- bitvector (BV) operations with LIA operations, and replacing all BV variables
-- with LIA variables. This transformation is not sound, but in practice it is a
-- useful heuristic. Then the result is transformed back into a BV result, and
-- checked for satisfiability. The satisfiability check is necessary because the
-- transformation is not sound, so LIA solution may not be a solution to the BV
-- CHCs.
runZ3Horn ::
  forall sym t st fs .
  sym ~ ExprBuilder t st fs =>
  sym ->
  Bool {- transform the BV CHCs into LIA CHCs -} ->
  LogData ->
  [SomeSymFn sym] ->
  [BoolExpr t] ->
  IO (SatResult (MapF (SymFnWrapper sym) (SymFnWrapper sym)) ())
runZ3Horn sym do_bv_to_lia_transform log_data inv_fns horn_clauses = do
  (lia_inv_fns, lia_horn_clauses, bv_to_lia_fn_subst) <- transformHornClausesForZ3
    sym
    do_bv_to_lia_transform
    inv_fns
    horn_clauses

  logSolverEvent sym
    (SolverStartSATQuery $ SolverStartSATQueryRec
      { satQuerySolverName = show Z3
      , satQueryReason = logReason log_data
      })

  path <- SMT2.defaultSolverPath Z3 sym
  get_value_result <- withZ3 sym path (log_data { logVerbosity = 2 }) $ \session -> do
    writeHornProblem sym (SMT2.sessionWriter session) lia_inv_fns lia_horn_clauses

    check_sat_result <- RSP.getLimitedSolverResponse "check-sat"
      (\case
        RSP.AckSat -> Just $ Sat ()
        RSP.AckUnsat -> Just $ Unsat ()
        RSP.AckUnknown -> Just Unknown
        _ -> Nothing)
      (SMT2.sessionWriter session)
      Syntax.checkSat

    logSolverEvent sym
      (SolverEndSATQuery $ SolverEndSATQueryRec
        { satQueryResult = check_sat_result
        , satQueryError = Nothing
        })

    traverseSatResult
      (\() -> do
        sexp <- RSP.getLimitedSolverResponse "get-value"
          (\case
            RSP.AckSuccessSExp sexp -> Just sexp
            _ -> Nothing)
          (SMT2.sessionWriter session)
          (Syntax.getValue [])
        SMT2.parseFnValues sym (SMT2.sessionWriter session) lia_inv_fns sexp)
      return
      check_sat_result

  let transform_result_lia_to_bv ::
        SatResult (MapF (SymFnWrapper sym) (SymFnWrapper sym)) () ->
        IO (SatResult (MapF (SymFnWrapper sym) (SymFnWrapper sym)) ())
      transform_result_lia_to_bv = \case
        Sat lia_defined_fns -> do
          defined_inv_fns <- MapF.fromList <$> mapM
            (\(SomeSymFn fn) ->
              if| Just (SomeSymFn lia_fn) <- Map.lookup (SomeSymFn fn) bv_to_lia_fn_subst
                , Just (SymFnWrapper lia_defined_fn) <- MapF.lookup (SymFnWrapper lia_fn) lia_defined_fns -> do
                  some_defined_fn <- transformSymFnLIA2BV sym $ SomeSymFn lia_defined_fn
                  case some_defined_fn of
                    SomeSymFn defined_fn
                      | Just Refl <- testEquality (fnArgTypes fn) (fnArgTypes defined_fn)
                      , Just Refl <- testEquality (fnReturnType fn) (fnReturnType defined_fn) ->
                        return $ MapF.Pair (SymFnWrapper fn) (SymFnWrapper defined_fn)
                    _ -> fail $ "runZ3Horn: function type mismatch in solver result: " ++ show fn
                | otherwise -> fail $ "runZ3Horn: function not found in solver result: " ++ show fn)
            inv_fns

          all_unsat <- and <$> mapM
            (\clause -> do
              defined_clause <- notPred sym =<< substituteSymFns sym defined_inv_fns clause
              runZ3InOverride sym (log_data { logVerbosity = 2 }) [defined_clause] $ return . isUnsat)
            horn_clauses

          return $ if all_unsat then Sat defined_inv_fns else Unknown

        _ -> return Unknown

  if do_bv_to_lia_transform then
    transform_result_lia_to_bv get_value_result
  else
    return get_value_result

writeZ3HornSMT2File ::
  sym ~ ExprBuilder t st fs =>
  sym ->
  Bool {- transform the BV CHCs into LIA CHCs -} ->
  Handle ->
  [SomeSymFn sym] ->
  [BoolExpr t] ->
  IO ()
writeZ3HornSMT2File sym do_bv_to_lia_transform h inv_fns horn_clauses = do
  (lia_inv_fns, lia_horn_clauses, _bv_to_lia_fn_subst) <- transformHornClausesForZ3
    sym
    do_bv_to_lia_transform
    inv_fns
    horn_clauses

  writer <- SMT2.defaultFileWriter
    Z3
    (show Z3)
    (SMT2.defaultFeatures Z3)
    (Just z3StrictParsing)
    sym
    h
  SMT2.setDefaultLogicAndOptions writer
  writeHornProblem sym writer lia_inv_fns lia_horn_clauses
  SMT2.writeExit writer

transformHornClausesForZ3 ::
  sym ~ ExprBuilder t st fs =>
  sym ->
  Bool ->
  [SomeSymFn sym] ->
  [BoolExpr t] ->
  IO ([SomeSymFn sym], [BoolExpr t], Map (SomeSymFn sym) (SomeSymFn sym))
transformHornClausesForZ3 sym do_bv_to_lia_transform inv_fns horn_clauses =
  if do_bv_to_lia_transform then do
    (lia_horn_clauses, bv_to_lia_fn_subst) <- transformPredBV2LIA sym horn_clauses
    let lia_inv_fns = Map.elems bv_to_lia_fn_subst
    return (lia_inv_fns, lia_horn_clauses, bv_to_lia_fn_subst)
  else return (inv_fns, horn_clauses, Map.empty)

writeHornProblem ::
  sym ~ ExprBuilder t st fs =>
  sym ->
  WriterConn t (SMT2.Writer Z3) ->
  [SomeSymFn sym] ->
  [BoolExpr t] ->
  IO ()
writeHornProblem sym writer inv_fns horn_clauses = do
  SMT2.setLogic writer Syntax.hornLogic
  implications <- mapM
    (\clause -> foldrM (viewSome $ forallPred sym) clause $ exprUninterpConstants sym clause)
    horn_clauses
  mapM_ (SMT2.assume writer) implications
  SMT2.writeCheckSat writer
  fn_name_bimap <- cacheLookupFnNameBimap writer $ map (\(SomeSymFn fn) -> SomeExprSymFn fn) inv_fns
  SMT2.writeGetValue writer $ map fromText $ Bimap.elems fn_name_bimap
