{- |
This module provides functions for counting symbolic formulas and creating cardinality-constraints.

Internally this converts each given 'Expr' 'BoolSort' into a numerical 'Expr' using 'ite', then sums them up:

@
  count :: forall t f. (Functor f, Foldable f, Num (Expr t)) => f (Expr BoolSort) -> Expr t
  count = sum . fmap (\\b -> ite b 1 0)
@

Therefore additional information for the temporal summation may need to be provided.

E.g. if your logic is \"QF_LIA\" you would want @'count'' \@'IntSort' $ ...@

It is worth noting that some cardinality constraints use optimized encodings for special cases.
-}
module Language.Hasmtlib.Counting
(
  -- * Count
  count, count'

  -- * At-Least
, atLeast

  -- * Exactly
, exactly

  -- * At-Most
, atMost
, amoSqrt
, amoQuad
)
where

import Prelude hiding (not, (&&), (||), and, or, all, any)
import Language.Hasmtlib.Type.SMTSort
import Language.Hasmtlib.Type.Expr
import Language.Hasmtlib.Boolean
import Data.Foldable (toList)
import Data.List (transpose)
import Data.Proxy
import Control.Lens

-- | Wrapper for 'count' which takes a 'Proxy'.
count' :: forall t f. (Functor f, Foldable f, Num (Expr t)) => Proxy t -> f (Expr BoolSort) -> Expr t
count' :: forall (t :: SMTSort) (f :: * -> *).
(Functor f, Foldable f, Num (Expr t)) =>
Proxy t -> f (Expr 'BoolSort) -> Expr t
count' Proxy t
_ = f (Expr t) -> Expr t
forall a. Num a => f a -> a
forall (t :: * -> *) a. (Foldable t, Num a) => t a -> a
sum (f (Expr t) -> Expr t)
-> (f (Expr 'BoolSort) -> f (Expr t))
-> f (Expr 'BoolSort)
-> Expr t
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Expr 'BoolSort -> Expr t) -> f (Expr 'BoolSort) -> f (Expr t)
forall a b. (a -> b) -> f a -> f b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (\Expr 'BoolSort
b -> Expr 'BoolSort -> Expr t -> Expr t -> Expr t
forall b a. Iteable b a => b -> a -> a -> a
ite Expr 'BoolSort
b Expr t
1 Expr t
0)
{-# INLINE count' #-}

-- | Out of many bool-expressions build a formula which encodes how many of them are 'true'.
count :: forall t f. (Functor f, Foldable f, Num (Expr t)) => f (Expr BoolSort) -> Expr t
count :: forall (t :: SMTSort) (f :: * -> *).
(Functor f, Foldable f, Num (Expr t)) =>
f (Expr 'BoolSort) -> Expr t
count = Proxy t -> f (Expr 'BoolSort) -> Expr t
forall (t :: SMTSort) (f :: * -> *).
(Functor f, Foldable f, Num (Expr t)) =>
Proxy t -> f (Expr 'BoolSort) -> Expr t
count' (forall {k} (t :: k). Proxy t
forall (t :: SMTSort). Proxy t
Proxy @t)
{-# INLINE count #-}

-- | Out of many bool-expressions build a formula which encodes that __at most__ @k@ of them are 'true'.
--
--   'atMost' is defined as follows:
--
-- @
-- 'atMost' 0 = 'nand'
-- 'atMost' 1 = 'amoSqrt'
-- 'atMost' k = ('<=?' k) . 'count'
-- @
atMost  :: forall t f. (Functor f, Foldable f, Num (Expr t), Orderable (Expr t)) => Expr t -> f (Expr BoolSort) -> Expr BoolSort
atMost :: forall (t :: SMTSort) (f :: * -> *).
(Functor f, Foldable f, Num (Expr t), Orderable (Expr t)) =>
Expr t -> f (Expr 'BoolSort) -> Expr 'BoolSort
atMost Expr t
0 = f (Expr 'BoolSort) -> Expr 'BoolSort
forall (t :: * -> *) b. (Foldable t, Boolean b) => t b -> b
nand
atMost Expr t
1 = f (Expr 'BoolSort) -> Expr 'BoolSort
forall (t :: * -> *) b. (Foldable t, Boolean b) => t b -> b
amoSqrt
atMost Expr t
k = (Expr t -> Expr t -> Expr 'BoolSort
forall a. Orderable a => a -> a -> Expr 'BoolSort
<=? Expr t
k) (Expr t -> Expr 'BoolSort)
-> (f (Expr 'BoolSort) -> Expr t)
-> f (Expr 'BoolSort)
-> Expr 'BoolSort
forall b c a. (b -> c) -> (a -> b) -> a -> c
. f (Expr 'BoolSort) -> Expr t
forall (t :: SMTSort) (f :: * -> *).
(Functor f, Foldable f, Num (Expr t)) =>
f (Expr 'BoolSort) -> Expr t
count
{-# INLINE atMost #-}

-- | Out of many bool-expressions build a formula which encodes that __at least__ @k@ of them are 'true'.
--
--   'atLeast' is defined as follows:
--
-- @
-- 'atLeast' 0 = 'const' 'true'
-- 'atLeast' 1 = 'or'
-- 'atLeast' k = ('>=?' k) . 'count'
-- @
atLeast :: forall t f. (Functor f, Foldable f, Num (Expr t), Orderable (Expr t)) => Expr t -> f (Expr BoolSort) -> Expr BoolSort
atLeast :: forall (t :: SMTSort) (f :: * -> *).
(Functor f, Foldable f, Num (Expr t), Orderable (Expr t)) =>
Expr t -> f (Expr 'BoolSort) -> Expr 'BoolSort
atLeast Expr t
0 = Expr 'BoolSort -> f (Expr 'BoolSort) -> Expr 'BoolSort
forall a b. a -> b -> a
const Expr 'BoolSort
forall b. Boolean b => b
true
atLeast Expr t
1 = f (Expr 'BoolSort) -> Expr 'BoolSort
forall (t :: * -> *) b. (Foldable t, Boolean b) => t b -> b
or
atLeast Expr t
k = (Expr t -> Expr t -> Expr 'BoolSort
forall a. Orderable a => a -> a -> Expr 'BoolSort
>=? Expr t
k) (Expr t -> Expr 'BoolSort)
-> (f (Expr 'BoolSort) -> Expr t)
-> f (Expr 'BoolSort)
-> Expr 'BoolSort
forall b c a. (b -> c) -> (a -> b) -> a -> c
. f (Expr 'BoolSort) -> Expr t
forall (t :: SMTSort) (f :: * -> *).
(Functor f, Foldable f, Num (Expr t)) =>
f (Expr 'BoolSort) -> Expr t
count
{-# INLINE atLeast #-}

-- | Out of many bool-expressions build a formula which encodes that __exactly__ @k@ of them are 'true'.
--
--   'exactly' is defined as follows:
--
-- @
-- 'exactly' 0 xs = 'nand' xs
-- 'exactly' 1 xs = 'atLeast' \@t 1 xs '&&' 'atMost' \@t 1 xs
-- 'exactly' k xs = 'count' xs '===' k
-- @
exactly :: forall t f. (Functor f, Foldable f, Num (Expr t), Orderable (Expr t)) => Expr t -> f (Expr BoolSort) -> Expr BoolSort
exactly :: forall (t :: SMTSort) (f :: * -> *).
(Functor f, Foldable f, Num (Expr t), Orderable (Expr t)) =>
Expr t -> f (Expr 'BoolSort) -> Expr 'BoolSort
exactly Expr t
0 f (Expr 'BoolSort)
xs = f (Expr 'BoolSort) -> Expr 'BoolSort
forall (t :: * -> *) b. (Foldable t, Boolean b) => t b -> b
nand f (Expr 'BoolSort)
xs
exactly Expr t
1 f (Expr 'BoolSort)
xs = forall (t :: SMTSort) (f :: * -> *).
(Functor f, Foldable f, Num (Expr t), Orderable (Expr t)) =>
Expr t -> f (Expr 'BoolSort) -> Expr 'BoolSort
atLeast @t Expr t
1 f (Expr 'BoolSort)
xs Expr 'BoolSort -> Expr 'BoolSort -> Expr 'BoolSort
forall b. Boolean b => b -> b -> b
&& forall (t :: SMTSort) (f :: * -> *).
(Functor f, Foldable f, Num (Expr t), Orderable (Expr t)) =>
Expr t -> f (Expr 'BoolSort) -> Expr 'BoolSort
atMost @t Expr t
1 f (Expr 'BoolSort)
xs
exactly Expr t
k f (Expr 'BoolSort)
xs = f (Expr 'BoolSort) -> Expr t
forall (t :: SMTSort) (f :: * -> *).
(Functor f, Foldable f, Num (Expr t)) =>
f (Expr 'BoolSort) -> Expr t
count f (Expr 'BoolSort)
xs Expr t -> Expr t -> Expr 'BoolSort
forall a. Equatable a => a -> a -> Expr 'BoolSort
=== Expr t
k
{-# INLINE exactly #-}

-- | The squareroot-encoding, also called product-encoding, is a special encoding for @atMost 1@.
--
--   The original product-encoding provided by /Jingchao Chen/ in /A New SAT Encoding of the At-Most-One Constraint (2010)/
--   used auxiliary variables and would therefore be monadic.
--   It requires \( 2 \sqrt{n} + \mathcal{O}(\sqrt[4]{n}) \) auxiliary variables and
--   \( 2n + 4\sqrt{n} + \mathcal{O}(\sqrt[4]{n}) \) clauses.
--
--   To make this encoding pure, all auxiliary variables are replaced with a disjunction of size \( \mathcal{O}(\sqrt{n}) \).
--   Therefore zero auxiliary variables are required and technically clause-count remains the same, although the clauses get bigger.
amoSqrt :: (Foldable f, Boolean b) => f b -> b
amoSqrt :: forall (t :: * -> *) b. (Foldable t, Boolean b) => t b -> b
amoSqrt f b
xs
  | f b -> Int
forall a. f a -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length f b
xs Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
10 = [b] -> b
forall b. Boolean b => [b] -> b
amoQuad ([b] -> b) -> [b] -> b
forall a b. (a -> b) -> a -> b
$ f b -> [b]
forall a. f a -> [a]
forall (t :: * -> *) a. Foldable t => t a -> [a]
toList f b
xs
  | Bool
otherwise =
      let n :: Integer
n = Int -> Integer
forall a. Integral a => a -> Integer
toInteger (Int -> Integer) -> Int -> Integer
forall a b. (a -> b) -> a -> b
$ f b -> Int
forall a. f a -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length f b
xs
          p :: Integer
p = Double -> Integer
forall b. Integral b => Double -> b
forall a b. (RealFrac a, Integral b) => a -> b
ceiling (Double -> Integer) -> Double -> Integer
forall a b. (a -> b) -> a -> b
$ Double -> Double
forall a. Floating a => a -> a
sqrt (Double -> Double) -> Double -> Double
forall a b. (a -> b) -> a -> b
$ Integer -> Double
forall a. Num a => Integer -> a
fromInteger Integer
n
          rows :: [[b]]
rows = Int -> [b] -> [[b]]
forall {a}. Int -> [a] -> [[a]]
splitEvery (Integer -> Int
forall a. Num a => Integer -> a
fromInteger Integer
p) ([b] -> [[b]]) -> [b] -> [[b]]
forall a b. (a -> b) -> a -> b
$ f b -> [b]
forall a. f a -> [a]
forall (t :: * -> *) a. Foldable t => t a -> [a]
toList f b
xs
          columns :: [[b]]
columns = [[b]] -> [[b]]
forall a. [[a]] -> [[a]]
transpose [[b]]
rows
          vs :: [b]
vs = [b] -> b
forall (t :: * -> *) b. (Foldable t, Boolean b) => t b -> b
or ([b] -> b) -> [[b]] -> [b]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> [[b]]
rows
          us :: [b]
us = [b] -> b
forall (t :: * -> *) b. (Foldable t, Boolean b) => t b -> b
or ([b] -> b) -> [[b]] -> [b]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> [[b]]
columns
       in [b] -> b
forall (t :: * -> *) b. (Foldable t, Boolean b) => t b -> b
amoSqrt [b]
vs b -> b -> b
forall b. Boolean b => b -> b -> b
&& [b] -> b
forall (t :: * -> *) b. (Foldable t, Boolean b) => t b -> b
amoSqrt [b]
us b -> b -> b
forall b. Boolean b => b -> b -> b
&&
          [b] -> b
forall (t :: * -> *) b. (Foldable t, Boolean b) => t b -> b
and ((Int -> [b] -> b) -> [[b]] -> [b]
forall a b. (Int -> a -> b) -> [a] -> [b]
forall i (f :: * -> *) a b.
FunctorWithIndex i f =>
(i -> a -> b) -> f a -> f b
imap (\Int
j [b]
r -> [b] -> b
forall (t :: * -> *) b. (Foldable t, Boolean b) => t b -> b
and ([b] -> b) -> [b] -> b
forall a b. (a -> b) -> a -> b
$ (Int -> b -> b) -> [b] -> [b]
forall a b. (Int -> a -> b) -> [a] -> [b]
forall i (f :: * -> *) a b.
FunctorWithIndex i f =>
(i -> a -> b) -> f a -> f b
imap (\Int
i b
x -> (b
x b -> b -> b
forall b. Boolean b => b -> b -> b
==> [b]
us [b] -> Int -> b
forall a. HasCallStack => [a] -> Int -> a
!! Int
i) b -> b -> b
forall b. Boolean b => b -> b -> b
&& (b
x b -> b -> b
forall b. Boolean b => b -> b -> b
==> [b]
vs [b] -> Int -> b
forall a. HasCallStack => [a] -> Int -> a
!! Int
j)) [b]
r) [[b]]
rows)
  where
    splitEvery :: Int -> [a] -> [[a]]
splitEvery Int
n = ([a] -> Bool) -> [[a]] -> [[a]]
forall a. (a -> Bool) -> [a] -> [a]
takeWhile (Bool -> Bool
forall b. Boolean b => b -> b
not (Bool -> Bool) -> ([a] -> Bool) -> [a] -> Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. [a] -> Bool
forall a. [a] -> Bool
forall (t :: * -> *) a. Foldable t => t a -> Bool
null) ([[a]] -> [[a]]) -> ([a] -> [[a]]) -> [a] -> [[a]]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ([a] -> [a]) -> [[a]] -> [[a]]
forall a b. (a -> b) -> [a] -> [b]
map (Int -> [a] -> [a]
forall a. Int -> [a] -> [a]
take Int
n) ([[a]] -> [[a]]) -> ([a] -> [[a]]) -> [a] -> [[a]]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ([a] -> [a]) -> [a] -> [[a]]
forall a. (a -> a) -> a -> [a]
iterate (Int -> [a] -> [a]
forall a. Int -> [a] -> [a]
drop Int
n)

-- | The quadratic-encoding, also called pairwise-encoding, is a special encoding for @atMost 1@.
--
--   It's the naive encoding for the at-most-one-constraint and produces \( \binom{n}{2} \) clauses and no auxiliary variables..
amoQuad :: Boolean b => [b] -> b
amoQuad :: forall b. Boolean b => [b] -> b
amoQuad [b]
as = [b] -> b
forall (t :: * -> *) b. (Foldable t, Boolean b) => t b -> b
and ([b] -> b) -> [b] -> b
forall a b. (a -> b) -> a -> b
$ do
  [b]
ys <- Int -> [b] -> [[b]]
forall {a}. Int -> [a] -> [[a]]
subs Int
2 [b]
as
  b -> [b]
forall a. a -> [a]
forall (m :: * -> *) a. Monad m => a -> m a
return (b -> [b]) -> b -> [b]
forall a b. (a -> b) -> a -> b
$ (b -> b) -> [b] -> b
forall (t :: * -> *) b a.
(Foldable t, Boolean b) =>
(a -> b) -> t a -> b
any b -> b
forall b. Boolean b => b -> b
not [b]
ys
  where
    subs :: Int -> [a] -> [[a]]
    subs :: forall {a}. Int -> [a] -> [[a]]
subs Int
0 [a]
_ = [[]]
    subs Int
_ [] = []
    subs Int
k (a
x : [a]
xs) = ([a] -> [a]) -> [[a]] -> [[a]]
forall a b. (a -> b) -> [a] -> [b]
map (a
x :) (Int -> [a] -> [[a]]
forall {a}. Int -> [a] -> [[a]]
subs (Int
k Int -> Int -> Int
forall a. Num a => a -> a -> a
-Int
1) [a]
xs) [[a]] -> [[a]] -> [[a]]
forall a. Semigroup a => a -> a -> a
<> Int -> [a] -> [[a]]
forall {a}. Int -> [a] -> [[a]]
subs Int
k [a]
xs
{-# INLINE amoQuad #-}