module Satchmo.Counting.Binary
( atleast
, atmost
, exactly
, count
)
where
import Prelude hiding ( and, or, not )
import Satchmo.Boolean
import Satchmo.Binary
import Satchmo.SAT ( SAT)
{-# specialize inline atleast :: Int -> [ Boolean] -> SAT Boolean #-}
{-# specialize inline atmost :: Int -> [ Boolean] -> SAT Boolean #-}
{-# specialize inline exactly :: Int -> [ Boolean] -> SAT Boolean #-}
{-# specialize inline count :: [ Boolean] -> SAT Number #-}
count :: MonadSAT m => [ Boolean ] -> m Number
count :: forall (m :: * -> *). MonadSAT m => [Boolean] -> m Number
count [Boolean]
bits
= m Number -> (Number -> Number -> m Number) -> [Number] -> m Number
forall (m :: * -> *) a.
Monad m =>
m a -> (a -> a -> m a) -> [a] -> m a
collect (Integer -> m Number
forall (m :: * -> *). MonadSAT m => Integer -> m Number
Satchmo.Binary.constant Integer
0) Number -> Number -> m Number
forall (m :: * -> *). MonadSAT m => Number -> Number -> m Number
Satchmo.Binary.add
([Number] -> m Number) -> [Number] -> m Number
forall a b. (a -> b) -> a -> b
$ (Boolean -> Number) -> [Boolean] -> [Number]
forall a b. (a -> b) -> [a] -> [b]
map ( \ Boolean
bit -> [Boolean] -> Number
Satchmo.Binary.make [Boolean
bit] )
([Boolean] -> [Number]) -> [Boolean] -> [Number]
forall a b. (a -> b) -> a -> b
$ [Boolean]
bits
data NumCarries =
NumCarries { NumCarries -> Number
num:: Number,NumCarries -> [Boolean]
carries:: [Boolean]}
zro :: NumCarries
zro = NumCarries {num :: Number
num=[Boolean] -> Number
make [], carries :: [Boolean]
carries=[] }
mke :: a -> Boolean -> NumCarries
mke a
0 Boolean
b = NumCarries {num :: Number
num=[Boolean] -> Number
make[],carries :: [Boolean]
carries=[Boolean
b]}
mke a
w Boolean
b | a
w a -> a -> Bool
forall a. Ord a => a -> a -> Bool
> a
0 = NumCarries {num :: Number
num=[Boolean] -> Number
make[Boolean
b],carries :: [Boolean]
carries=[]}
pls :: Int -> NumCarries -> NumCarries -> m NumCarries
pls Int
w NumCarries
x NumCarries
y = do
Number
z <- Number -> Number -> m Number
forall (m :: * -> *). MonadSAT m => Number -> Number -> m Number
Satchmo.Binary.add (NumCarries -> Number
num NumCarries
x) (NumCarries -> Number
num NumCarries
y)
let ([Boolean]
pre,[Boolean]
post) = Int -> [Boolean] -> ([Boolean], [Boolean])
forall a. Int -> [a] -> ([a], [a])
splitAt Int
w ([Boolean] -> ([Boolean], [Boolean]))
-> [Boolean] -> ([Boolean], [Boolean])
forall a b. (a -> b) -> a -> b
$ Number -> [Boolean]
bits Number
z
NumCarries -> m NumCarries
forall a. a -> m a
forall (m :: * -> *) a. Monad m => a -> m a
return (NumCarries -> m NumCarries) -> NumCarries -> m NumCarries
forall a b. (a -> b) -> a -> b
$ NumCarries
{ num :: Number
num = [Boolean] -> Number
make [Boolean]
pre
, carries :: [Boolean]
carries = [Boolean]
post [Boolean] -> [Boolean] -> [Boolean]
forall a. [a] -> [a] -> [a]
++ NumCarries -> [Boolean]
carries NumCarries
x [Boolean] -> [Boolean] -> [Boolean]
forall a. [a] -> [a] -> [a]
++ NumCarries -> [Boolean]
carries NumCarries
y
}
count_and_carry :: Int -> [Boolean] -> m NumCarries
count_and_carry Int
width [Boolean]
bits
= m NumCarries
-> (NumCarries -> NumCarries -> m NumCarries)
-> [NumCarries]
-> m NumCarries
forall (m :: * -> *) a.
Monad m =>
m a -> (a -> a -> m a) -> [a] -> m a
collect (NumCarries -> m NumCarries
forall a. a -> m a
forall (m :: * -> *) a. Monad m => a -> m a
return NumCarries
zro) (Int -> NumCarries -> NumCarries -> m NumCarries
forall {m :: * -> *}.
MonadSAT m =>
Int -> NumCarries -> NumCarries -> m NumCarries
pls Int
width) ([NumCarries] -> m NumCarries) -> [NumCarries] -> m NumCarries
forall a b. (a -> b) -> a -> b
$ (Boolean -> NumCarries) -> [Boolean] -> [NumCarries]
forall a b. (a -> b) -> [a] -> [b]
map (Int -> Boolean -> NumCarries
forall {a}. (Num a, Ord a) => a -> Boolean -> NumCarries
mke Int
width) [Boolean]
bits
collect :: Monad m => m a -> (a -> a -> m a) -> [a] -> m a
collect :: forall (m :: * -> *) a.
Monad m =>
m a -> (a -> a -> m a) -> [a] -> m a
collect m a
z a -> a -> m a
b [a]
xs = case [a]
xs of
[] -> m a
z
[a
x] -> a -> m a
forall a. a -> m a
forall (m :: * -> *) a. Monad m => a -> m a
return a
x
(a
x:a
y:[a]
zs) -> a -> a -> m a
b a
x a
y m a -> (a -> m a) -> m a
forall a b. m a -> (a -> m b) -> m b
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= \ a
c -> m a -> (a -> a -> m a) -> [a] -> m a
forall (m :: * -> *) a.
Monad m =>
m a -> (a -> a -> m a) -> [a] -> m a
collect m a
z a -> a -> m a
b ([a]
zs [a] -> [a] -> [a]
forall a. [a] -> [a] -> [a]
++ [a
c])
atleast :: MonadSAT m => Int -> [ Boolean ] -> m Boolean
atleast :: forall (m :: * -> *). MonadSAT m => Int -> [Boolean] -> m Boolean
atleast Int
k [Boolean]
xs = Bool
-> (Number -> Number -> m Boolean) -> Int -> [Boolean] -> m Boolean
forall (m :: * -> *).
MonadSAT m =>
Bool
-> (Number -> Number -> m Boolean) -> Int -> [Boolean] -> m Boolean
common Bool
True Number -> Number -> m Boolean
forall (m :: * -> *). MonadSAT m => Number -> Number -> m Boolean
ge Int
k [Boolean]
xs
atmost :: MonadSAT m => Int -> [ Boolean ] -> m Boolean
atmost :: forall (m :: * -> *). MonadSAT m => Int -> [Boolean] -> m Boolean
atmost Int
k [Boolean]
xs = Bool
-> (Number -> Number -> m Boolean) -> Int -> [Boolean] -> m Boolean
forall (m :: * -> *).
MonadSAT m =>
Bool
-> (Number -> Number -> m Boolean) -> Int -> [Boolean] -> m Boolean
common Bool
False Number -> Number -> m Boolean
forall (m :: * -> *). MonadSAT m => Number -> Number -> m Boolean
le Int
k [Boolean]
xs
exactly :: MonadSAT m => Int -> [ Boolean ] -> m Boolean
exactly :: forall (m :: * -> *). MonadSAT m => Int -> [Boolean] -> m Boolean
exactly Int
k [Boolean]
xs = Bool
-> (Number -> Number -> m Boolean) -> Int -> [Boolean] -> m Boolean
forall (m :: * -> *).
MonadSAT m =>
Bool
-> (Number -> Number -> m Boolean) -> Int -> [Boolean] -> m Boolean
common Bool
False Number -> Number -> m Boolean
forall (m :: * -> *). MonadSAT m => Number -> Number -> m Boolean
eq Int
k [Boolean]
xs
common :: MonadSAT m
=> Bool
-> (Number -> Number -> m Boolean)
-> Int -> [ Boolean ] -> m Boolean
common :: forall (m :: * -> *).
MonadSAT m =>
Bool
-> (Number -> Number -> m Boolean) -> Int -> [Boolean] -> m Boolean
common Bool
may_overflow Number -> Number -> m Boolean
cmp Int
k [Boolean]
xs = do
let bk :: [Bool]
bk = Integer -> [Bool]
Satchmo.Binary.toBinary (Integer -> [Bool]) -> Integer -> [Bool]
forall a b. (a -> b) -> a -> b
$ Int -> Integer
forall a b. (Integral a, Num b) => a -> b
fromIntegral Int
k
NumCarries { num :: NumCarries -> Number
num=Number
n,carries :: NumCarries -> [Boolean]
carries=[Boolean]
cs} <-
Int -> [Boolean] -> m NumCarries
forall {m :: * -> *}.
MonadSAT m =>
Int -> [Boolean] -> m NumCarries
count_and_carry ([Bool] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [Bool]
bk) [Boolean]
xs
Number
goal <- Integer -> m Number
forall (m :: * -> *). MonadSAT m => Integer -> m Number
Satchmo.Binary.constant (Integer -> m Number) -> Integer -> m Number
forall a b. (a -> b) -> a -> b
$ Int -> Integer
forall a b. (Integral a, Num b) => a -> b
fromIntegral Int
k
Boolean
ok <- Number -> Number -> m Boolean
cmp Number
n Number
goal
if Bool
may_overflow
then [Boolean] -> m Boolean
forall (m :: * -> *). MonadSAT m => [Boolean] -> m Boolean
or ([Boolean] -> m Boolean) -> [Boolean] -> m Boolean
forall a b. (a -> b) -> a -> b
$ Boolean
ok Boolean -> [Boolean] -> [Boolean]
forall a. a -> [a] -> [a]
: [Boolean]
cs
else [Boolean] -> m Boolean
forall (m :: * -> *). MonadSAT m => [Boolean] -> m Boolean
and ([Boolean] -> m Boolean) -> [Boolean] -> m Boolean
forall a b. (a -> b) -> a -> b
$ Boolean
ok Boolean -> [Boolean] -> [Boolean]
forall a. a -> [a] -> [a]
: (Boolean -> Boolean) -> [Boolean] -> [Boolean]
forall a b. (a -> b) -> [a] -> [b]
map Boolean -> Boolean
not [Boolean]
cs