SMT.hs

-- |
-- Module      :  Cryptol.TypeCheck.Solver.SMT
-- Copyright   :  (c) 2013-2016 Galois, Inc.
-- License     :  BSD3
-- Maintainer  :  cryptol@galois.com
-- Stability   :  provisional
-- Portability :  portable
{-# LANGUAGE Trustworthy #-}
{-# LANGUAGE RecordWildCards #-}
{-# LANGUAGE FlexibleContexts #-}
{-# Language FlexibleInstances #-}
{-# LANGUAGE PatternGuards #-}

module Cryptol.TypeCheck.Solver.SMT
  ( -- * Setup
    Solver
  , SolverConfig
  , withSolver
  , startSolver
  , stopSolver
  , killSolver
  , isNumeric
  , resetSolver

    -- * Debugging
  , debugBlock
  , debugLog

    -- * Proving stuff
  , proveImp
  , checkUnsolvable
  , tryGetModel
  , shrinkModel

    -- * Lower level interactions
  , inNewFrame, TVars, declareVars, assume, unsolvable
  , push, pop
  ) where

import           SimpleSMT (SExpr)
import qualified SimpleSMT as SMT
import           Data.Map ( Map )
import qualified Data.Map as Map
import qualified Data.Set as Set
import           Data.Maybe(catMaybes,isJust)
import           Data.List(partition)
import           Control.Exception
import           Control.Monad(msum,zipWithM,void)
import           Data.Char(isSpace)
import           Text.Read(readMaybe)
import           System.IO(IOMode(..), hClose, openFile)
import qualified System.IO.Strict as StrictIO
import           System.FilePath((</>))
import           System.Directory(doesFileExist)

import Cryptol.Prelude(cryptolTcContents)
import Cryptol.TypeCheck.Type
import Cryptol.TypeCheck.InferTypes
import Cryptol.TypeCheck.Solver.InfNat(Nat'(..))
import Cryptol.TypeCheck.TypePat hiding ((~>),(~~>))
import Cryptol.TypeCheck.Subst(Subst)
import Cryptol.Utils.Panic
import Cryptol.Utils.PP ( Doc, pp )




-- | An SMT solver packed with a logger for debugging.
data Solver = Solver
  { solver    :: SMT.Solver
    -- ^ The actual solver

  , logger    :: SMT.Logger
    -- ^ For debugging
  }

setupSolver :: Solver -> SolverConfig -> IO ()
setupSolver s cfg = do
  _ <- SMT.setOptionMaybe (solver s) ":global-decls" "false"
  SMT.setLogic (solver s) "ALL"
  loadTcPrelude s (solverPreludePath cfg)

-- | Start a fresh solver instance
startSolver :: IO () -> SolverConfig -> IO Solver
startSolver onExit sCfg =
   do let smtFileEnabled = isJust (solverSmtFile sCfg)
      (logger, mbLoggerCloseHdl) <-
        -- There are two scenarios under which we want to explicitly log SMT
        -- solver interactions:
        --
        -- 1. The user wants to debug-print interactions with the `tcDebug`
        --    option
        -- 2. The user wants to write interactions to the `tcSmtFile` option
        --
        -- We enable logging if either one is true.
        if (solverVerbose sCfg) > 0 || smtFileEnabled
        then case solverSmtFile sCfg of
               Nothing ->
                 do logger <- SMT.newLogger 0
                    pure (logger, Nothing)
               Just file ->
                 do loggerHdl <- openFile file WriteMode
                    logger <- SMT.newLoggerWithHandle loggerHdl 0
                    pure (logger, Just (hClose loggerHdl))
        else pure (quietLogger, Nothing)
      let smtDbg = if (solverVerbose sCfg) > 1 || smtFileEnabled
                   then Just logger
                   else Nothing
      solver <- SMT.newSolverWithConfig
                  (SMT.defaultConfig (solverPath sCfg) (solverArgs sCfg))
                    { SMT.solverOnExit =
                        Just $ \_exitCode ->
                        do onExit
                           sequence_ mbLoggerCloseHdl
                    , SMT.solverLogger =
                        maybe SMT.noSolverLogger SMT.smtSolverLogger smtDbg
                    }
      let sol = Solver solver logger
      setupSolver sol sCfg
      return sol

  where
  quietLogger = SMT.Logger { SMT.logMessage = \_ -> return ()
                           , SMT.logLevel   = return 0
                           , SMT.logSetLevel= \_ -> return ()
                           , SMT.logTab     = return ()
                           , SMT.logUntab   = return ()
                           }

-- | Kill the process running the solver
killSolver :: Solver -> IO ()
killSolver s = void $ SMT.forceStop (solver s)

-- | Shut down a solver instance
stopSolver :: Solver -> IO ()
stopSolver s = void $ SMT.stop (solver s)

resetSolver :: Solver -> SolverConfig -> IO ()
resetSolver s sCfg = do
  _ <- SMT.simpleCommand (solver s) ["reset"]
  setupSolver s sCfg

-- | Execute a computation with a fresh solver instance.
withSolver :: IO () -> SolverConfig -> (Solver -> IO a) -> IO a
withSolver onExit cfg = bracket (startSolver onExit cfg) stopSolver

-- | Load the definitions used for type checking.
loadTcPrelude :: Solver -> [FilePath] {- ^ Search in this paths -} -> IO ()
loadTcPrelude s [] = loadString s cryptolTcContents
loadTcPrelude s (p : ps) =
  do let file = p </> "CryptolTC.smt2"
     yes <- doesFileExist file
     if yes then loadFile s file
            else loadTcPrelude s ps


loadFile :: Solver -> FilePath -> IO ()
loadFile s file = loadString s =<< StrictIO.readFile file

loadString :: Solver -> String -> IO ()
loadString s str = go (dropComments str)
  where
  go txt
    | all isSpace txt = return ()
    | otherwise =
      case SMT.readSExpr txt of
        Just (e,rest) -> SMT.command (solver s) e >> go rest
        Nothing       -> panic "loadFile" [ "Failed to parse SMT file."
                                          , txt
                                          ]

  dropComments = unlines . map dropComment . lines
  dropComment xs = case break (== ';') xs of
                     (as,_:_) -> as
                     _ -> xs




--------------------------------------------------------------------------------
-- Debugging


debugBlock :: Solver -> String -> IO a -> IO a
debugBlock s@Solver { .. } name m =
  do debugLog s (";;; " ++ name)
     SMT.logTab logger
     a <- m
     SMT.logUntab logger
     return a

class DebugLog t where
  debugLog :: Solver -> t -> IO ()

  debugLogList :: Solver -> [t] -> IO ()
  debugLogList s ts = case ts of
                        [] -> debugLog s "(none)"
                        _  -> mapM_ (debugLog s) ts

instance DebugLog Char where
  debugLog s x     = SMT.logMessage (logger s) (show x)
  debugLogList s x = SMT.logMessage (logger s) x

instance DebugLog a => DebugLog [a] where
  debugLog = debugLogList

instance DebugLog a => DebugLog (Maybe a) where
  debugLog s x = case x of
                   Nothing -> debugLog s "(nothing)"
                   Just a  -> debugLog s a

instance DebugLog Doc where
  debugLog s x = debugLog s (show x)

instance DebugLog Type where
  debugLog s x = debugLog s (pp x)

instance DebugLog Goal where
  debugLog s x = debugLog s (goal x)

instance DebugLog Subst where
  debugLog s x = debugLog s (pp x)
--------------------------------------------------------------------------------


-- | Returns goals that were not proved
proveImp :: Solver -> [Prop] -> [Goal] -> IO [Goal]
proveImp sol ps gs0 =
  debugBlock sol "PROVE IMP" $
  do let gs1       = concatMap flatGoal gs0
         (gs,rest) = partition (isNumeric . goal) gs1
         numAsmp   = filter isNumeric (concatMap pSplitAnd ps)
         vs        = Set.toList (fvs (numAsmp, map goal gs))
     tvs <- debugBlock sol "VARIABLES" $
       do SMT.push (solver sol)
          Map.fromList <$> zipWithM (declareVar sol) [ 0 .. ] vs
     debugBlock sol "ASSUMPTIONS" $
       mapM_ (assume sol tvs) numAsmp
     gs' <- mapM (prove sol tvs) gs
     SMT.pop (solver sol)
     return (catMaybes gs' ++ rest)

-- | Check if the given goals are known to be unsolvable.
checkUnsolvable :: Solver -> [Goal] -> IO Bool
checkUnsolvable sol gs0 =
  debugBlock sol "CHECK UNSOLVABLE" $
  do let ps = filter isNumeric
            $ map goal
            $ concatMap flatGoal gs0
         vs = Set.toList (fvs ps)
     tvs <- debugBlock sol "VARIABLES" $
       do push sol
          Map.fromList <$> zipWithM (declareVar sol) [ 0 .. ] vs
     ans <- unsolvable sol tvs ps
     pop sol
     return ans

tryGetModel :: Solver -> [TVar] -> [Prop] -> IO (Maybe [(TVar,Nat')])
tryGetModel sol as ps =
  debugBlock sol "TRY GET MODEL" $
  do push sol
     tvs <- Map.fromList <$> zipWithM (declareVar sol) [ 0 .. ] as
     mapM_ (assume sol tvs) ps
     sat <- SMT.check (solver sol)
     su <- case sat of
             SMT.Sat ->
               case as of
                 [] -> return (Just [])
                 _ -> do res <- SMT.getExprs (solver sol) (Map.elems tvs)
                         let parse x = do e <- Map.lookup x tvs
                                          t <- parseNum =<< lookup e res
                                          return (x, t)
                         return (mapM parse as)
             _ -> return Nothing
     pop sol
     return su

  where
  parseNum a
    | SMT.Other s <- a
    , SMT.List [con,val,isFin,isErr] <- s
    , SMT.Atom "mk-infnat" <- con
    , SMT.Atom "false"     <- isErr
    , SMT.Atom fin         <- isFin
    , SMT.Atom v           <- val
    , Just n               <- readMaybe v
    = Just (if fin == "false" then Inf else Nat n)

  parseNum _ = Nothing

shrinkModel :: Solver -> [TVar] -> [Prop] -> [(TVar,Nat')] -> IO [(TVar,Nat')]
shrinkModel sol as ps0 mdl = go [] ps0 mdl
  where
  go done ps ((x,Nat k) : more) =
    do k1 <- shrink1 ps x k
       go ((x,Nat k1) : done) ((tNum k1 >== TVar x) : ps) more

  go done ps ((x,i) : more) = go ((x,i) : done) ps more
  go done _ [] = return done

  shrink1 ps x k
    | k == 0 = return 0
    | otherwise =
      do let k1 = div k 2
             p1 = tNum k1 >== TVar x
         mb <- tryGetModel sol as (p1 : ps)
         case mb of
           Nothing     -> return k
           Just newMdl ->
             case lookup x newMdl of
               Just (Nat k2) -> shrink1 ps x k2
               _ -> panic "shrink" ["model is missing variable", show x]



--------------------------------------------------------------------------------

push :: Solver -> IO ()
push sol = SMT.push (solver sol)

pop :: Solver -> IO ()
pop sol = SMT.pop (solver sol)

inNewFrame :: Solver -> IO a -> IO a
inNewFrame sol m =
  do push sol
     a <- m
     pop sol
     pure a


declareVar :: Solver -> Int -> TVar -> IO (TVar, SExpr)
declareVar s x v =
  do let name = (if isFreeTV v then "fv" else "kv") ++ show x
     e <- SMT.declare (solver s) name cryInfNat
     SMT.assert (solver s) (SMT.fun "cryVar" [ e ])
     return (v,e)


declareVars :: Solver -> [TVar] -> IO TVars
declareVars sol vs =
  Map.fromList <$> zipWithM (declareVar sol) [ 0 .. ]
                                             [ v | v <- vs, kindOf v == KNum ]

assume :: Solver -> TVars -> Prop -> IO ()
assume s tvs p = SMT.assert (solver s) (SMT.fun "cryAssume" [ toSMT tvs p ])

prove :: Solver -> TVars -> Goal -> IO (Maybe Goal)
prove sol tvs g =
  debugBlock sol "PROVE" $
  do let s = solver sol
     push sol
     SMT.assert s (SMT.fun "cryProve" [ toSMT tvs (goal g) ])
     res <- SMT.check s
     pop sol
     case res of
       SMT.Unsat -> return Nothing
       _         -> return (Just g)


-- | Check if some numeric goals are known to be unsolvable.
unsolvable :: Solver -> TVars -> [Prop] -> IO Bool
unsolvable sol tvs ps =
  debugBlock sol "UNSOLVABLE" $
  do SMT.push (solver sol)
     mapM_ (assume sol tvs) ps
     res <- SMT.check (solver sol)
     SMT.pop (solver sol)
     case res of
       SMT.Unsat -> return True
       _         -> return False



--------------------------------------------------------------------------------

-- | Split up the 'And' in a goal
flatGoal :: Goal -> [Goal]
flatGoal g = [ g { goal = p } | p <- pSplitAnd (goal g) ]


-- | Assumes no 'And'
isNumeric :: Prop -> Bool
isNumeric ty = matchDefault False $ msum [ is (|=|), is (|/=|), is (|>=|)
                                         , is aFin, is aPrime ]
  where
  is f = f ty >> return True


--------------------------------------------------------------------------------

type TVars = Map TVar SExpr

cryInfNat :: SExpr
cryInfNat = SMT.const "InfNat"


toSMT :: TVars -> Type -> SExpr
toSMT tvs ty = matchDefault (panic "toSMT" [ "Unexpected type", show ty ])
  $ msum $ map (\f -> f tvs ty)
  [ aInf            ~> "cryInf"
  , aNat            ~> "cryNat"

  , aFin            ~> "cryFin"
  , aPrime          ~> "cryPrime"
  , (|=|)           ~> "cryEq"
  , (|/=|)          ~> "cryNeq"
  , (|>=|)          ~> "cryGeq"
  , aAnd            ~> "cryAnd"
  , aTrue           ~> "cryTrue"

  , anAdd           ~> "cryAdd"
  , (|-|)           ~> "crySub"
  , aMul            ~> "cryMul"
  , (|^|)           ~> "cryExp"
  , (|/|)           ~> "cryDiv"
  , (|%|)           ~> "cryMod"
  , aMin            ~> "cryMin"
  , aMax            ~> "cryMax"
  , aWidth          ~> "cryWidth"
  , aCeilDiv        ~> "cryCeilDiv"
  , aCeilMod        ~> "cryCeilMod"
  , aLenFromThenTo  ~> "cryLenFromThenTo"

  , anError KNum    ~> "cryErr"
  , anError KProp   ~> "cryErrProp"

  , aTVar           ~> "(unused)"
  ]

--------------------------------------------------------------------------------

(~>) :: Mk a => (Type -> Match a) -> String -> TVars -> Type -> Match SExpr
(m ~> f) tvs t = m t >>= \a -> return (mk tvs f a)

class Mk t where
  mk :: TVars -> String -> t -> SExpr

instance Mk () where
  mk _ f _ = SMT.const f

instance Mk Integer where
  mk _ f x = SMT.fun f [ SMT.int x ]

instance Mk TVar where
  mk tvs _ x = tvs Map.! x

instance Mk Type where
  mk tvs f x = SMT.fun f [toSMT tvs x]

instance Mk (Type,Type) where
  mk tvs f (x,y) = SMT.fun f [ toSMT tvs x, toSMT tvs y]

instance Mk (Type,Type,Type) where
  mk tvs f (x,y,z) = SMT.fun f [ toSMT tvs x, toSMT tvs y, toSMT tvs z ]

--------------------------------------------------------------------------------