module Ntha.Type.Refined where

import           Ntha.Core.Ast
import           Ntha.Type.Type
import           Ntha.Type.TypeScope
import           Ntha.Z3.Assertion
import           Ntha.Z3.Class
import           Ntha.Z3.Context
import           Ntha.Z3.Logic

import           Z3.Monad
import           Prelude                hiding (lookup)
import           Control.Monad          (mapM_)
import           Control.Monad.IO.Class (liftIO)

genPred :: Term -> Z3Pred
genPred term = case term of
  TmLT t1 t2 -> PAssert $ Less t1 t2
  TmGT t1 t2 -> PAssert $ Greater t1 t2
  TmLE t1 t2 -> PAssert $ LessE t1 t2
  TmGE t1 t2 -> PAssert $ GreaterE t1 t2
  TmEqual t1 t2 -> PAssert $ Equal t1 t2
  TmAnd t1 t2 -> PConj (genPred t1) (genPred t2)
  TmOr t1 t2 -> PDisj (genPred t1) (genPred t2)
  TmNot t -> PNeg (genPred t)
  _ -> error $ "not support term: " ++ show term

replaceRtnTerm :: String -> Term -> Term -> Term
replaceRtnTerm rtnName rtnTerm predTerm = case predTerm of
  TmVar n -> if n == rtnName then rtnTerm else predTerm
  TmNum _ -> predTerm
  TmLT t1 t2 -> TmLT (replaceRtnTerm' t1) (replaceRtnTerm' t2)
  TmGT t1 t2 -> TmGT (replaceRtnTerm' t1) (replaceRtnTerm' t2)
  TmLE t1 t2 -> TmLE (replaceRtnTerm' t1) (replaceRtnTerm' t2)
  TmGE t1 t2 -> TmGE (replaceRtnTerm' t1) (replaceRtnTerm' t2)
  TmSub t1 t2 -> TmSub (replaceRtnTerm' t1) (replaceRtnTerm' t2)
  TmAdd t1 t2 -> TmAdd (replaceRtnTerm' t1) (replaceRtnTerm' t2)
  TmMul t1 t2 -> TmMul (replaceRtnTerm' t1) (replaceRtnTerm' t2)
  TmDiv t1 t2 -> TmDiv (replaceRtnTerm' t1) (replaceRtnTerm' t2)
  TmEqual t1 t2 -> TmEqual (replaceRtnTerm' t1) (replaceRtnTerm' t2)
  TmAnd t1 t2 -> TmAnd (replaceRtnTerm' t1) (replaceRtnTerm' t2)
  TmOr t1 t2 -> TmOr (replaceRtnTerm' t1) (replaceRtnTerm' t2)
  TmNot t -> TmNot (replaceRtnTerm' t)
  TmIf t1 t2 t3 -> TmIf (replaceRtnTerm' t1) (replaceRtnTerm' t2) (replaceRtnTerm' t3)
  where replaceRtnTerm' = replaceRtnTerm rtnName rtnTerm

genRtnPred :: String -> Term -> Term -> Z3Pred
-- use neg to find counterexamples
genRtnPred rtnName rtnTerm = PNeg . genPred . (replaceRtnTerm rtnName rtnTerm)

convertProg' :: Expr -> Term
convertProg' expr = case expr of
  ENum n -> TmNum n
  EVar name -> TmVar name
  EApp fn arg ->
    case fn of
      EApp (EVar op) arg' -> opConstruct argTerm' argTerm
        where argTerm' = convertProg' arg'
              argTerm = convertProg' arg
              opConstruct =
                case op of
                  "+" -> TmAdd
                  "-" -> TmSub
                  "*" -> TmMul
                  "/" -> TmDiv
                  "<" -> TmLT
                  ">" -> TmGT
                  "≤" -> TmLE
                  "≥" -> TmGE
                  "=" -> TmEqual
                  "∧" -> TmAnd
                  "∨" -> TmOr
                  _ -> error $ "not support op: " ++ op
      EVar op -> case op of
                   "¬" -> let argTerm = convertProg' arg
                         in TmNot argTerm
                   _ -> error $ "not support op: " ++ op
      _ -> error $ "not support fn: " ++ show fn
  EIf cond (thenInstruction:[]) (elseInstruction:[]) -> TmIf condTerm thenTerm elseTerm
    where condTerm = convertProg' cond
          thenTerm = convertProg' thenInstruction
          elseTerm = convertProg' elseInstruction
  _ -> error $ "not support expr: " ++ show expr

convertProg :: Expr -> TypeScope -> IO Z3Pred
convertProg expr scope = case expr of
  only support exists and  for now
  EDestructLetBinding main args (instruction:[]) -> do
    let name = case main of
                 IdPattern n -> n ++ "-sig"
                 _ -> ""
    let typeSig = lookup name scope
    let argNames = map (\pat -> case pat of
                                  IdPattern n -> n
                                  _ -> show pat)
                       args
    case typeSig of
      Just (TSig ta) -> do
        let terms = extractTerm ta
        let predNames = getPredNames ta
        case predNames of
          -- (¬ ⊥) always satisfied
          [] -> return PFalse
          _ -> case (argNames, terms) of
                ([n], [rtnTerm']) ->
                  return $ PExists n RTInt $ genRtnPred' rtnTerm'
                ([n1, n2], [rtnTerm']) ->
                  return $ PExists2 n1 n2 RTInt $ genRtnPred' rtnTerm'
                ([n], [argTerm, rtnTerm']) ->
                  return $ PExists n RTInt $ PConj (genPred argTerm) $ genRtnPred' rtnTerm'
                ([n1, n2], [argTerm1, argTerm2, rtnTerm']) ->
                  return $ PExists2 n1 n2 RTInt $ PConj (PConj (genPred argTerm1) $ genPred argTerm2) $ genRtnPred' rtnTerm'
                _ -> error $ "not support args: " ++ show argNames ++ " and terms: " ++ show terms
              where rtnName = last predNames
                    rtnTerm = convertProg' instruction
                    genRtnPred' :: Term -> Z3Pred
                    genRtnPred' = genRtnPred rtnName rtnTerm
      -- (¬ ⊥) always satisfied
      _ -> return PFalse
  EProgram (instruction:_) -> convertProg instruction scope
  _ -> error $ "not support expr: " ++ show expr

checkPre :: Z3Pred -> Z3SMT () (Result, Maybe Model)
checkPre pre = local $ do
  ast <- encode pre
  local (assert ast >> getModel)

checker :: Expr -> TypeScope -> IO ()
checker expr scope = case expr of
  EDestructLetBinding _ _ _ -> do
    progPred <- convertProg expr scope
    -- trade off
    let adts = [("", [("", [("", RTInt)])])]
    ret <- runSMT adts () $ do
             (r, _mm) <- checkPre progPred
             case r of
               Unsat -> do
                   core <- getUnsatCore
                   liftIO $ sequence_ (map print core)
                   return r
               other -> return other
    if ret == Right Unsat
    then return ()
    else error "refined type check failed"
  EProgram instructions -> mapM_ (\instr -> checker instr scope) instructions
  _ -> return ()

 use neg to find counterexamples
 (¬ ⊥) always satisfied
 (¬ ⊥) always satisfied
 trade off
module Ntha.Type.Refined where

import           Ntha.Core.Ast
import           Ntha.Type.Type
import           Ntha.Type.TypeScope
import           Ntha.Z3.Assertion
import           Ntha.Z3.Class
import           Ntha.Z3.Context
import           Ntha.Z3.Logic

import           Z3.Monad
import           Prelude                hiding (lookup)
import           Control.Monad          (mapM_)
import           Control.Monad.IO.Class (liftIO)

genPred :: Term -> Z3Pred
genPred term = case term of
  TmLT t1 t2 -> PAssert $ Less t1 t2
  TmGT t1 t2 -> PAssert $ Greater t1 t2
  TmLE t1 t2 -> PAssert $ LessE t1 t2
  TmGE t1 t2 -> PAssert $ GreaterE t1 t2
  TmEqual t1 t2 -> PAssert $ Equal t1 t2
  TmAnd t1 t2 -> PConj (genPred t1) (genPred t2)
  TmOr t1 t2 -> PDisj (genPred t1) (genPred t2)
  TmNot t -> PNeg (genPred t)
  _ -> error $ "not support term: " ++ show term

replaceRtnTerm :: String -> Term -> Term -> Term
replaceRtnTerm rtnName rtnTerm predTerm = case predTerm of
  TmVar n -> if n == rtnName then rtnTerm else predTerm
  TmNum _ -> predTerm
  TmLT t1 t2 -> TmLT (replaceRtnTerm' t1) (replaceRtnTerm' t2)
  TmGT t1 t2 -> TmGT (replaceRtnTerm' t1) (replaceRtnTerm' t2)
  TmLE t1 t2 -> TmLE (replaceRtnTerm' t1) (replaceRtnTerm' t2)
  TmGE t1 t2 -> TmGE (replaceRtnTerm' t1) (replaceRtnTerm' t2)
  TmSub t1 t2 -> TmSub (replaceRtnTerm' t1) (replaceRtnTerm' t2)
  TmAdd t1 t2 -> TmAdd (replaceRtnTerm' t1) (replaceRtnTerm' t2)
  TmMul t1 t2 -> TmMul (replaceRtnTerm' t1) (replaceRtnTerm' t2)
  TmDiv t1 t2 -> TmDiv (replaceRtnTerm' t1) (replaceRtnTerm' t2)
  TmEqual t1 t2 -> TmEqual (replaceRtnTerm' t1) (replaceRtnTerm' t2)
  TmAnd t1 t2 -> TmAnd (replaceRtnTerm' t1) (replaceRtnTerm' t2)
  TmOr t1 t2 -> TmOr (replaceRtnTerm' t1) (replaceRtnTerm' t2)
  TmNot t -> TmNot (replaceRtnTerm' t)
  TmIf t1 t2 t3 -> TmIf (replaceRtnTerm' t1) (replaceRtnTerm' t2) (replaceRtnTerm' t3)
  where replaceRtnTerm' = replaceRtnTerm rtnName rtnTerm

genRtnPred :: String -> Term -> Term -> Z3Pred
genRtnPred rtnName rtnTerm = PNeg . genPred . (replaceRtnTerm rtnName rtnTerm)

convertProg' :: Expr -> Term
convertProg' expr = case expr of
  ENum n -> TmNum n
  EVar name -> TmVar name
  EApp fn arg ->
    case fn of
      EApp (EVar op) arg' -> opConstruct argTerm' argTerm
        where argTerm' = convertProg' arg'
              argTerm = convertProg' arg
              opConstruct =
                case op of
                  "+" -> TmAdd
                  "-" -> TmSub
                  "*" -> TmMul
                  "/" -> TmDiv
                  "<" -> TmLT
                  ">" -> TmGT
                  "≤" -> TmLE
                  "≥" -> TmGE
                  "=" -> TmEqual
                  "∧" -> TmAnd
                  "∨" -> TmOr
                  _ -> error $ "not support op: " ++ op
      EVar op -> case op of
                   "¬" -> let argTerm = convertProg' arg
                         in TmNot argTerm
                   _ -> error $ "not support op: " ++ op
      _ -> error $ "not support fn: " ++ show fn
  EIf cond (thenInstruction:[]) (elseInstruction:[]) -> TmIf condTerm thenTerm elseTerm
    where condTerm = convertProg' cond
          thenTerm = convertProg' thenInstruction
          elseTerm = convertProg' elseInstruction
  _ -> error $ "not support expr: " ++ show expr

convertProg :: Expr -> TypeScope -> IO Z3Pred
convertProg expr scope = case expr of
  only support exists and  for now
  EDestructLetBinding main args (instruction:[]) -> do
    let name = case main of
                 IdPattern n -> n ++ "-sig"
                 _ -> ""
    let typeSig = lookup name scope
    let argNames = map (\pat -> case pat of
                                  IdPattern n -> n
                                  _ -> show pat)
                       args
    case typeSig of
      Just (TSig ta) -> do
        let terms = extractTerm ta
        let predNames = getPredNames ta
        case predNames of
          [] -> return PFalse
          _ -> case (argNames, terms) of
                ([n], [rtnTerm']) ->
                  return $ PExists n RTInt $ genRtnPred' rtnTerm'
                ([n1, n2], [rtnTerm']) ->
                  return $ PExists2 n1 n2 RTInt $ genRtnPred' rtnTerm'
                ([n], [argTerm, rtnTerm']) ->
                  return $ PExists n RTInt $ PConj (genPred argTerm) $ genRtnPred' rtnTerm'
                ([n1, n2], [argTerm1, argTerm2, rtnTerm']) ->
                  return $ PExists2 n1 n2 RTInt $ PConj (PConj (genPred argTerm1) $ genPred argTerm2) $ genRtnPred' rtnTerm'
                _ -> error $ "not support args: " ++ show argNames ++ " and terms: " ++ show terms
              where rtnName = last predNames
                    rtnTerm = convertProg' instruction
                    genRtnPred' :: Term -> Z3Pred
                    genRtnPred' = genRtnPred rtnName rtnTerm
      _ -> return PFalse
  EProgram (instruction:_) -> convertProg instruction scope
  _ -> error $ "not support expr: " ++ show expr

checkPre :: Z3Pred -> Z3SMT () (Result, Maybe Model)
checkPre pre = local $ do
  ast <- encode pre
  local (assert ast >> getModel)

checker :: Expr -> TypeScope -> IO ()
checker expr scope = case expr of
  EDestructLetBinding _ _ _ -> do
    progPred <- convertProg expr scope
    let adts = [("", [("", [("", RTInt)])])]
    ret <- runSMT adts () $ do
             (r, _mm) <- checkPre progPred
             case r of
               Unsat -> do
                   core <- getUnsatCore
                   liftIO $ sequence_ (map print core)
                   return r
               other -> return other
    if ret == Right Unsat
    then return ()
    else error "refined type check failed"
  EProgram instructions -> mapM_ (\instr -> checker instr scope) instructions
  _ -> return ()
