{-# LANGUAGE BangPatterns #-}
{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# OPTIONS_GHC -Wall #-}
{-# OPTIONS_HADDOCK show-extensions #-}
-----------------------------------------------------------------------------
-- |
-- Module      :  ToySolver.SAT.Encoder.PB.Internal.Adder
-- Copyright   :  (c) Masahiro Sakai 2016
-- License     :  BSD-style
--
-- Maintainer  :  masahiro.sakai@gmail.com
-- Stability   :  provisional
-- Portability :  non-portable
--
-- References:
--
-- * [ES06] N. Eén and N. Sörensson. Translating Pseudo-Boolean
--   Constraints into SAT. JSAT 2:1–26, 2006.
--
-----------------------------------------------------------------------------
module ToySolver.SAT.Encoder.PB.Internal.Adder
  ( addPBLinAtLeastAdder
  , encodePBLinAtLeastWithPolarityAdder
  ) where

import Control.Monad
import Control.Monad.Primitive
import Data.Bits
import Data.Maybe
import Data.Primitive.MutVar
import Data.Sequence (Seq)
import qualified Data.Sequence as Seq
import ToySolver.Data.Boolean
import qualified ToySolver.Internal.Data.SeqQueue as SQ
import qualified ToySolver.SAT.Types as SAT
import qualified ToySolver.SAT.Encoder.Tseitin as Tseitin
import ToySolver.SAT.Encoder.Tseitin (Formula (..))

addPBLinAtLeastAdder :: forall m. PrimMonad m => Tseitin.Encoder m -> SAT.PBLinAtLeast -> m ()
addPBLinAtLeastAdder :: forall (m :: * -> *).
PrimMonad m =>
Encoder m -> PBLinAtLeast -> m ()
addPBLinAtLeastAdder Encoder m
enc PBLinAtLeast
constr = do
  Formula
formula <- Encoder m -> PBLinAtLeast -> m Formula
forall (m :: * -> *).
PrimMonad m =>
Encoder m -> PBLinAtLeast -> m Formula
encodePBLinAtLeastAdder' Encoder m
enc PBLinAtLeast
constr
  Encoder m -> Formula -> m ()
forall (m :: * -> *). PrimMonad m => Encoder m -> Formula -> m ()
Tseitin.addFormula Encoder m
enc Formula
formula

encodePBLinAtLeastWithPolarityAdder :: PrimMonad m => Tseitin.Encoder m -> Tseitin.Polarity -> SAT.PBLinAtLeast -> m SAT.Lit
encodePBLinAtLeastWithPolarityAdder :: forall (m :: * -> *).
PrimMonad m =>
Encoder m -> Polarity -> PBLinAtLeast -> m Lit
encodePBLinAtLeastWithPolarityAdder Encoder m
enc Polarity
polarity PBLinAtLeast
constr = do
  Formula
formula <- Encoder m -> PBLinAtLeast -> m Formula
forall (m :: * -> *).
PrimMonad m =>
Encoder m -> PBLinAtLeast -> m Formula
encodePBLinAtLeastAdder' Encoder m
enc PBLinAtLeast
constr
  Encoder m -> Polarity -> Formula -> m Lit
forall (m :: * -> *).
PrimMonad m =>
Encoder m -> Polarity -> Formula -> m Lit
Tseitin.encodeFormulaWithPolarity Encoder m
enc Polarity
polarity Formula
formula

encodePBLinAtLeastAdder' :: PrimMonad m => Tseitin.Encoder m -> SAT.PBLinAtLeast -> m Tseitin.Formula
encodePBLinAtLeastAdder' :: forall (m :: * -> *).
PrimMonad m =>
Encoder m -> PBLinAtLeast -> m Formula
encodePBLinAtLeastAdder' Encoder m
_ (PBLinSum
_,Integer
rhs) | Integer
rhs Integer -> Integer -> Bool
forall a. Ord a => a -> a -> Bool
<= Integer
0 = Formula -> m Formula
forall a. a -> m a
forall (m :: * -> *) a. Monad m => a -> m a
return Formula
forall a. MonotoneBoolean a => a
true
encodePBLinAtLeastAdder' Encoder m
enc (PBLinSum
lhs,Integer
rhs) = do
  [Lit]
lhs1 <- Encoder m -> PBLinSum -> m [Lit]
forall (m :: * -> *).
PrimMonad m =>
Encoder m -> PBLinSum -> m [Lit]
encodePBLinSumAdder Encoder m
enc PBLinSum
lhs
  let rhs1 :: [Bool]
rhs1 = Integer -> [Bool]
bits Integer
rhs
  if [Lit] -> Lit
forall a. [a] -> Lit
forall (t :: * -> *) a. Foldable t => t a -> Lit
length [Lit]
lhs1 Lit -> Lit -> Bool
forall a. Ord a => a -> a -> Bool
< [Bool] -> Lit
forall a. [a] -> Lit
forall (t :: * -> *) a. Foldable t => t a -> Lit
length [Bool]
rhs1 then do
    Formula -> m Formula
forall a. a -> m a
forall (m :: * -> *) a. Monad m => a -> m a
return Formula
forall a. MonotoneBoolean a => a
false
  else do
    let lhs2 :: [Lit]
lhs2 = [Lit] -> [Lit]
forall a. [a] -> [a]
reverse [Lit]
lhs1
        rhs2 :: [Bool]
rhs2 = Lit -> Bool -> [Bool]
forall a. Lit -> a -> [a]
replicate ([Lit] -> Lit
forall a. [a] -> Lit
forall (t :: * -> *) a. Foldable t => t a -> Lit
length [Lit]
lhs1 Lit -> Lit -> Lit
forall a. Num a => a -> a -> a
- [Bool] -> Lit
forall a. [a] -> Lit
forall (t :: * -> *) a. Foldable t => t a -> Lit
length [Bool]
rhs1) Bool
False [Bool] -> [Bool] -> [Bool]
forall a. [a] -> [a] -> [a]
++ [Bool] -> [Bool]
forall a. [a] -> [a]
reverse [Bool]
rhs1
        f :: [(Lit, Bool)] -> Formula
f [] = Formula
forall a. MonotoneBoolean a => a
true
        f ((Lit
x,Bool
False) : [(Lit, Bool)]
xs) = Lit -> Formula
Atom Lit
x Formula -> Formula -> Formula
forall a. MonotoneBoolean a => a -> a -> a
.||. [(Lit, Bool)] -> Formula
f [(Lit, Bool)]
xs
        f ((Lit
x,Bool
True) : [(Lit, Bool)]
xs) = Lit -> Formula
Atom Lit
x Formula -> Formula -> Formula
forall a. MonotoneBoolean a => a -> a -> a
.&&. [(Lit, Bool)] -> Formula
f [(Lit, Bool)]
xs
    Formula -> m Formula
forall a. a -> m a
forall (m :: * -> *) a. Monad m => a -> m a
return (Formula -> m Formula) -> Formula -> m Formula
forall a b. (a -> b) -> a -> b
$ [(Lit, Bool)] -> Formula
f ([Lit] -> [Bool] -> [(Lit, Bool)]
forall a b. [a] -> [b] -> [(a, b)]
zip [Lit]
lhs2 [Bool]
rhs2)
  where
    bits :: Integer -> [Bool]
    bits :: Integer -> [Bool]
bits Integer
n = Integer -> Lit -> [Bool]
forall {t}. (Num t, Bits t) => t -> Lit -> [Bool]
f Integer
n Lit
0
      where
        f :: t -> Lit -> [Bool]
f t
0 !Lit
_ = []
        f t
n Lit
i = t -> Lit -> Bool
forall a. Bits a => a -> Lit -> Bool
testBit t
n Lit
i Bool -> [Bool] -> [Bool]
forall a. a -> [a] -> [a]
: t -> Lit -> [Bool]
f (t -> Lit -> t
forall a. Bits a => a -> Lit -> a
clearBit t
n Lit
i) (Lit
iLit -> Lit -> Lit
forall a. Num a => a -> a -> a
+Lit
1)

encodePBLinSumAdder :: forall m. PrimMonad m => Tseitin.Encoder m -> SAT.PBLinSum -> m [SAT.Lit]
encodePBLinSumAdder :: forall (m :: * -> *).
PrimMonad m =>
Encoder m -> PBLinSum -> m [Lit]
encodePBLinSumAdder Encoder m
enc PBLinSum
lhs = do
  (MutVar (PrimState m) (Seq (SeqQueue m Lit))
buckets :: MutVar (PrimState m) (Seq (SQ.SeqQueue m SAT.Lit))) <- Seq (SeqQueue m Lit)
-> m (MutVar (PrimState m) (Seq (SeqQueue m Lit)))
forall (m :: * -> *) a.
PrimMonad m =>
a -> m (MutVar (PrimState m) a)
newMutVar Seq (SeqQueue m Lit)
forall a. Seq a
Seq.empty
  let insert :: Int -> Int -> m ()
      insert :: Lit -> Lit -> m ()
insert Lit
i Lit
x = do
        Seq (SeqQueue m Lit)
bs <- MutVar (PrimState m) (Seq (SeqQueue m Lit))
-> m (Seq (SeqQueue m Lit))
forall (m :: * -> *) a.
PrimMonad m =>
MutVar (PrimState m) a -> m a
readMutVar MutVar (PrimState m) (Seq (SeqQueue m Lit))
buckets
        let n :: Lit
n = Seq (SeqQueue m Lit) -> Lit
forall a. Seq a -> Lit
Seq.length Seq (SeqQueue m Lit)
bs
        SeqQueue m Lit
q <- if Lit
i Lit -> Lit -> Bool
forall a. Ord a => a -> a -> Bool
< Lit
n then do
               SeqQueue m Lit -> m (SeqQueue m Lit)
forall a. a -> m a
forall (m :: * -> *) a. Monad m => a -> m a
return (SeqQueue m Lit -> m (SeqQueue m Lit))
-> SeqQueue m Lit -> m (SeqQueue m Lit)
forall a b. (a -> b) -> a -> b
$ Seq (SeqQueue m Lit) -> Lit -> SeqQueue m Lit
forall a. Seq a -> Lit -> a
Seq.index Seq (SeqQueue m Lit)
bs Lit
i
             else do
               [SeqQueue m Lit]
qs <- Lit -> m (SeqQueue m Lit) -> m [SeqQueue m Lit]
forall (m :: * -> *) a. Applicative m => Lit -> m a -> m [a]
replicateM (Lit
iLit -> Lit -> Lit
forall a. Num a => a -> a -> a
+Lit
1 Lit -> Lit -> Lit
forall a. Num a => a -> a -> a
- Lit
n) m (SeqQueue m Lit)
forall q (m :: * -> *). NewFifo q m => m q
SQ.newFifo
               let bs' :: Seq (SeqQueue m Lit)
bs' = Seq (SeqQueue m Lit)
bs Seq (SeqQueue m Lit)
-> Seq (SeqQueue m Lit) -> Seq (SeqQueue m Lit)
forall a. Seq a -> Seq a -> Seq a
Seq.>< [SeqQueue m Lit] -> Seq (SeqQueue m Lit)
forall a. [a] -> Seq a
Seq.fromList [SeqQueue m Lit]
qs
               MutVar (PrimState m) (Seq (SeqQueue m Lit))
-> Seq (SeqQueue m Lit) -> m ()
forall (m :: * -> *) a.
PrimMonad m =>
MutVar (PrimState m) a -> a -> m ()
writeMutVar MutVar (PrimState m) (Seq (SeqQueue m Lit))
buckets Seq (SeqQueue m Lit)
bs'
               SeqQueue m Lit -> m (SeqQueue m Lit)
forall a. a -> m a
forall (m :: * -> *) a. Monad m => a -> m a
return (SeqQueue m Lit -> m (SeqQueue m Lit))
-> SeqQueue m Lit -> m (SeqQueue m Lit)
forall a b. (a -> b) -> a -> b
$ Seq (SeqQueue m Lit) -> Lit -> SeqQueue m Lit
forall a. Seq a -> Lit -> a
Seq.index Seq (SeqQueue m Lit)
bs' Lit
i
        SeqQueue m Lit -> Lit -> m ()
forall q (m :: * -> *) a. Enqueue q m a => q -> a -> m ()
SQ.enqueue SeqQueue m Lit
q Lit
x

      bits :: Integer -> [Int]
      bits :: Integer -> [Lit]
bits Integer
n = Integer -> Lit -> [Lit]
forall {t}. (Num t, Bits t) => t -> Lit -> [Lit]
f Integer
n Lit
0
        where
          f :: t -> Lit -> [Lit]
f t
0 !Lit
_ = []
          f t
n Lit
i
            | t -> Lit -> Bool
forall a. Bits a => a -> Lit -> Bool
testBit t
n Lit
i = Lit
i Lit -> [Lit] -> [Lit]
forall a. a -> [a] -> [a]
: t -> Lit -> [Lit]
f (t -> Lit -> t
forall a. Bits a => a -> Lit -> a
clearBit t
n Lit
i) (Lit
iLit -> Lit -> Lit
forall a. Num a => a -> a -> a
+Lit
1)
            | Bool
otherwise = t -> Lit -> [Lit]
f t
n (Lit
iLit -> Lit -> Lit
forall a. Num a => a -> a -> a
+Lit
1)

  PBLinSum -> (PBLinTerm -> m ()) -> m ()
forall (t :: * -> *) (m :: * -> *) a b.
(Foldable t, Monad m) =>
t a -> (a -> m b) -> m ()
forM_ PBLinSum
lhs ((PBLinTerm -> m ()) -> m ()) -> (PBLinTerm -> m ()) -> m ()
forall a b. (a -> b) -> a -> b
$ \(Integer
c,Lit
x) -> do
    [Lit] -> (Lit -> m ()) -> m ()
forall (t :: * -> *) (m :: * -> *) a b.
(Foldable t, Monad m) =>
t a -> (a -> m b) -> m ()
forM_ (Integer -> [Lit]
bits Integer
c) ((Lit -> m ()) -> m ()) -> (Lit -> m ()) -> m ()
forall a b. (a -> b) -> a -> b
$ \Lit
i -> Lit -> Lit -> m ()
insert Lit
i Lit
x

  let loop :: Lit -> [Lit] -> m [Lit]
loop Lit
i [Lit]
ret = do
        Seq (SeqQueue m Lit)
bs <- MutVar (PrimState m) (Seq (SeqQueue m Lit))
-> m (Seq (SeqQueue m Lit))
forall (m :: * -> *) a.
PrimMonad m =>
MutVar (PrimState m) a -> m a
readMutVar MutVar (PrimState m) (Seq (SeqQueue m Lit))
buckets
        let n :: Lit
n = Seq (SeqQueue m Lit) -> Lit
forall a. Seq a -> Lit
Seq.length Seq (SeqQueue m Lit)
bs
        if Lit
i Lit -> Lit -> Bool
forall a. Ord a => a -> a -> Bool
>= Lit
n then do
          [Lit] -> m [Lit]
forall a. a -> m a
forall (m :: * -> *) a. Monad m => a -> m a
return ([Lit] -> m [Lit]) -> [Lit] -> m [Lit]
forall a b. (a -> b) -> a -> b
$ [Lit] -> [Lit]
forall a. [a] -> [a]
reverse [Lit]
ret
        else do
          let q :: SeqQueue m Lit
q = Seq (SeqQueue m Lit) -> Lit -> SeqQueue m Lit
forall a. Seq a -> Lit -> a
Seq.index Seq (SeqQueue m Lit)
bs Lit
i
          Lit
m <- SeqQueue m Lit -> m Lit
forall q (m :: * -> *). QueueSize q m => q -> m Lit
SQ.queueSize SeqQueue m Lit
q
          case Lit
m of
            Lit
0 -> do
              Lit
b <- Encoder m -> [Lit] -> m Lit
forall (m :: * -> *). PrimMonad m => Encoder m -> [Lit] -> m Lit
Tseitin.encodeDisj Encoder m
enc [] -- False
              Lit -> [Lit] -> m [Lit]
loop (Lit
iLit -> Lit -> Lit
forall a. Num a => a -> a -> a
+Lit
1) (Lit
b Lit -> [Lit] -> [Lit]
forall a. a -> [a] -> [a]
: [Lit]
ret)
            Lit
1 -> do
              Lit
b <- Maybe Lit -> Lit
forall a. HasCallStack => Maybe a -> a
fromJust (Maybe Lit -> Lit) -> m (Maybe Lit) -> m Lit
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> SeqQueue m Lit -> m (Maybe Lit)
forall q (m :: * -> *) a. Dequeue q m a => q -> m (Maybe a)
SQ.dequeue SeqQueue m Lit
q
              Lit -> [Lit] -> m [Lit]
loop (Lit
iLit -> Lit -> Lit
forall a. Num a => a -> a -> a
+Lit
1) (Lit
b Lit -> [Lit] -> [Lit]
forall a. a -> [a] -> [a]
: [Lit]
ret)
            Lit
2 -> do
              Lit
b1 <- Maybe Lit -> Lit
forall a. HasCallStack => Maybe a -> a
fromJust (Maybe Lit -> Lit) -> m (Maybe Lit) -> m Lit
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> SeqQueue m Lit -> m (Maybe Lit)
forall q (m :: * -> *) a. Dequeue q m a => q -> m (Maybe a)
SQ.dequeue SeqQueue m Lit
q
              Lit
b2 <- Maybe Lit -> Lit
forall a. HasCallStack => Maybe a -> a
fromJust (Maybe Lit -> Lit) -> m (Maybe Lit) -> m Lit
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> SeqQueue m Lit -> m (Maybe Lit)
forall q (m :: * -> *) a. Dequeue q m a => q -> m (Maybe a)
SQ.dequeue SeqQueue m Lit
q
              Lit
s <- Encoder m -> Lit -> Lit -> m Lit
forall (m :: * -> *).
PrimMonad m =>
Encoder m -> Lit -> Lit -> m Lit
encodeHASum Encoder m
enc Lit
b1 Lit
b2
              Lit
c <- Encoder m -> Lit -> Lit -> m Lit
forall (m :: * -> *).
PrimMonad m =>
Encoder m -> Lit -> Lit -> m Lit
encodeHACarry Encoder m
enc Lit
b1 Lit
b2
              Lit -> Lit -> m ()
insert (Lit
iLit -> Lit -> Lit
forall a. Num a => a -> a -> a
+Lit
1) Lit
c
              Lit -> [Lit] -> m [Lit]
loop (Lit
iLit -> Lit -> Lit
forall a. Num a => a -> a -> a
+Lit
1) (Lit
s Lit -> [Lit] -> [Lit]
forall a. a -> [a] -> [a]
: [Lit]
ret)
            Lit
_ -> do
              Lit
b1 <- Maybe Lit -> Lit
forall a. HasCallStack => Maybe a -> a
fromJust (Maybe Lit -> Lit) -> m (Maybe Lit) -> m Lit
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> SeqQueue m Lit -> m (Maybe Lit)
forall q (m :: * -> *) a. Dequeue q m a => q -> m (Maybe a)
SQ.dequeue SeqQueue m Lit
q
              Lit
b2 <- Maybe Lit -> Lit
forall a. HasCallStack => Maybe a -> a
fromJust (Maybe Lit -> Lit) -> m (Maybe Lit) -> m Lit
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> SeqQueue m Lit -> m (Maybe Lit)
forall q (m :: * -> *) a. Dequeue q m a => q -> m (Maybe a)
SQ.dequeue SeqQueue m Lit
q
              Lit
b3 <- Maybe Lit -> Lit
forall a. HasCallStack => Maybe a -> a
fromJust (Maybe Lit -> Lit) -> m (Maybe Lit) -> m Lit
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> SeqQueue m Lit -> m (Maybe Lit)
forall q (m :: * -> *) a. Dequeue q m a => q -> m (Maybe a)
SQ.dequeue SeqQueue m Lit
q
              Lit
s <- Encoder m -> Lit -> Lit -> Lit -> m Lit
forall (m :: * -> *).
PrimMonad m =>
Encoder m -> Lit -> Lit -> Lit -> m Lit
Tseitin.encodeFASum Encoder m
enc Lit
b1 Lit
b2 Lit
b3
              Lit
c <- Encoder m -> Lit -> Lit -> Lit -> m Lit
forall (m :: * -> *).
PrimMonad m =>
Encoder m -> Lit -> Lit -> Lit -> m Lit
Tseitin.encodeFACarry Encoder m
enc Lit
b1 Lit
b2 Lit
b3
              Lit -> Lit -> m ()
insert Lit
i Lit
s
              Lit -> Lit -> m ()
insert (Lit
iLit -> Lit -> Lit
forall a. Num a => a -> a -> a
+Lit
1) Lit
c
              Lit -> [Lit] -> m [Lit]
loop Lit
i [Lit]
ret
  Lit -> [Lit] -> m [Lit]
loop Lit
0 []

encodeHASum :: PrimMonad m => Tseitin.Encoder m -> SAT.Lit -> SAT.Lit -> m SAT.Lit
encodeHASum :: forall (m :: * -> *).
PrimMonad m =>
Encoder m -> Lit -> Lit -> m Lit
encodeHASum = Encoder m -> Lit -> Lit -> m Lit
forall (m :: * -> *).
PrimMonad m =>
Encoder m -> Lit -> Lit -> m Lit
Tseitin.encodeXOR

encodeHACarry :: PrimMonad m => Tseitin.Encoder m -> SAT.Lit -> SAT.Lit -> m SAT.Lit
encodeHACarry :: forall (m :: * -> *).
PrimMonad m =>
Encoder m -> Lit -> Lit -> m Lit
encodeHACarry Encoder m
enc Lit
a Lit
b = Encoder m -> [Lit] -> m Lit
forall (m :: * -> *). PrimMonad m => Encoder m -> [Lit] -> m Lit
Tseitin.encodeConj Encoder m
enc [Lit
a,Lit
b]