{-# 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
  , encodePBLinAtLeastAdder
  ) 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 <- forall (m :: * -> *).
PrimMonad m =>
Encoder m -> PBLinAtLeast -> m Formula
encodePBLinAtLeastAdder' Encoder m
enc PBLinAtLeast
constr
  forall (m :: * -> *). PrimMonad m => Encoder m -> Formula -> m ()
Tseitin.addFormula Encoder m
enc Formula
formula

encodePBLinAtLeastAdder :: PrimMonad m => Tseitin.Encoder m -> SAT.PBLinAtLeast -> m SAT.Lit
encodePBLinAtLeastAdder :: forall (m :: * -> *).
PrimMonad m =>
Encoder m -> PBLinAtLeast -> m Lit
encodePBLinAtLeastAdder Encoder m
enc PBLinAtLeast
constr = do
  Formula
formula <- forall (m :: * -> *).
PrimMonad m =>
Encoder m -> PBLinAtLeast -> m Formula
encodePBLinAtLeastAdder' Encoder m
enc PBLinAtLeast
constr
  forall (m :: * -> *). PrimMonad m => Encoder m -> Formula -> m Lit
Tseitin.encodeFormula Encoder m
enc 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 forall a. Ord a => a -> a -> Bool
<= Integer
0 = forall (m :: * -> *) a. Monad m => a -> m a
return forall a. MonotoneBoolean a => a
true
encodePBLinAtLeastAdder' Encoder m
enc (PBLinSum
lhs,Integer
rhs) = do
  [Lit]
lhs1 <- 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 forall (t :: * -> *) a. Foldable t => t a -> Lit
length [Lit]
lhs1 forall a. Ord a => a -> a -> Bool
< forall (t :: * -> *) a. Foldable t => t a -> Lit
length [Bool]
rhs1 then do
    forall (m :: * -> *) a. Monad m => a -> m a
return forall a. MonotoneBoolean a => a
false
  else do
    let lhs2 :: [Lit]
lhs2 = forall a. [a] -> [a]
reverse [Lit]
lhs1
        rhs2 :: [Bool]
rhs2 = forall a. Lit -> a -> [a]
replicate (forall (t :: * -> *) a. Foldable t => t a -> Lit
length [Lit]
lhs1 forall a. Num a => a -> a -> a
- forall (t :: * -> *) a. Foldable t => t a -> Lit
length [Bool]
rhs1) Bool
False forall a. [a] -> [a] -> [a]
++ forall a. [a] -> [a]
reverse [Bool]
rhs1
        f :: [(Lit, Bool)] -> Formula
f [] = forall a. MonotoneBoolean a => a
true
        f ((Lit
x,Bool
False) : [(Lit, Bool)]
xs) = Lit -> Formula
Atom Lit
x 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 forall a. MonotoneBoolean a => a -> a -> a
.&&. [(Lit, Bool)] -> Formula
f [(Lit, Bool)]
xs
    forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$ [(Lit, Bool)] -> Formula
f (forall a b. [a] -> [b] -> [(a, b)]
zip [Lit]
lhs2 [Bool]
rhs2)
  where
    bits :: Integer -> [Bool]
    bits :: Integer -> [Bool]
bits Integer
n = 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 = forall a. Bits a => a -> Lit -> Bool
testBit t
n Lit
i forall a. a -> [a] -> [a]
: t -> Lit -> [Bool]
f (forall a. Bits a => a -> Lit -> a
clearBit t
n Lit
i) (Lit
iforall 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))) <- forall (m :: * -> *) a.
PrimMonad m =>
a -> m (MutVar (PrimState m) a)
newMutVar 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 <- forall (m :: * -> *) a.
PrimMonad m =>
MutVar (PrimState m) a -> m a
readMutVar MutVar (PrimState m) (Seq (SeqQueue m Lit))
buckets
        let n :: Lit
n = forall a. Seq a -> Lit
Seq.length Seq (SeqQueue m Lit)
bs
        SeqQueue m Lit
q <- if Lit
i forall a. Ord a => a -> a -> Bool
< Lit
n then do
               forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$ forall a. Seq a -> Lit -> a
Seq.index Seq (SeqQueue m Lit)
bs Lit
i
             else do
               [SeqQueue m Lit]
qs <- forall (m :: * -> *) a. Applicative m => Lit -> m a -> m [a]
replicateM (Lit
iforall a. Num a => a -> a -> a
+Lit
1 forall a. Num a => a -> a -> a
- Lit
n) forall q (m :: * -> *). NewFifo q m => m q
SQ.newFifo
               let bs' :: Seq (SeqQueue m Lit)
bs' = Seq (SeqQueue m Lit)
bs forall a. Seq a -> Seq a -> Seq a
Seq.>< forall a. [a] -> Seq a
Seq.fromList [SeqQueue m Lit]
qs
               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'
               forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$ forall a. Seq a -> Lit -> a
Seq.index Seq (SeqQueue m Lit)
bs' Lit
i
        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 = 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
            | forall a. Bits a => a -> Lit -> Bool
