{-# OPTIONS_HADDOCK show-extensions #-}
{-# LANGUAGE TypeFamilies #-}
-----------------------------------------------------------------------------
-- |
-- Module      :  ToySolver.SAT.Encoder.Integer
-- Copyright   :  (c) Masahiro Sakai 2012
-- License     :  BSD-style
--
-- Maintainer  :  masahiro.sakai@gmail.com
-- Stability   :  provisional
-- Portability :  non-portable
--
-----------------------------------------------------------------------------
module ToySolver.SAT.Encoder.Integer
  ( Expr (..)
  , newVar
  , linearize
  , addConstraint
  , addConstraintSoft
  , eval
  ) where

import Control.Monad
import Control.Monad.Primitive
import Data.Array.IArray
import Data.VectorSpace
import Text.Printf

import ToySolver.Data.OrdRel
import qualified ToySolver.SAT.Types as SAT
import qualified ToySolver.SAT.Encoder.PBNLC as PBNLC

newtype Expr = Expr SAT.PBSum
  deriving (Expr -> Expr -> Bool
(Expr -> Expr -> Bool) -> (Expr -> Expr -> Bool) -> Eq Expr
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
/= :: Expr -> Expr -> Bool
$c/= :: Expr -> Expr -> Bool
== :: Expr -> Expr -> Bool
$c== :: Expr -> Expr -> Bool
Eq, Int -> Expr -> ShowS
[Expr] -> ShowS
Expr -> String
(Int -> Expr -> ShowS)
-> (Expr -> String) -> ([Expr] -> ShowS) -> Show Expr
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
showList :: [Expr] -> ShowS
$cshowList :: [Expr] -> ShowS
show :: Expr -> String
$cshow :: Expr -> String
showsPrec :: Int -> Expr -> ShowS
$cshowsPrec :: Int -> Expr -> ShowS
Show, ReadPrec [Expr]
ReadPrec Expr
Int -> ReadS Expr
ReadS [Expr]
(Int -> ReadS Expr)
-> ReadS [Expr] -> ReadPrec Expr -> ReadPrec [Expr] -> Read Expr
forall a.
(Int -> ReadS a)
-> ReadS [a] -> ReadPrec a -> ReadPrec [a] -> Read a
readListPrec :: ReadPrec [Expr]
$creadListPrec :: ReadPrec [Expr]
readPrec :: ReadPrec Expr
$creadPrec :: ReadPrec Expr
readList :: ReadS [Expr]
$creadList :: ReadS [Expr]
readsPrec :: Int -> ReadS Expr
$creadsPrec :: Int -> ReadS Expr
Read)

newVar :: SAT.AddPBNL m enc => enc -> Integer -> Integer -> m Expr
newVar :: enc -> Integer -> Integer -> m Expr
newVar enc
enc Integer
lo Integer
hi
  | Integer
lo Integer -> Integer -> Bool
forall a. Ord a => a -> a -> Bool
> Integer
hi = do
      enc -> Clause -> m ()
forall (m :: * -> *) a. AddClause m a => a -> Clause -> m ()
SAT.addClause enc
enc [] -- assert inconsistency
      Expr -> m Expr
forall (m :: * -> *) a. Monad m => a -> m a
return Expr
0
  | Integer
lo Integer -> Integer -> Bool
forall a. Eq a => a -> a -> Bool
== Integer
hi = Expr -> m Expr
forall (m :: * -> *) a. Monad m => a -> m a
return (Expr -> m Expr) -> Expr -> m Expr
forall a b. (a -> b) -> a -> b
$ Integer -> Expr
forall a. Num a => Integer -> a
fromInteger Integer
lo
  | Bool
otherwise = do
      let hi' :: Integer
hi' = Integer
hi Integer -> Integer -> Integer
forall a. Num a => a -> a -> a
- Integer
lo
          bitWidth :: Int
bitWidth = Clause -> Int
forall a. [a] -> a
head (Clause -> Int) -> Clause -> Int
forall a b. (a -> b) -> a -> b
$ [Int
w | Int
w <- [Int
1..], let mx :: Integer
mx = Integer
2 Integer -> Int -> Integer
forall a b. (Num a, Integral b) => a -> b -> a
^ Int
w Integer -> Integer -> Integer
forall a. Num a => a -> a -> a
- Integer
1, Integer
hi' Integer -> Integer -> Bool
forall a. Ord a => a -> a -> Bool
<= Integer
mx]
      Clause
vs <- enc -> Int -> m Clause
forall (m :: * -> *) a. NewVar m a => a -> Int -> m Clause
SAT.newVars enc
enc Int
bitWidth
      let xs :: [(Integer, Int)]
xs = [Integer] -> Clause -> [(Integer, Int)]
forall a b. [a] -> [b] -> [(a, b)]
zip ((Integer -> Integer) -> Integer -> [Integer]
forall a. (a -> a) -> a -> [a]
iterate (Integer
2Integer -> Integer -> Integer
forall a. Num a => a -> a -> a
*) Integer
1) Clause
vs
      enc -> [(Integer, Int)] -> Integer -> m ()
forall (m :: * -> *) a.
AddPBLin m a =>
a -> [(Integer, Int)] -> Integer -> m ()
SAT.addPBAtMost enc
enc [(Integer, Int)]
xs Integer
hi'
      Expr -> m Expr
forall (m :: * -> *) a. Monad m => a -> m a
return (Expr -> m Expr) -> Expr -> m Expr
forall a b. (a -> b) -> a -> b
$ PBSum -> Expr
Expr ((Integer
lo,[]) (Integer, Clause) -> PBSum -> PBSum
forall a. a -> [a] -> [a]
: [(Integer
c,[Int
x]) | (Integer
c,Int
x) <- [(Integer, Int)]
xs])

instance AdditiveGroup Expr where
  Expr PBSum
xs1 ^+^ :: Expr -> Expr -> Expr
^+^ Expr PBSum
xs2 = PBSum -> Expr
Expr (PBSum
xs1PBSum -> PBSum -> PBSum
forall a. [a] -> [a] -> [a]
++PBSum
xs2)
  zeroV :: Expr
zeroV = PBSum -> Expr
Expr []
  negateV :: Expr -> Expr
negateV = ((-Integer
1) Scalar Expr -> Expr -> Expr
forall v. VectorSpace v => Scalar v -> v -> v
*^)

instance VectorSpace Expr where
  type Scalar Expr = Integer
  Scalar Expr
n *^ :: Scalar Expr -> Expr -> Expr
*^ Expr PBSum
xs = PBSum -> Expr
Expr [(Integer
Scalar Expr
nInteger -> Integer -> Integer
forall a. Num a => a -> a -> a
*Integer
m,Clause
lits) | (Integer
m,Clause
lits) <- PBSum
xs]

instance Num Expr where
  Expr PBSum
xs1 + :: Expr -> Expr -> Expr
+ Expr PBSum
xs2 = PBSum -> Expr
Expr (PBSum
xs1PBSum -> PBSum -> PBSum
forall a. [a] -> [a] -> [a]
++PBSum
xs2)
  Expr PBSum
xs1 * :: Expr -> Expr -> Expr
* Expr PBSum
xs2 = PBSum -> Expr
Expr [(Integer
c1Integer -> Integer -> Integer
forall a. Num a => a -> a -> a
*Integer
c2, Clause
lits1Clause -> Clause -> Clause
forall a. [a] -> [a] -> [a]
++Clause
lits2) | (Integer
c1,Clause
lits1) <- PBSum
xs1, (Integer
c2,Clause
lits2) <- PBSum
xs2]
  negate :: Expr -> Expr
negate (Expr PBSum
xs) = PBSum -> Expr
Expr [(-Integer
c,Clause
lits) | (Integer
c,Clause
lits) <- PBSum
xs]
  abs :: Expr -> Expr
abs      = Expr -> Expr
forall a. a -> a
id
  signum :: Expr -> Expr
signum Expr
_ = Expr
1
  fromInteger :: Integer -> Expr
fromInteger Integer
c = PBSum -> Expr
Expr [(Integer
c,[])]

linearize :: PrimMonad m => PBNLC.Encoder m -> Expr -> m (SAT.PBLinSum, Integer)
linearize :: Encoder m -> Expr -> m ([(Integer, Int)], Integer)
linearize Encoder m
enc (Expr PBSum
xs) = do
  let ys :: PBSum
ys = [(Integer
c,Clause
lits) | (Integer
c,lits :: Clause
lits@(Int
_:Clause
_)) <- PBSum
xs]
      c :: Integer
c  = [Integer] -> Integer
forall (t :: * -> *) a. (Foldable t, Num a) => t a -> a
sum [Integer
c | (Integer
c,[]) <- PBSum
xs]
  [(Integer, Int)]
zs <- Encoder m -> PBSum -> m [(Integer, Int)]
forall (m :: * -> *).
PrimMonad m =>
Encoder m -> PBSum -> m [(Integer, Int)]
PBNLC.linearizePBSum Encoder m
enc PBSum
ys
  ([(Integer, Int)], Integer) -> m ([(Integer, Int)], Integer)
forall (m :: * -> *) a. Monad m => a -> m a
return ([(Integer, Int)]
zs, Integer
c)

addConstraint :: SAT.AddPBNL m enc => enc -> OrdRel Expr -> m ()
addConstraint :: enc -> OrdRel Expr -> m ()
addConstraint enc
enc (OrdRel Expr
lhs RelOp
op Expr
rhs) = do
  let Expr PBSum
e = Expr
lhs Expr -> Expr -> Expr
forall a. Num a => a -> a -> a
- Expr
rhs
  case RelOp
op of
    RelOp
Le  -> enc -> PBSum -> Integer -> m ()
forall (m :: * -> *) a.
AddPBNL m a =>
a -> PBSum -> Integer -> m ()
SAT.addPBNLAtMost  enc
enc PBSum
e Integer
0
    RelOp
Lt  -> enc -> PBSum -> Integer -> m ()
forall (m :: * -> *) a.
AddPBNL m a =>
a -> PBSum -> Integer -> m ()
SAT.addPBNLAtMost  enc
enc PBSum
e (-Integer
1)
    RelOp
Ge  -> enc -> PBSum -> Integer -> m ()
forall (m :: * -> *) a.
AddPBNL m a =>
a -> PBSum -> Integer -> m ()
SAT.addPBNLAtLeast enc
enc PBSum
e Integer
0
    RelOp
Gt  -> enc -> PBSum -> Integer -> m ()
forall (m :: * -> *) a.
AddPBNL m a =>
a -> PBSum -> Integer -> m ()
SAT.addPBNLAtLeast enc
enc PBSum
e Integer
1
    RelOp
Eql -> enc -> PBSum -> Integer -> m ()
forall (m :: * -> *) a.
AddPBNL m a =>
a -> PBSum -> Integer -> m ()
SAT.addPBNLExactly enc
enc PBSum
e Integer
0
    RelOp
NEq -> do
      Int
sel <- enc -> m Int
forall (m :: * -> *) a. NewVar m a => a -> m Int
SAT.newVar enc
enc
      enc -> Int -> PBSum -> Integer -> m ()
forall (m :: * -> *) a.
AddPBNL m a =>
a -> Int -> PBSum -> Integer -> m ()
SAT.addPBNLAtLeastSoft enc
enc Int
sel PBSum
e Integer
1
      enc -> Int -> PBSum -> Integer -> m ()
forall (m :: * -> *) a.
AddPBNL m a =>
a -> Int -> PBSum -> Integer -> m ()
SAT.addPBNLAtMostSoft  enc
enc (-Int
sel) PBSum
e (-Integer
1)

addConstraintSoft :: SAT.AddPBNL m enc => enc -> SAT.Lit -> OrdRel Expr -> m ()
addConstraintSoft :: enc -> Int -> OrdRel Expr -> m ()
addConstraintSoft enc
enc Int
sel (OrdRel Expr
lhs RelOp
op Expr
rhs) = do
  let Expr PBSum
e = Expr
lhs Expr -> Expr -> Expr
forall a. Num a => a -> a -> a
- Expr
rhs
  case RelOp
op of
    RelOp
Le  -> enc -> Int -> PBSum -> Integer -> m ()
forall (m :: * -> *) a.
AddPBNL m a =>
a -> Int -> PBSum -> Integer -> m ()
SAT.addPBNLAtMostSoft  enc
enc Int
sel PBSum
e Integer
0
    RelOp
Lt  -> enc -> Int -> PBSum -> Integer -> m ()
forall (m :: * -> *) a.
AddPBNL m a =>
a -> Int -> PBSum -> Integer -> m ()
SAT.addPBNLAtMostSoft  enc
enc Int
sel PBSum
e (-Integer
1)
    RelOp
Ge  -> enc -> Int -> PBSum -> Integer -> m ()
forall (m :: * -> *) a.
AddPBNL m a =>
a -> Int -> PBSum -> Integer -> m ()
SAT.addPBNLAtLeastSoft enc
enc Int
sel PBSum
e Integer
0
    RelOp
Gt  -> enc -> Int -> PBSum -> Integer -> m ()
forall (m :: * -> *) a.
AddPBNL m a =>
a -> Int -> PBSum -> Integer -> m ()
SAT.addPBNLAtLeastSoft enc
enc Int
sel PBSum
e Integer
1
    RelOp
Eql -> enc -> Int -> PBSum -> Integer -> m ()
forall (m :: * -> *) a.
AddPBNL m a =>
a -> Int -> PBSum -> Integer -> m ()
SAT.addPBNLExactlySoft enc
enc Int
sel PBSum
e Integer
0
    RelOp
NEq -> do
      Int
sel2 <- enc -> m Int
forall (m :: * -> *) a. NewVar m a => a -> m Int
SAT.newVar enc
enc
      Int
sel3 <- enc -> m Int
forall (m :: * -> *) a. NewVar m a => a -> m Int
SAT.newVar enc
enc
      Int
sel4 <- enc -> m Int
forall (m :: * -> *) a. NewVar m a => a -> m Int
SAT.newVar enc
enc
      enc -> Clause -> m ()
forall (m :: * -> *) a. AddClause m a => a -> Clause -> m ()
SAT.addClause enc
enc [-Int
sel, -Int
sel2, Int
sel3] -- sel ∧  sel2 → sel3
      enc -> Clause -> m ()
forall (m :: * -> *) a. AddClause m a => a -> Clause -> m ()
SAT.addClause enc
enc [-Int
sel,  Int
sel2, Int
sel4] -- sel ∧ ¬sel2 → sel4
      enc -> Int -> PBSum -> Integer -> m ()
forall (m :: * -> *) a.
AddPBNL m a =>
a -> Int -> PBSum -> Integer -> m ()
SAT.addPBNLAtLeastSoft enc
enc Int
sel3 PBSum
e Integer
1
      enc -> Int -> PBSum -> Integer -> m ()
forall (m :: * -> *) a.
AddPBNL m a =>
a -> Int -> PBSum -> Integer -> m ()
SAT.addPBNLAtMostSoft  enc
enc Int
sel4 PBSum
e (-Integer
1)

eval :: SAT.IModel m => m -> Expr -> Integer
eval :: m -> Expr -> Integer
eval m
m (Expr PBSum
ts) = [Integer] -> Integer
forall (t :: * -> *) a. (Foldable t, Num a) => t a -> a
sum [if [Bool] -> Bool
forall (t :: * -> *). Foldable t => t Bool -> Bool
and [m -> Int -> Bool
forall m. IModel m => m -> Int -> Bool
SAT.evalLit m
m Int
lit | Int
lit <- Clause
lits] then Integer
n else Integer
0 | (Integer
n,Clause
lits) <- PBSum
ts]