{-# OPTIONS_GHC -Wall #-}
{-# OPTIONS_HADDOCK show-extensions #-}
{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE FlexibleInstances #-}
{-# LANGUAGE MultiParamTypeClasses #-}
-----------------------------------------------------------------------------
-- |
-- Module      :  ToySolver.SAT.Store.PB
-- Copyright   :  (c) Masahiro Sakai 2016
-- License     :  BSD-style
--
-- Maintainer  :  masahiro.sakai@gmail.com
-- Stability   :  experimental
-- Portability :  non-portable
--
-----------------------------------------------------------------------------
module ToySolver.SAT.Store.PB
  ( PBStore
  , newPBStore
  , getPBFormula
  ) where

import Control.Monad.Primitive
import Data.Foldable (toList)
import Data.Primitive.MutVar
import Data.Sequence (Seq, (|>))
import qualified Data.Sequence as Seq
import qualified Data.PseudoBoolean as PBFile
import qualified ToySolver.SAT.Types as SAT

data PBStore m = PBStore (MutVar (PrimState m) Int) (MutVar (PrimState m) (Seq PBFile.Constraint))

instance PrimMonad m => SAT.NewVar m (PBStore m) where
  newVar :: PBStore m -> m Var
newVar (PBStore MutVar (PrimState m) Var
ref MutVar (PrimState m) (Seq Constraint)
_) = do
    MutVar (PrimState m) Var -> (Var -> Var) -> m ()
forall (m :: * -> *) a.
PrimMonad m =>
MutVar (PrimState m) a -> (a -> a) -> m ()
modifyMutVar' MutVar (PrimState m) Var
ref (Var -> Var -> Var
forall a. Num a => a -> a -> a
+Var
1)
    MutVar (PrimState m) Var -> m Var
forall (m :: * -> *) a.
PrimMonad m =>
MutVar (PrimState m) a -> m a
readMutVar MutVar (PrimState m) Var
ref

instance PrimMonad m => SAT.AddClause m (PBStore m) where
  addClause :: PBStore m -> [Var] -> m ()
addClause PBStore m
enc [Var]
clause = PBStore m -> PBSum -> Integer -> m ()
forall (m :: * -> *) a.
AddPBNL m a =>
a -> PBSum -> Integer -> m ()
SAT.addPBNLAtLeast PBStore m
enc [(Integer
1,[Var
l]) | Var
l <- [Var]
clause] Integer
1

instance PrimMonad m => SAT.AddCardinality m (PBStore m) where
  addAtLeast :: PBStore m -> [Var] -> Var -> m ()
addAtLeast PBStore m
enc [Var]
lhs Var
rhs = PBStore m -> PBSum -> Integer -> m ()
forall (m :: * -> *) a.
AddPBNL m a =>
a -> PBSum -> Integer -> m ()
SAT.addPBNLAtLeast PBStore m
enc [(Integer
1,[Var
l]) | Var
l <- [Var]
lhs] (Var -> Integer
forall a b. (Integral a, Num b) => a -> b
fromIntegral Var
rhs)

instance PrimMonad m => SAT.AddPBLin m (PBStore m) where
  addPBAtLeast :: PBStore m -> PBLinSum -> Integer -> m ()
addPBAtLeast PBStore m
enc PBLinSum
lhs Integer
rhs = PBStore m -> PBSum -> Integer -> m ()
forall (m :: * -> *) a.
AddPBNL m a =>
a -> PBSum -> Integer -> m ()
SAT.addPBNLAtLeast PBStore m
enc [(Integer
c,[Var
l]) | (Integer
c,Var
l) <- PBLinSum
lhs] Integer
rhs
  addPBAtMost :: PBStore m -> PBLinSum -> Integer -> m ()
addPBAtMost PBStore m
enc PBLinSum
lhs Integer
rhs  = PBStore m -> PBSum -> Integer -> m ()
forall (m :: * -> *) a.
AddPBNL m a =>
a -> PBSum -> Integer -> m ()
SAT.addPBNLAtMost PBStore m
enc  [(Integer
c,[Var
l]) | (Integer
c,Var
l) <- PBLinSum
lhs] Integer
rhs
  addPBExactly :: PBStore m -> PBLinSum -> Integer -> m ()
addPBExactly PBStore m
enc PBLinSum
lhs Integer
rhs = PBStore m -> PBSum -> Integer -> m ()
forall (m :: * -> *) a.
AddPBNL m a =>
a -> PBSum -> Integer -> m ()
SAT.addPBNLExactly PBStore m
enc [(Integer
c,[Var
l]) | (Integer
c,Var
l) <- PBLinSum
lhs] Integer
rhs

instance PrimMonad m => SAT.AddPBNL m (PBStore m) where
  addPBNLAtLeast :: PBStore m -> PBSum -> Integer -> m ()
addPBNLAtLeast (PBStore MutVar (PrimState m) Var
_ MutVar (PrimState m) (Seq Constraint)
ref) PBSum
lhs Integer
rhs = do
    let lhs' :: PBSum
lhs' = [(Integer
c,[Var]
ls) | (Integer
c,ls :: [Var]
ls@(Var
_:[Var]
_)) <- PBSum
lhs]
        rhs' :: Integer
rhs' = Integer
rhs Integer -> Integer -> Integer
forall a. Num a => a -> a -> a
- [Integer] -> Integer
forall a. Num a => [a] -> a
forall (t :: * -> *) a. (Foldable t, Num a) => t a -> a
sum [Integer
c | (Integer
c,[]) <- PBSum
lhs]
    MutVar (PrimState m) (Seq Constraint)
-> (Seq Constraint -> Seq Constraint) -> m ()
forall (m :: * -> *) a.
PrimMonad m =>
MutVar (PrimState m) a -> (a -> a) -> m ()
modifyMutVar' MutVar (PrimState m) (Seq Constraint)
ref (Seq Constraint -> Constraint -> Seq Constraint
forall a. Seq a -> a -> Seq a
|> (PBSum
lhs', Op
PBFile.Ge, Integer
rhs'))
  addPBNLExactly :: PBStore m -> PBSum -> Integer -> m ()
addPBNLExactly (PBStore MutVar (PrimState m) Var
_ MutVar (PrimState m) (Seq Constraint)
ref) PBSum
lhs Integer
rhs = do
    let lhs' :: PBSum
lhs' = [(Integer
c,[Var]
ls) | (Integer
c,ls :: [Var]
ls@(Var
_:[Var]
_)) <- PBSum
lhs]
        rhs' :: Integer
rhs' = Integer
rhs Integer -> Integer -> Integer
forall a. Num a => a -> a -> a
- [Integer] -> Integer
forall a. Num a => [a] -> a
forall (t :: * -> *) a. (Foldable t, Num a) => t a -> a
sum [Integer
c | (Integer
c,[]) <- PBSum
lhs]
    MutVar (PrimState m) (Seq Constraint)
-> (Seq Constraint -> Seq Constraint) -> m ()
forall (m :: * -> *) a.
PrimMonad m =>
MutVar (PrimState m) a -> (a -> a) -> m ()
modifyMutVar' MutVar (PrimState m) (Seq Constraint)
ref (Seq Constraint -> Constraint -> Seq Constraint
forall a. Seq a -> a -> Seq a
|> (PBSum
lhs', Op
PBFile.Eq, Integer
rhs'))

newPBStore :: PrimMonad m => m (PBStore m)
newPBStore :: forall (m :: * -> *). PrimMonad m => m (PBStore m)
newPBStore = do
  MutVar (PrimState m) Var
ref1 <- Var -> m (MutVar (PrimState m) Var)
forall (m :: * -> *) a.
PrimMonad m =>
a -> m (MutVar (PrimState m) a)
newMutVar Var
0
  MutVar (PrimState m) (Seq Constraint)
ref2 <- Seq Constraint -> m (MutVar (PrimState m) (Seq Constraint))
forall (m :: * -> *) a.
PrimMonad m =>
a -> m (MutVar (PrimState m) a)
newMutVar Seq Constraint
forall a. Seq a
Seq.empty
  PBStore m -> m (PBStore m)
forall a. a -> m a
forall (m :: * -> *) a. Monad m => a -> m a
return (MutVar (PrimState m) Var
-> MutVar (PrimState m) (Seq Constraint) -> PBStore m
forall (m :: * -> *).
MutVar (PrimState m) Var
-> MutVar (PrimState m) (Seq Constraint) -> PBStore m
PBStore MutVar (PrimState m) Var
ref1 MutVar (PrimState m) (Seq Constraint)
ref2)

getPBFormula :: PrimMonad m => PBStore m -> m (PBFile.Formula)
getPBFormula :: forall (m :: * -> *). PrimMonad m => PBStore m -> m Formula
getPBFormula (PBStore MutVar (PrimState m) Var
ref1 MutVar (PrimState m) (Seq Constraint)
ref2) = do
  Var
nv <- MutVar (PrimState m) Var -> m Var
forall (m :: * -> *) a.
PrimMonad m =>
MutVar (PrimState m) a -> m a
readMutVar MutVar (PrimState m) Var
ref1
  Seq Constraint
cs <- MutVar (PrimState m) (Seq Constraint) -> m (Seq Constraint)
forall (m :: * -> *) a.
PrimMonad m =>
MutVar (PrimState m) a -> m a
readMutVar MutVar (PrimState m) (Seq Constraint)
ref2
  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
$
    PBFile.Formula
    { pbObjectiveFunction :: Maybe PBSum
PBFile.pbObjectiveFunction = Maybe PBSum
forall a. Maybe a
Nothing
    , pbConstraints :: [Constraint]
PBFile.pbConstraints = Seq Constraint -> [Constraint]
forall a. Seq a -> [a]
forall (t :: * -> *) a. Foldable t => t a -> [a]
toList Seq Constraint
cs
    , pbNumVars :: Var
PBFile.pbNumVars = Var
nv
    , pbNumConstraints :: Var
PBFile.pbNumConstraints = Seq Constraint -> Var
forall a. Seq a -> Var
Seq.length Seq Constraint
cs
    }