-- | functions in this module have no extra variables but exponential cost.

module Satchmo.Counting.Direct 

( atleast
, atmost
, exactly
, assert_implies_atmost
, assert_implies_exactly
)

where

import Satchmo.Boolean ( Boolean, MonadSAT )  
import qualified Satchmo.Boolean as B

import Control.Monad ( forM, forM_ )

select :: Int -> [a] -> [[a]]
select :: forall a. Int -> [a] -> [[a]]
select Int
0 [a]
xs = [[]]
select Int
k [] = []
select Int
k (a
x:[a]
xs) =
  Int -> [a] -> [[a]]
forall a. Int -> [a] -> [[a]]
select Int
k [a]
xs [[a]] -> [[a]] -> [[a]]
forall a. [a] -> [a] -> [a]
++ (([a] -> [a]) -> [[a]] -> [[a]]
forall a b. (a -> b) -> [a] -> [b]
map (a
xa -> [a] -> [a]
forall a. a -> [a] -> [a]
:) ([[a]] -> [[a]]) -> [[a]] -> [[a]]
forall a b. (a -> b) -> a -> b
$ Int -> [a] -> [[a]]
forall a. Int -> [a] -> [[a]]
select (Int
kInt -> Int -> Int
forall a. Num a => a -> a -> a
-Int
1) [a]
xs)

atleast :: MonadSAT m => Int -> [ Boolean ] -> m Boolean
atleast :: forall (m :: * -> *). MonadSAT m => Int -> [Boolean] -> m Boolean
atleast Int
k [Boolean]
xs = [Boolean] -> m Boolean
forall (m :: * -> *). MonadSAT m => [Boolean] -> m Boolean
B.or ([Boolean] -> m Boolean) -> m [Boolean] -> m Boolean
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< [[Boolean]] -> ([Boolean] -> m Boolean) -> m [Boolean]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
t a -> (a -> m b) -> m (t b)
forM (Int -> [Boolean] -> [[Boolean]]
forall a. Int -> [a] -> [[a]]
select Int
k [Boolean]
xs) [Boolean] -> m Boolean
forall (m :: * -> *). MonadSAT m => [Boolean] -> m Boolean
B.and

atmost :: MonadSAT m => Int -> [ Boolean ] -> m Boolean
atmost :: forall (m :: * -> *). MonadSAT m => Int -> [Boolean] -> m Boolean
atmost Int
k [Boolean]
xs = Int -> [Boolean] -> m Boolean
forall (m :: * -> *). MonadSAT m => Int -> [Boolean] -> m Boolean
atleast ([Boolean] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [Boolean]
xs Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
k) ([Boolean] -> m Boolean) -> [Boolean] -> m Boolean
forall a b. (a -> b) -> a -> b
$ (Boolean -> Boolean) -> [Boolean] -> [Boolean]
forall a b. (a -> b) -> [a] -> [b]
map Boolean -> Boolean
B.not [Boolean]
xs

exactly :: MonadSAT m => Int -> [ Boolean ] -> m Boolean
exactly :: forall (m :: * -> *). MonadSAT m => Int -> [Boolean] -> m Boolean
exactly Int
k [Boolean]
xs = do
  Boolean
this <- Int -> [Boolean] -> m Boolean
forall (m :: * -> *). MonadSAT m => Int -> [Boolean] -> m Boolean
atleast Int
k [Boolean]
xs
  Boolean
that <- Int -> [Boolean] -> m Boolean
forall (m :: * -> *). MonadSAT m => Int -> [Boolean] -> m Boolean
atmost Int
k [Boolean]
xs
  Boolean
this Boolean -> Boolean -> m Boolean
forall {m :: * -> *}. MonadSAT m => Boolean -> Boolean -> m Boolean
B.&& Boolean
that

-- | (and ys) implies (atmost k xs)
assert_implies_atmost :: [Boolean] -> Int -> [Boolean] -> m ()
assert_implies_atmost [Boolean]
ys Int
k [Boolean]
xs | Int
k Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
>= Int
0 = 
  [[Boolean]] -> ([Boolean] -> m ()) -> m ()
