{-# 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
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, Lit -> Expr -> ShowS
[Expr] -> ShowS
Expr -> String
forall a.
(Lit -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
showList :: [Expr] -> ShowS
$cshowList :: [Expr] -> ShowS
show :: Expr -> String
$cshow :: Expr -> String
showsPrec :: Lit -> Expr -> ShowS
$cshowsPrec :: Lit -> Expr -> ShowS
Show, ReadPrec [Expr]
ReadPrec Expr
Lit -> ReadS Expr
ReadS [Expr]
forall a.
(Lit -> 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 :: Lit -> ReadS Expr
$creadsPrec :: Lit -> ReadS Expr
Read)

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

instance AdditiveGroup Expr where
  Expr PBSum
xs1 ^+^ :: Expr -> Expr -> Expr
^+^ Expr PBSum
xs2 = PBSum -> Expr
Expr (PBSum
xs1forall a. [a] -> [a] -> [a]
++PBSum
xs2)
  zeroV :: Expr
zeroV = PBSum -> Expr
Expr []
  negateV :: Expr -> Expr
negateV = ((-Integer
1) 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 [(Scalar Expr
nforall a. Num a => a -> a -> a
*Integer
m,[Lit]
lits) | (Integer
m,[Lit]
lits) <- PBSum
xs]

instance Num Expr where
  Expr PBSum
xs1 + :: Expr -> Expr -> Expr
+ Expr PBSum
xs2 = PBSum -> Expr
Expr (PBSum
xs1forall a. [a] -> [a] -> [a]
++PBSum
xs2)
  Expr PBSum
xs1 * :: Expr -> Expr -> Expr
* Expr PBSum
xs2 = PBSum -> Expr
Expr [(Integer
c1forall a. Num a => a -> a -> a
*Integer
c2, [Lit]
lits1forall a. [a] -> [a] -> [a]
++[Lit]
lits2) | (Integer
c1,[Lit]
lits1) <- PBSum
xs1, (Integer
c2,[Lit]
lits2) <- PBSum
xs2]
  negate :: Expr -> Expr
negate (Expr PBSum
xs) = PBSum -> Expr
Expr [(-Integer
c,[Lit]
lits) | (Integer
c,[Lit]
lits) <- PBSum
xs]
  abs :: Expr -> Expr
abs      = 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 :: forall (m :: * -> *).
PrimMonad m =>
Encoder m -> Expr -> m ([(Integer, Lit)], Integer)
linearize Encoder m
enc (Expr PBSum
xs) = do
  let ys :: PBSum
ys = [(Integer
c,[Lit]
lits) | (Integer
c,lits :: [Lit]
lits@(Lit
_:[Lit]
_)) <- PBSum
xs]
      c :: Integer
c  = forall (t :: * -> *) a. (Foldable t, Num a) => t a -> a
sum [Integer
c | (Integer
c,[]) <- PBSum
xs]
  [(Integer, Lit)]
zs <- forall (m :: * -> *).
PrimMonad m =>
Encoder m -> PBSum -> m [(Integer, Lit)]
PBNLC.linearizePBSum Encoder m
enc PBSum
ys
  forall (m :: * -> *) a. Monad m => a -> m a
return ([(Integer, Lit)]
zs, Integer
c)

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

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

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