testBit t
n Lit
i = Lit
i forall a. a -> [a] -> [a]
: t -> Lit -> [Lit]
f (forall a. Bits a => a -> Lit -> a
clearBit t
n Lit
i) (Lit
iforall a. Num a => a -> a -> a
+Lit
1)
            | Bool
otherwise = t -> Lit -> [Lit]
f t
n (Lit
iforall a. Num a => a -> a -> a
+Lit
1)

  forall (t :: * -> *) (m :: * -> *) a b.
(Foldable t, Monad m) =>
t a -> (a -> m b) -> m ()
forM_ PBLinSum
lhs forall a b. (a -> b) -> a -> b
$ \(Integer
c,Lit
x) -> do
    forall (t :: * -> *) (m :: * -> *) a b.
(Foldable t, Monad m) =>
t a -> (a -> m b) -> m ()
forM_ (Integer -> [Lit]
bits Integer
c) 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 <- forall (m :: * -> *) a.
PrimMonad m =>
MutVar (PrimState m) a -> m a
readMutVar MutVar (PrimState m) (Seq (SeqQueue m Lit))
buckets
        let n :: Lit
n = forall a. Seq a -> Lit
Seq.length Seq (SeqQueue m Lit)
bs
        if Lit
i forall a. Ord a => a -> a -> Bool
>= Lit
n then do
          forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$ forall a. [a] -> [a]
reverse [Lit]
ret
        else do
          let q :: SeqQueue m Lit
q = forall a. Seq a -> Lit -> a
Seq.index Seq (SeqQueue m Lit)
bs Lit
i
          Lit
m <- forall q (m :: * -> *). QueueSize q m => q -> m Lit
SQ.queueSize SeqQueue m Lit
q
          case Lit
m of
            Lit
0 -> do
              Lit
b <- forall (m :: * -> *). PrimMonad m => Encoder m -> [Lit] -> m Lit
Tseitin.encodeDisj Encoder m
enc [] -- False
              Lit -> [Lit] -> m [Lit]
loop (Lit
iforall a. Num a => a -> a -> a
+Lit
1) (Lit
b forall a. a -> [a] -> [a]
: [Lit]
ret)
            Lit
1 -> do
              Lit
b <- forall a. HasCallStack => Maybe a -> a
fromJust forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> forall q (m :: * -> *) a. Dequeue q m a => q -> m (Maybe a)
SQ.dequeue SeqQueue m Lit
q
              Lit -> [Lit] -> m [Lit]
loop (Lit
iforall a. Num a => a -> a -> a
+Lit
1) (Lit
b forall a. a -> [a] -> [a]
: [Lit]
ret)
            Lit
2 -> do
              Lit
b1 <- forall a. HasCallStack => Maybe a -> a
fromJust forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> forall q (m :: * -> *) a. Dequeue q m a => q -> m (Maybe a)
SQ.dequeue SeqQueue m Lit
q
              Lit
b2 <- forall a. HasCallStack => Maybe a -> a
fromJust forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> forall q (m :: * -> *) a. Dequeue q m a => q -> m (Maybe a)
SQ.dequeue SeqQueue m Lit
q
              Lit
s <- forall (m :: * -> *).
PrimMonad m =>
Encoder m -> Lit -> Lit -> m Lit
encodeHASum Encoder m
enc Lit
b1 Lit
b2
              Lit
c <- forall (m :: * -> *).
PrimMonad m =>
Encoder m -> Lit -> Lit -> m Lit
encodeHACarry Encoder m
enc Lit
b1 Lit
b2
              Lit -> Lit -> m ()
insert (Lit
iforall a. Num a => a -> a -> a
+Lit
1) Lit
c
              Lit -> [Lit] -> m [Lit]
loop (Lit
iforall a. Num a => a -> a -> a
+Lit
1) (Lit
s forall a. a -> [a] -> [a]
: [Lit]
ret)
            Lit
_ -> do
              Lit
b1 <- forall a. HasCallStack => Maybe a -> a
fromJust forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> forall q (m :: * -> *) a. Dequeue q m a => q -> m (Maybe a)
SQ.dequeue SeqQueue m Lit
q
              Lit
b2 <- forall a. HasCallStack => Maybe a -> a
fromJust forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> forall q (m :: * -> *) a. Dequeue q m a => q -> m (Maybe a)
SQ.dequeue SeqQueue m Lit
q
              Lit
b3 <- forall a. HasCallStack => Maybe a -> a
fromJust forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> forall q (m :: * -> *) a. Dequeue q m a => q -> m (Maybe a)
SQ.dequeue SeqQueue m Lit
q
              Lit
s <- forall (m :: * -> *).
PrimMonad m =>
Encoder m -> Lit -> Lit -> Lit -> m Lit
Tseitin.encodeFASum Encoder m
enc Lit
b1 Lit
b2 Lit
b3
              Lit
c <- 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
iforall 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 = 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 = forall (m :: * -> *). PrimMonad m => Encoder m -> [Lit] -> m Lit
Tseitin.encodeConj Encoder m
enc [Lit
a,Lit
b]