{-# OPTIONS_GHC -Wall #-}
{-# OPTIONS_HADDOCK show-extensions #-}
{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE FlexibleInstances #-}
{-# LANGUAGE MultiParamTypeClasses #-}
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
forall (m :: * -> *) a.
PrimMonad m =>
MutVar (PrimState m) a -> (a -> a) -> m ()
modifyMutVar' MutVar (PrimState m) Var
ref (forall a. Num a => a -> a -> a
+Var
1)
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 = 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 = forall (m :: * -> *) a.
AddPBNL m a =>
a -> PBSum -> Integer -> m ()
SAT.addPBNLAtLeast PBStore m
enc [(Integer
1,[Var
l]) | Var
l <- [Var]
lhs] (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 = 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 = 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 = 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 forall a. Num a => a -> a -> a
- forall (t :: * -> *) a. (Foldable t, Num a) => t a -> a
sum [Integer
c | (Integer
c,[]) <- PBSum
lhs]
forall (m :: * -> *) a.
PrimMonad m =>
MutVar (PrimState m) a -> (a -> a) -> m ()
modifyMutVar' MutVar (PrimState m) (Seq Constraint)
ref (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 forall a. Num a => a -> a -> a
- forall (t :: * -> *) a. (Foldable t, Num a) => t a -> a
sum [Integer
c | (Integer
c,[]) <- PBSum
lhs]
forall (m :: * -> *) a.
PrimMonad m =>
MutVar (PrimState m) a -> (a -> a) -> m ()
modifyMutVar' MutVar (PrimState m) (Seq Constraint)
ref (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 <- forall (m :: * -> *) a.
PrimMonad m =>
a -> m (MutVar (PrimState m) a)
newMutVar Var
0
MutVar (PrimState m) (Seq Constraint)
ref2 <- forall (m :: * -> *) a.
PrimMonad m =>
a -> m (MutVar (PrimState m) a)
newMutVar forall a. Seq a
Seq.empty
forall (m :: * -> *) a. Monad m => a -> m a
return (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 <- forall (m :: * -> *) a.
PrimMonad m =>
MutVar (PrimState m) a -> m a
readMutVar MutVar (PrimState m) Var
ref1
Seq Constraint
cs <- forall (m :: * -> *) a.
PrimMonad m =>
MutVar (PrimState m) a -> m a
readMutVar MutVar (PrimState m) (Seq Constraint)
ref2
forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$
PBFile.Formula
{ pbObjectiveFunction :: Maybe PBSum
PBFile.pbObjectiveFunction = forall a. Maybe a
Nothing
, pbConstraints :: [Constraint]
PBFile.pbConstraints = forall (t :: * -> *) a. Foldable t => t a -> [a]
toList Seq Constraint
cs
, pbNumVars :: Var
PBFile.pbNumVars = Var
nv
, pbNumConstraints :: Var
PBFile.pbNumConstraints = forall a. Seq a -> Var
Seq.length Seq Constraint
cs
}