content: {-# LANGUAGE DataKinds #-}
  {-# LANGUAGE GADTs #-}
  {-# LANGUAGE LambdaCase #-}
  {-# LANGUAGE RankNTypes #-}
  {-# LANGUAGE TypeApplications #-}
  {-# LANGUAGE TypeOperators #-}
  
  import           ProbeSolvers
  import           Test.Tasty
  import           Test.Tasty.ExpectedFailure
  import           Test.Tasty.HUnit
  
  import           Data.Maybe
  import           System.Environment
  
  import           Data.Parameterized.Context
  import           Data.Parameterized.Map (MapF)
  import           Data.Parameterized.Nonce
  
  import           What4.Config
  import           What4.Expr
  import           What4.Interface
  import           What4.SatResult
  import           What4.Solver.Adapter
  import qualified What4.Solver.CVC5 as CVC5
  import qualified What4.Solver.Z3 as Z3
  
  type SimpleExprBuilder t fs = ExprBuilder t EmptyExprBuilderState fs
  
  logData :: LogData
  logData = defaultLogData { logCallbackVerbose = (\_ -> putStrLn) }
  
  withSym :: FloatModeRepr fm -> (forall t . SimpleExprBuilder t (Flags fm) -> IO a) -> IO a
  withSym float_mode action = withIONonceGenerator $ \gen -> do
    sym <- newExprBuilder float_mode EmptyExprBuilderState gen
    extendConfig CVC5.cvc5Options (getConfiguration sym)
    extendConfig Z3.z3Options (getConfiguration sym)
    action sym
  
  intProblem :: IsSymExprBuilder sym => sym -> IO ([SomeSymFn sym], [Pred sym], Pred sym)
  intProblem sym = do
    inv <- freshTotalUninterpFn sym (safeSymbol "inv") knownRepr knownRepr
    i <- freshConstant sym (safeSymbol "i") knownRepr
    n <- freshConstant sym (safeSymbol "n") knownRepr
    zero <- intLit sym 0
    one <- intLit sym 1
    lt_1_n <- intLt sym one n
    inv_0_n <- applySymFn sym inv $ Empty :> zero :> n
    -- 1 < n ==> inv(0, n)
    impl0 <- impliesPred sym lt_1_n inv_0_n
    inv_i_n <- applySymFn sym inv $ Empty :> i :> n
    add_i_1 <- intAdd sym i one
    lt_add_i_1_n <- intLt sym add_i_1 n
    conj0 <- andPred sym inv_i_n lt_add_i_1_n
    inv_add_i_1_n <- applySymFn sym inv $ Empty :> add_i_1 :> n
    -- inv(i, n) /\ i+1 < n ==> inv(i+1, n)
    impl1 <- impliesPred sym conj0 inv_add_i_1_n
    le_0_i <- intLe sym zero i
    lt_i_n <- intLt sym i n
    conj1 <- andPred sym le_0_i lt_i_n
    -- inv(i, n) ==> 0 <= i /\ i < n
    impl2 <- impliesPred sym inv_i_n conj1
  
    -- inv(i, n) /\ not (i + 1 < n) ==> i + 1 == n
    not_lt_add_i_1_n <- notPred sym lt_add_i_1_n
    conj2 <- andPred sym inv_i_n not_lt_add_i_1_n
    eq_add_i_1_n <- intEq sym add_i_1 n
    impl3 <- notPred sym =<< impliesPred sym conj2 eq_add_i_1_n
  
    return ([SomeSymFn inv], [impl0, impl1, impl2], impl3)
  
  bvProblem :: IsSymExprBuilder sym => sym -> IO ([SomeSymFn sym], [Pred sym], Pred sym)
  bvProblem sym = do
    inv <- freshTotalUninterpFn sym (safeSymbol "inv") knownRepr knownRepr
    i <- freshConstant sym (safeSymbol "i") $ BaseBVRepr $ knownNat @64
    n <- freshConstant sym (safeSymbol "n") knownRepr
    zero <- bvZero sym knownNat
    one <- bvOne sym knownNat
    ult_1_n <- bvUlt sym one n
    inv_0_n <- applySymFn sym inv $ Empty :> zero :> n
    -- 1 < n ==> inv(0, n)
    impl0 <- impliesPred sym ult_1_n inv_0_n
    inv_i_n <- applySymFn sym inv $ Empty :> i :> n
    add_i_1 <- bvAdd sym i one
    ult_add_i_1_n <- bvUlt sym add_i_1 n
    conj0 <- andPred sym inv_i_n ult_add_i_1_n
    inv_add_i_1_n <- applySymFn sym inv $ Empty :> add_i_1 :> n
    -- inv(i, n) /\ i+1 < n ==> inv(i+1, n)
    impl1 <- impliesPred sym conj0 inv_add_i_1_n
    ule_0_i <- bvUle sym zero i -- trivially true, here for similarity with int test
    ult_i_n <- bvUlt sym i n
    conj1 <- andPred sym ule_0_i ult_i_n
    -- inv(i, n) ==> 0 <= i /\ i < n
    impl2 <- impliesPred sym inv_i_n conj1
  
    -- inv(i, n) /\ not (i + 1 < n) ==> i + 1 == n
    not_ult_add_i_1_n <- notPred sym ult_add_i_1_n
    conj2 <- andPred sym inv_i_n not_ult_add_i_1_n
    eq_add_i_1_n <- bvEq sym add_i_1 n
    impl3 <- notPred sym =<< impliesPred sym conj2 eq_add_i_1_n
  
    return ([SomeSymFn inv], [impl0, impl1, impl2], impl3)
  
  synthesis_test ::
    String ->
    (forall sym . IsSymExprBuilder sym => sym -> IO ([SomeSymFn sym], [Pred sym], Pred sym)) ->
    String ->
    (forall sym t fs .
      sym ~ SimpleExprBuilder t fs =>
      sym ->
      LogData ->
      [SomeSymFn sym] ->
      [BoolExpr t] ->
      IO (SatResult (MapF (SymFnWrapper sym) (SymFnWrapper sym)) ())) ->
    (forall t fs a .
      SimpleExprBuilder t fs ->
      LogData ->
      [BoolExpr t] ->
      (SatResult (GroundEvalFn t, Maybe (ExprRangeBindings t)) () -> IO a) ->
      IO a) ->
    TestTree
  synthesis_test test_name synthesis_problem solver_name run_solver_synthesis run_solver_in_override =
    testCase (test_name ++ " " ++ solver_name ++ " test") $ withSym FloatIEEERepr $ \sym -> do
      (synth_fns, constraints, goal) <- synthesis_problem sym
  
      run_solver_in_override sym logData [goal] $ \res -> isSat res @? "sat"
  
      subst <- run_solver_synthesis sym logData synth_fns constraints >>= \case
        Sat res -> return res
        Unsat{} -> fail "Infeasible"
        Unknown -> fail "Fail"
  
      goal' <- substituteSymFns sym subst goal
      run_solver_in_override sym logData [goal'] $ \res -> isUnsat res @? "unsat"
  
  main :: IO ()
  main = do
    testLevel <- TestLevel . fromMaybe "0" <$> lookupEnv "CI_TEST_LEVEL"
    let solverNames = map SolverName [ "cvc5", "z3" ]
    solvers <- reportSolverVersions testLevel id
      =<< (zip solverNames <$> mapM getSolverVersion solverNames)
    let skipPre4_8_9 why =
          let shouldSkip = case lookup (SolverName "z3") solvers of
                Just (SolverVersion v) -> any (`elem` [ "4.8.8" ]) $ words v
                Nothing -> True
          in if shouldSkip then expectFailBecause why else id
        failureZ3 = "failure with older Z3 versions; upgrade to at least 4.8.9"
    defaultMain $ testGroup "Tests" $
      [ synthesis_test "int" intProblem "cvc5" CVC5.runCVC5SyGuS CVC5.runCVC5InOverride
      , skipPre4_8_9 failureZ3 $ synthesis_test "int" intProblem "z3" (\sym -> Z3.runZ3Horn sym False) Z3.runZ3InOverride
      , synthesis_test "bv" bvProblem "cvc5" CVC5.runCVC5SyGuS CVC5.runCVC5InOverride
      , skipPre4_8_9 failureZ3 $ synthesis_test "bv->int" bvProblem  "z3" (\sym -> Z3.runZ3Horn sym True) Z3.runZ3InOverride
      ]
  