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) -- for specializations

{-# 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