Main.hs

{-# language DataKinds                  #-}
{-# language FlexibleContexts           #-}
{-# language FlexibleInstances          #-}
{-# language GADTs                      #-}
{-# language LambdaCase                 #-}
{-# language PatternSynonyms            #-}
{-# language PolyKinds                  #-}
{-# language TypeApplications           #-}
{-# language TypeOperators              #-}

module Main where

import           Control.Category
import           Control.Monad          ((<=<))
import           Control.Monad.IO.Class (liftIO)
import           Data.SBV               hiding (name, (<+>))
import           Data.SBV.Control       hiding (Name)
import           Prelude                hiding (id, (.), (>>))

data IntegerOp = Add | Sub | Mul | Div

data LogicalOp a b where
  And :: LogicalOp (Bool ': Bool ': a) (Bool ': a)
  Or  :: LogicalOp (Bool ': Bool ': a) (Bool ': a)
  Not :: LogicalOp (Bool ': a)         (Bool ': a)

data StrOp a where
  StrConcat  :: StrOp String
  StrCompare :: StrOp Bool

data HList :: (k -> *) -> [k] -> * where
  RNil :: HList f '[]
  (:<) :: f a -> HList f as -> HList f (a ': as)

infixr 5 :<

instance Mergeable (HList SBV a) where
  symbolicMerge _     _    RNil      RNil      = RNil
  symbolicMerge force test (x :< xs) (y :< ys)
    | x == y    = x :< symbolicMerge force test xs ys
    | otherwise = error "No least-upper-bound for HList SBV"

data Term dom codom where
  IntegerOp
    :: IntegerOp
    -> Term (Integer ': Integer ': a) (Integer ': a)

  LogicalOp
    :: LogicalOp a b
    -> Term a b

  StrOp
    :: StrOp a
    -> Term (String ': String ': b) (a ': b)

  IfThenElse
    :: Term a b
    -> Term a b
    -> Term (Bool ': a) b

  PushLit
    :: (SymWord a, Show a)
    => a
    -> Term dom (a ': dom)

  -- Used for sbv quantified values
  PushSbvVal :: SBV x -> Term a (x ': a)

  Noop  :: Term a                 a
  Pop   :: Term (x ': a)          a
  Dup   :: Term (x ': a)          (x ': x ': a)
  Swap0 :: Term (x ': y ': a)     (y ': x ': a)
  Swap1 :: Term (x ': y : z ': a) (z ': y ': x ': a)

  Sequence :: Term a b -> Term b c -> Term a c

instance Category Term where
  id  = Noop
  (.) = flip Sequence

pattern And' :: Term (Bool ': Bool ': a) (Bool ': a)
pattern And' = LogicalOp And

pattern Or' :: Term (Bool ': Bool ': a) (Bool ': a)
pattern Or' = LogicalOp Or

pattern Not' :: Term (Bool ': a) (Bool ': a)
pattern Not' = LogicalOp Not

pattern StrConcat' :: Term (String ': String ': a) (String ': a)
pattern StrConcat'  = StrOp StrConcat

pattern StrCompare' :: Term (String ': String ': a) (Bool ': a)
pattern StrCompare' = StrOp StrCompare

(>>) :: Category k => k a b -> k b c -> k a c
(>>) = flip (.)

-- evaluation

newtype Eval dom codom = Eval { runEval :: dom -> Maybe codom }

instance Functor (Eval dom) where
  fmap f (Eval g) = Eval (fmap f . g)

instance Applicative (Eval dom) where
  pure a            = Eval (const (Just a))
  Eval f <*> Eval a = Eval (\dom -> f dom <*> a dom)

instance Category Eval where
  id              = Eval Just
  Eval f . Eval g = Eval (f <=< g)

instance Mergeable codom => Mergeable (Eval dom codom) where
  symbolicMerge force test (Eval left) (Eval right) = Eval $
    symbolicMerge force test left right

lift1
  :: (SBV x -> SBV y)
  -> Eval (HList SBV (x ': a)) (HList SBV (y ': a))
lift1 f = Eval $ \(x :< stack) -> Just (f x :< stack)

lift2
  :: (SBV x -> SBV y -> SBV z)
  -> Eval (HList SBV (x ': y ': a)) (HList SBV (z ': a))
lift2 f = Eval $ \(x :< y :< stack) -> Just (f x y :< stack)

arr :: (a -> b) -> Eval a b
arr x = Eval $ Just . x

eval
  :: Term dom codom
  -> Eval (HList SBV dom) (HList SBV codom)
eval = \case
  PushLit i    -> arr (literal i :<)
  PushSbvVal a -> arr (a :<)

  IntegerOp op -> lift2 $ case op of
    Add -> (+)
    Sub -> (-)
    Mul -> (*)
    Div -> sDiv

  LogicalOp op -> case op of
    Not -> lift1 bnot
    And -> lift2 (&&&)
    Or  -> lift2 (|||)

  StrOp StrConcat  -> lift2 (.++)
  StrOp StrCompare -> lift2 (.==)

  IfThenElse trueCase falseCase -> Eval $ \(x :< stack) -> ite x
    (runEval (eval trueCase)  stack)
    (runEval (eval falseCase) stack)

  Noop  -> id
  Pop   -> arr $ \(          _ :< stack) ->                 stack
  Dup   -> arr $ \(          x :< stack) ->      (x :< x :< stack)
  Swap0 -> arr $ \(     x :< y :< stack) ->      (y :< x :< stack)
  Swap1 -> arr $ \(x :< y :< z :< stack) -> (z :< y :< x :< stack)

  Sequence g f -> eval f . eval g

runTerm :: Term '[] '[Bool] -> Symbolic ()
runTerm term =
  case runEval (eval term) RNil of
    Nothing -> liftIO $ putStrLn "analysis failed"
    Just (sbvTerm :< RNil) -> do
      constrain $ sbvTerm .== true
      query $ do
        cs <- checkSat
        case cs of
          Unk -> liftIO $ putStrLn "solver returned unknown"
          Unsat -> liftIO $ putStrLn "not satisfiable"
          Sat -> do
            liftIO $ putStrLn "input:"
            liftIO . print =<< getModel

main :: IO ()
main = do
  runSMT $ do
    [x, y, z] <- traverse (exists @Bool) ["x", "y", "z"]

    let term :: Term '[Bool, Bool, Bool] '[Bool]
        term = IfThenElse
          (Not' >> And')
          (Or'  >> Not')

    runTerm (PushSbvVal x >> PushSbvVal y >> PushSbvVal z >> term)

--   runSMT $ do
--     [x, y, z, w] <- traverse (exists @String) ["x", "y", "z", "w"]
--     runTerm $ And'
--       [ (StrCompare' (StrConcat' x y) (StrConcat' z w))
--       , (StrCompare' (StrConcat' x y) (PushLit "monic"))
--       ]