```-- Copyright (C) 2016 Peter Selinger.
--
-- This file is free software and may be distributed under the terms

-- | This module extends the SAT solver to arbitrary boolean formulas
-- (not necessarily in conjunctive normal form). This is done by the
-- standard trick of translating each boolean formula to a CNF,
-- preserving satisfiability and the number of solutions. The
-- translation is carefully done in such a way that the size of the
-- CNF is linear in the size of the formula.

module SAT.MiniSat.Formula where

import SAT.MiniSat.Variable
import SAT.MiniSat.Literals

import qualified Data.Map as Map
import Data.Map (Map)
import Data.Maybe
import Prelude hiding (all)

-- ----------------------------------------------------------------------
-- * Boolean formulas

-- | The type of boolean formulas. It is parametric over a set of
-- variables /v/. We provide the usual boolean operators, including
-- implications and exclusive or. For convenience, we also provide the
-- list quantifiers 'All', 'Some', 'None', 'ExactlyOne', and
-- 'AtMostOne', as well as a general 'Let' operator that can be
-- used to reduce the size of repetitive formulas.

data Formula v =
Var v                       -- ^ A variable.
| Yes                       -- ^ The formula /true/.
| No                        -- ^ The formula /false/.
| Not (Formula v)           -- ^ Negation.
| Formula v :&&: Formula v  -- ^ Conjunction.
| Formula v :||: Formula v  -- ^ Disjunction.
| Formula v :++: Formula v  -- ^ Exclusive or.
| Formula v :->: Formula v  -- ^ Implication.
| Formula v :<->: Formula v -- ^ If and only if.
| All [Formula v]           -- ^ All are true.
| Some [Formula v]          -- ^ At least one is true.
| None [Formula v]          -- ^ None are true.
| ExactlyOne [Formula v]    -- ^ Exactly one is true.
| AtMostOne [Formula v]     -- ^ At most one is true.
| Let (Formula v) (Formula v -> Formula v)
-- ^ 'Let' /a/ /f/ is the formula \"let /x/=/a/ in /fx/\",
-- which permits the subexpression /a/ to be re-used. It
-- is logically equivalent to /fa/, but typically smaller.
| Bound Integer             -- ^ For internal use only.

infixr 6 :&&:
infixr 5 :||:
infixr 4 :++:
infixr 2 :->:
infix  1 :<->:

instance (Eq v) => Eq (Formula v) where
a == b = exformula 0 a == exformula 0 b

instance (Ord v) => Ord (Formula v) where
compare a b = compare (exformula 0 a) (exformula 0 b)

instance (Show v) => Show (Formula v) where
showsPrec d a = showsPrec d (exformula 0 a)

-- ----------------------------------------------------------------------
-- * Explicit formulas

-- | Same as 'Formula', except it uses explicit bound variables. This is
-- only used to define 'Eq', 'Ord', and 'Show' instances for
-- 'Formula'.
data ExFormula v =
ExVar v                             -- ^ A variable.
| ExTrue                            -- ^ The formula /true/.
| ExFalse                           -- ^ The formula /false/.
| ExNot (ExFormula v)               -- ^ Negation.
| ExAnd (ExFormula v) (ExFormula v) -- ^ Conjunction.
| ExOr (ExFormula v) (ExFormula v)  -- ^ Disjunction.
| ExXOr (ExFormula v) (ExFormula v) -- ^ Exclusive or.
| ExImp (ExFormula v) (ExFormula v) -- ^ Implication.
| ExIff (ExFormula v) (ExFormula v) -- ^ If and only if.
| ExAll [ExFormula v]               -- ^ All are true.
| ExSome [ExFormula v]              -- ^ At least one is true.
| ExNone [ExFormula v]              -- ^ None are true.
| ExExactlyOne [ExFormula v]        -- ^ Exactly one is true.
| ExAtMostOne [ExFormula v]         -- ^ At most one is true.
| ExLet Integer (ExFormula v) (ExFormula v) -- ^ Let.
| ExBound Integer                   -- ^ Bound variable.
deriving (Eq, Ord)

-- | Convert a 'Formula' to an 'ExFormula'. The first argument is the
-- current nesting depth of let binders.
exformula :: Integer -> Formula v -> ExFormula v
exformula i (Var v) = ExVar v
exformula i Yes = ExTrue
exformula i No = ExFalse
exformula i (Not a) = ExNot (exformula i a)
exformula i (a :&&: b) = ExAnd (exformula i a) (exformula i b)
exformula i (a :||: b) = ExOr (exformula i a) (exformula i b)
exformula i (a :++: b) = ExXOr (exformula i a) (exformula i b)
exformula i (a :->: b) = ExImp (exformula i a) (exformula i b)
exformula i (a :<->: b) = ExIff (exformula i a) (exformula i b)
exformula i (All vs) = ExAll [exformula i v | v <- vs]
exformula i (Some vs) = ExSome [exformula i v | v <- vs]
exformula i (None vs) = ExNone [exformula i v | v <- vs]
exformula i (ExactlyOne vs) = ExExactlyOne [exformula i v | v <- vs]
exformula i (AtMostOne vs) = ExAtMostOne [exformula i v | v <- vs]
exformula i (Let a f) = ExLet i (exformula i a) (exformula (i+1) (f (Bound i)))
exformula i (Bound x)
| 0 <= x && x < i = ExBound x
| otherwise = error "MiniSat.Formula: 'Bound' is for internal use only"

-- | Names for bound variables.
variables :: [String]
variables = vs ++ [ v ++ show n | n <- [1..], v <- vs ]
where
vs = ["x", "y", "z", "u", "v", "w", "r", "s", "t"]

-- | Convert an integer to a bound variable name.
showBound :: Integer -> ShowS
showBound i = showString (variables !! (fromInteger i))

-- | The 'Show' instance of 'ExFormula' actually shows the corresponding
-- 'Formula'. We omit parentheses between associative identical
-- operators, but show some additional parentheses between different
-- operators for clarity.
instance (Show v) => Show (ExFormula v) where
showsPrec d (ExVar v) = showParen (d > 10) \$
showString "Var " . showsPrec d v
showsPrec d ExTrue = showString "Yes"
showsPrec d ExFalse = showString "No"
showsPrec d (ExNot a) = showParen (d > 10) \$
showString "Not " . showsPrec 11 a
showsPrec d (ExAnd a b) = showParen (d > 6 || d == 4 || d == 5) \$
showsPrec 6 a . showString " :&&: " . showsPrec 6 b
showsPrec d (ExOr a b) = showParen (d > 5 || d == 4) \$
showsPrec 5 a . showString " :||: " . showsPrec 5 b
showsPrec d (ExXOr a b) = showParen (d > 4) \$
showsPrec 4 a . showString " :++: " . showsPrec 4 b
showsPrec d (ExImp a b) = showParen (d > 2) \$
showsPrec 3 a . showString " :->: " . showsPrec 2 b
showsPrec d (ExIff a b) = showParen (d > 1) \$
showsPrec 3 a . showString " :<->: " . showsPrec 3 b
showsPrec d (ExAll as) = showParen (d > 10) \$
showString "All " . showsPrec 11 as
showsPrec d (ExSome as) = showParen (d > 10) \$
showString "Some " . showsPrec 11 as
showsPrec d (ExNone as) = showParen (d > 10) \$
showString "None " . showsPrec 11 as
showsPrec d (ExExactlyOne as) = showParen (d > 10) \$
showString "ExactlyOne " . showsPrec 11 as
showsPrec d (ExAtMostOne as) = showParen (d > 10) \$
showString "AtMostOne " . showsPrec 11 as
showsPrec d (ExLet i a b) = showParen (d > 10) \$
showString "Let " . showsPrec 11 a . showString " (\\"
. showBound i . showString " -> " . showsPrec 0 b . showString ")"
showsPrec d (ExBound i) = showBound i

-- ----------------------------------------------------------------------
-- * Desugaring

-- | A version of 'Let' that binds a list of variables.
let_list :: [Formula v] -> ([Formula v] -> Formula v) -> Formula v
let_list [] f = f []
let_list (a:as) f = Let a (\x -> let_list as (\xs -> f (x:xs)))

-- | Implementation of 'All'.
all :: [Formula v] -> Formula v
all vs = foldl (:&&:) Yes vs

-- | Implementation of 'Some'.
some :: [Formula v] -> Formula v
some vs = foldl (:||:) No vs

-- | Implementation of 'None'.
none :: [Formula v] -> Formula v
none vs = Not (some vs)

-- | Implementation of 'AtMostOne'. Note that this translation is linear.
-- Without 'Let', the size of the translated formula would be /O/(/n/^2).
atMostOne :: [Formula v] -> Formula v
atMostOne [] = Yes
atMostOne [a] = Yes
atMostOne (a1:a2:as) =
Let a1 (\x1 ->
Let a2 (\x2 ->
Not (x1 :&&: x2) :&&: atMostOne ((x1 :||: x2) : as)))

-- | Implementation of 'ExactlyOne'.
exactlyOne :: [Formula v] -> Formula v
exactlyOne as = let_list as (\xs -> some xs :&&: atMostOne xs)

-- ----------------------------------------------------------------------
-- * Reduced boolean formulas

-- | A reduced boolean formula doesn't contain the constants \"true\" or
-- \"false\", nor the connectives \"and\", \"implies\", or \"iff\",
-- nor list quantifiers. We use reduced formulas as an intermediate
-- representation between boolean formulas and De Morgan normal forms.
-- The primary purpose of this is to get rid of the constants \"true\"
-- and \"false\" in absorbing positions (i.e., \"/a/ or true\", \"/b/
-- and false\").
data RFormula v =
RFVar v                           -- ^ A variable.
| RFBound Integer                 -- ^ A bound variable.
| RFNot (RFormula v)              -- ^ Negation.
| RFOr (RFormula v) (RFormula v)  -- ^ Disjunction.
| RFXOr (RFormula v) (RFormula v) -- ^ Exclusive or.
| RFLet Integer (RFormula v) (RFormula v) -- ^ Let /x/ = /a/ in /b/.
deriving (Show)

-- | Translate a boolean formula to reduced form, or one of the
-- constants 'True' or 'False'. The integer argument is the current
-- nesting depth of let binders, and is used to number local bound
-- variables.
rformula :: Integer -> Formula v -> Either Bool (RFormula v)
rformula i (Var v) = Right (RFVar v)
rformula i Yes = Left True
rformula i No = Left False
rformula i (Not a) =
case rformula i a of
Left x -> Left (not x)
Right a'' -> Right (RFNot a'')
rformula i (a :||: b) =
case (rformula i a, rformula i b) of
(Left True, b'') -> Left True
(Left False, b'') -> b''
(a'', Left True) -> Left True
(a'', Left False) -> a''
(Right a'', Right b'') -> Right (RFOr a'' b'')
rformula i (a :++: b) =
case (rformula i a, rformula i b) of
(Left x, Left y) -> Left (x /= y)
(Left True, Right b'') -> Right (RFNot b'')
(Left False, Right b'') -> Right b''
(Right a'', Left True) -> Right (RFNot a'')
(Right a'', Left False) -> Right a''
(Right a'', Right b'') -> Right (RFXOr a'' b'')
rformula i (a :&&: b) = rformula i (Not (Not a :||: Not b))
rformula i (a :->: b) = rformula i (Not a :||: b)
rformula i (a :<->: b) = rformula i (Not a :++: b)
rformula i (All vs) = rformula i (all vs)
rformula i (Some vs) = rformula i (some vs)
rformula i (None vs) = rformula i (none vs)
rformula i (AtMostOne vs) = rformula i (atMostOne vs)
rformula i (ExactlyOne vs) = rformula i (exactlyOne vs)
rformula i (Let a f) =
case rformula i a of
Left True -> rformula i (f Yes)
Left False -> rformula i (f No)
Right a'' ->
case rformula (i+1) (f (Bound i)) of
Left x -> Left x
Right b'' -> Right (RFLet i a'' b'')
rformula i (Bound x) = Right (RFBound x)

-- ----------------------------------------------------------------------
-- * De Morgan formulas

-- | Formulas in De Morgan standard form. We omit the constant formulas
-- 'True' and 'False', as they should only occur at the top level.
data DMFormula v =
DMPos v                                     -- ^ A positive literal.
| DMNeg v                                   -- ^ A negative literal.
| DMPosBound Integer                        -- ^ A positive bound variable.
| DMNegBound Integer                        -- ^ A negative bound variable.
| DMAnd (DMFormula v) (DMFormula v)         -- ^ Conjunction.
| DMOr (DMFormula v) (DMFormula v)          -- ^ Disjunction.
| DMXOr (DMFormula v) (DMFormula v)         -- ^ Exclusive or.
| DMLet Integer (DMFormula v) (DMFormula v) -- ^ let /x/ = /a/ in /b/.

-- | Translate a reduced formula to De Morgan standard form.
demorgan :: RFormula v -> DMFormula v
demorgan (RFVar v) = DMPos v
demorgan (RFBound x) = DMPosBound x
demorgan (RFNot a) = demorgan_neg a
demorgan (RFOr a b) = DMOr (demorgan a) (demorgan b)
demorgan (RFXOr a b) = DMXOr (demorgan a) (demorgan b)
demorgan (RFLet x a b) = DMLet x (demorgan a) (demorgan b)

-- | Translate the negation of a reduced formula to De Morgan standard
-- form.
demorgan_neg :: RFormula v -> DMFormula v
demorgan_neg (RFVar v) = DMNeg v
demorgan_neg (RFBound x) = DMNegBound x
demorgan_neg (RFNot a) = demorgan a
demorgan_neg (RFOr a b) = DMAnd (demorgan_neg a) (demorgan_neg b)
demorgan_neg (RFXOr a b) = DMXOr (demorgan_neg a) (demorgan b)
demorgan_neg (RFLet x a b) = DMLet x (demorgan a) (demorgan_neg b)

-- ----------------------------------------------------------------------
-- * A monad for translation to CNF

-- | Literals over a set /v/, extended with countably many additional elements.
type ELit v = Lit (Either Integer v)

-- | A monad for translation to CNF. This monad keeps track of two kinds
-- of state: an integer counter to provide a supply of fresh
-- variables, and a list of definitional clauses.
data Trans v a = Trans (Integer -> (a, Integer, [[ELit v]]))

return a = Trans (\n -> (a, n, []))
(Trans f) >>= g = Trans (\n ->
let (a1, n1, l1) = f n in
let Trans h = g a1 in
let (a2, n2, l2) = h n1 in
(a2, n2, l1 ++ l2))

instance Applicative (Trans v) where
pure = return
(<*>) = ap

instance Functor (Trans v) where
fmap = liftM

-- | Run the 'Trans' monad.
runTrans :: Trans v a -> (a, [[ELit v]])
runTrans (Trans f) = (a, l)
where
(a, _, l) = f 0

-- | Return a fresh literal.
fresh_lit :: Trans v (ELit v)
fresh_lit = Trans (\n -> (Pos (Left n), n+1, []))

-- | Add some definitional clauses.
add_cnf :: [[ELit v]] -> Trans v ()
add_cnf cs = Trans (\n -> ((), n, cs))

-- ----------------------------------------------------------------------
-- * Translation from De Morgan to CNF

-- \$ The simplest way to translate a formula, say
--
-- > (A ∧ B) ∨ (C ∧ (D ∨ E ∨ F)),
--
-- is to introduce a new variable for each intermediate result, i.e.,
-- to give clauses
--
-- > X, X ↔ Y ∨ Z, Y ↔ A ∧ B, Z ↔ C ∧ W, W ↔ D ∨ V, V ↔ E ∨ F.
--
-- (Each equivalence \"↔\" can be translated into CNF
-- straightforwardly).  Our translation is basically an optimized
-- version of this. The optimizations are: introduce additional
-- variables only when necessary; if the input formula is a CNF, leave
-- it unchanged. We give specialized translations of some formulas
-- depending on the context; for example, an exclusive or can be
-- translated in several different ways. For the above example, our
-- translation yields:
--
-- > Y ∨ Z, X ↔ A ∧ B, Z ↔ C ∧ W, W ↔ D ∨ E ∨ F.

-- | An environment is a mapping from local bound variables to
-- literals.
type Env v = Map Integer (ELit v)

-- | Translate a De Morgan formula to a SAT-equivalent CNF.  Since this
-- requires the introduction of fresh variables, the CNF is over a
-- larger variable set, and we work in the 'Trans' monad. Moreover,
-- the translation depends on an environment.
trans_cnf :: Env v -> DMFormula v -> Trans v [[ELit v]]
trans_cnf env (DMPos v) = return [[Pos (Right v)]]
trans_cnf env (DMNeg v) = return [[Neg (Right v)]]
trans_cnf env (DMPosBound n) = return [[x]]
where
x = env Map.! n
trans_cnf env (DMNegBound n) = return [[neg x]]
where
x = env Map.! n
trans_cnf env (DMAnd a b) = do
a' <- trans_cnf env a
b' <- trans_cnf env b
return (a' ++ b')
trans_cnf env (DMOr a b) = do
a' <- trans_or env a
b' <- trans_or env b
return [a' ++ b']
trans_cnf env (DMXOr a b) = do
x <- trans_lit env a
y <- trans_lit env b
return [[x, y], [neg x, neg y]]
trans_cnf env (DMLet n a b) = do
do_let env n a b trans_cnf

-- | Translate a De Morgan formula to a disjunction of literals.
trans_or :: Env v -> DMFormula v -> Trans v [ELit v]
trans_or env (DMXOr a b) = do
x <- trans_lit env (DMXOr a b)
return [x]
trans_or env (DMLet n a b) = do
do_let env n a b trans_or
trans_or env a = do
c <- trans_cnf env a
or_of_cnf c

-- | Translate a De Morgan formula to a single literal.
trans_lit :: Env v -> DMFormula v -> Trans v (ELit v)
trans_lit env (DMXOr a b) = do
x <- trans_lit env a
y <- trans_lit env b
z <- lit_of_xor x y
return z
trans_lit env (DMLet n a b) = do
do_let env n a b trans_lit
trans_lit env a = do
c <- trans_cnf env a
lit_of_cnf c

-- | Polymorphically translate a 'DMLet' expression. The fifth
-- argument is usually one of 'trans_cnf', 'trans_or', or 'trans_lit',
-- and is used to translate the body of the \"let\".
do_let :: Env v -> Integer -> DMFormula v -> DMFormula v -> (Env v -> DMFormula v -> Trans v a) -> Trans v a
do_let env n a b cont = do
x <- trans_lit env a
let env' = Map.insert n x env
cont env' b

-- | Convert a CNF to a disjunction of literals.
or_of_cnf :: [[ELit v]] -> Trans v [ELit v]
or_of_cnf [d] = return d
or_of_cnf ds = do
x <- lit_of_cnf ds
return [x]

-- | Convert a CNF to a single literal.
lit_of_cnf :: [[ELit v]] -> Trans v (ELit v)
lit_of_cnf ds = do
xs <- sequence (map lit_of_or ds)
y <- lit_of_and xs
return y

-- | Convert a conjunction of literals to a single literal.
lit_of_and :: [ELit v] -> Trans v (ELit v)
lit_of_and [l] = return l
lit_of_and cs = do
x <- fresh_lit
-- Define x <-> c1 ∧ ... ∧ cn
add_cnf [[neg x, c] | c <- cs ]
add_cnf [[x] ++ [neg c | c <- cs]]
return x

-- | Convert a disjunction of literals to a single literal.
lit_of_or :: [ELit v] -> Trans v (ELit v)
lit_of_or [l] = return l
lit_of_or ds = do
x <- fresh_lit
-- Define x <-> d1 ∨ ... ∨ dn
add_cnf [ [x, neg d] | d <- ds ]
add_cnf [[neg x] ++ [d | d <- ds]]
return x

-- | Convert an exclusive or of two literals to a single literal.
lit_of_xor :: ELit v -> ELit v -> Trans v (ELit v)
lit_of_xor x y = do
z <- fresh_lit
-- Define z <-> x ⊕ y
add_cnf [[neg z, neg x, neg y]]
return z

-- ----------------------------------------------------------------------
-- * Top-level functions

-- | Translate a De Morgan formula to a SAT-equivalent CNF.
cnf_of_dm :: DMFormula v -> [[ELit v]]
cnf_of_dm f = l ++ a
where
env = Map.empty
(a, l) = runTrans (trans_cnf env f)

-- | Translate a boolean formula to a SAT-equivalent CNF.
cnf_of_formula :: Formula v -> [[ELit v]]
cnf_of_formula f =
case rformula 0 f of
Left True -> []
Left False -> [[]]
Right rf -> cnf_of_dm (demorgan rf)

-- ----------------------------------------------------------------------
-- * SAT solver API for boolean formulas

-- | Check whether a boolean formula is satisfiable.
satisfiable :: (Ord v) => Formula v -> Bool
satisfiable a = isJust (solve a)

-- | Return a satisfying assignment for the boolean formula, if any.
solve :: (Ord v) => Formula v -> Maybe (Map v Bool)
solve a = listToMaybe (solve_all a)

-- | Lazily enumerate all satisfying assignments for the boolean formula.
solve_all :: (Ord v) => Formula v -> [Map v Bool]
solve_all a = res
where
a' = cnf_of_formula a
res' = cnf_solve_all a'
res = map restrict res'
restrict ans = Map.fromList [ (v, b) | (Right v, b) <- Map.toList ans ]
```