forall (t :: * -> *) (m :: * -> *) a b.
(Foldable t, Monad m) =>
t a -> (a -> m b) -> m ()
forM_ (Int -> [Boolean] -> [[Boolean]]
forall a. Int -> [a] -> [[a]]
select (Int
kInt -> Int -> Int
forall a. Num a => a -> a -> a
+Int
1) [Boolean]
xs) (([Boolean] -> m ()) -> m ()) -> ([Boolean] -> m ()) -> m ()
forall a b. (a -> b) -> a -> b
$ \ [Boolean]
sub -> do
    [Boolean] -> m ()
forall (m :: * -> *). MonadSAT m => [Boolean] -> m ()
B.assert ([Boolean] -> m ()) -> [Boolean] -> m ()
forall a b. (a -> b) -> a -> b
$ (Boolean -> Boolean) -> [Boolean] -> [Boolean]
forall a b. (a -> b) -> [a] -> [b]
map Boolean -> Boolean
B.not [Boolean]
ys [Boolean] -> [Boolean] -> [Boolean]
forall a. [a] -> [a] -> [a]
++ (Boolean -> Boolean) -> [Boolean] -> [Boolean]
forall a b. (a -> b) -> [a] -> [b]
map Boolean -> Boolean
B.not [Boolean]
sub
assert_implies_atmost [Boolean]
ys Int
k [Boolean]
_ =
  [Boolean] -> m ()
forall (m :: * -> *). MonadSAT m => [Boolean] -> m ()
B.assert ([Boolean] -> m ()) -> [Boolean] -> m ()
forall a b. (a -> b) -> a -> b
$ (Boolean -> Boolean) -> [Boolean] -> [Boolean]
forall a b. (a -> b) -> [a] -> [b]
map Boolean -> Boolean
B.not [Boolean]
ys

assert_implies_atleast :: [Boolean] -> Int -> [Boolean] -> m ()
assert_implies_atleast [Boolean]
ys Int
k [Boolean]
xs =
  [Boolean] -> Int -> [Boolean] -> m ()
forall {m :: * -> *}.
MonadSAT m =>
[Boolean] -> Int -> [Boolean] -> m ()
assert_implies_atmost [Boolean]
ys ([Boolean] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [Boolean]
xs Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
k) ((Boolean -> Boolean) -> [Boolean] -> [Boolean]
forall a b. (a -> b) -> [a] -> [b]
map Boolean -> Boolean
B.not [Boolean]
xs)

-- | asserting that  (and ys)  implies  (exactly k xs)
assert_implies_exactly :: [Boolean] -> Int -> [Boolean] -> m ()
assert_implies_exactly [Boolean]
ys Int
k [Boolean]
xs = do
  [Boolean] -> Int -> [Boolean] -> m ()
forall {m :: * -> *}.
MonadSAT m =>
[Boolean] -> Int -> [Boolean] -> m ()
assert_implies_atmost [Boolean]
ys Int
k [Boolean]
xs
  [Boolean] -> Int -> [Boolean] -> m ()
forall {m :: * -> *}.
MonadSAT m =>
[Boolean] -> Int -> [Boolean] -> m ()
assert_implies_atleast [Boolean]
ys Int
k [Boolean]
xs

-- | (atmost k xs) implies (or ys)
assert_atmost_implies :: [Boolean] -> Int -> [Boolean] -> m ()
assert_atmost_implies [Boolean]
xs Int
k [Boolean]
ys =
  [Boolean] -> Int -> [Boolean] -> m ()
forall {m :: * -> *}.
MonadSAT m =>
[Boolean] -> Int -> [Boolean] -> m ()
assert_implies_atleast ((Boolean -> Boolean) -> [Boolean] -> [Boolean]
forall a b. (a -> b) -> [a] -> [b]
map Boolean -> Boolean
B.not [Boolean]
ys) (Int
kInt -> Int -> Int
forall a. Num a => a -> a -> a
+Int
1) [Boolean]
xs

assert_atleast_implies :: [Boolean] -> Int -> [Boolean] -> m ()
assert_atleast_implies [Boolean]
xs Int
k [Boolean]
ys =
  [Boolean] -> Int -> [Boolean] -> m ()
forall {m :: * -> *}.
MonadSAT m =>
[Boolean] -> Int -> [Boolean] -> m ()
assert_implies_atmost ((Boolean -> Boolean) -> [Boolean] -> [Boolean]
forall a b. (a -> b) -> [a] -> [b]
map Boolean -> Boolean
B.not [Boolean]
ys) (Int
kInt -> Int -> Int
forall a. Num a => a -> a -> a
+Int
1) [Boolean]
xs