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
  = forall (m :: * -> *) a.
Monad m =>
m a -> (a -> a -> m a) -> [a] -> m a
collect (forall (m :: * -> *). MonadSAT m => Integer -> m Number
Satchmo.Binary.constant Integer
0) forall (m :: * -> *). MonadSAT m => Number -> Number -> m Number
Satchmo.Binary.add
  forall a b. (a -> b) -> a -> b
$ forall a b. (a -> b) -> [a] -> [b]
map ( \ Boolean
bit -> [Boolean] -> Number
Satchmo.Binary.make [Boolean
bit] )
  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 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 <- 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) = forall a. Int -> [a] -> ([a], [a])
splitAt Int
w forall a b. (a -> b) -> a -> b
$ Number -> [Boolean]
bits Number
z
  forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$ NumCarries
     { num :: Number
num = [Boolean] -> Number
make [Boolean]
pre
     , carries :: [Boolean]
carries = [Boolean]
post forall a. [a] -> [a] -> [a]
++ NumCarries -> [Boolean]
carries NumCarries
x forall a. [a] -> [a] -> [a]
++ NumCarries -> [Boolean]
carries NumCarries
y
     }

count_and_carry :: Int -> [Boolean] -> m NumCarries
count_and_carry Int
width [Boolean]
bits 
  = forall (m :: * -> *) a.
Monad m =>
m a -> (a -> a -> m a) -> [a] -> m a
collect (forall (m :: * -> *) a. Monad m => a -> m a
return NumCarries
zro) (forall {m :: * -> *}.
MonadSAT m =>
Int -> NumCarries -> NumCarries -> m NumCarries
pls Int
width) forall a b. (a -> b) -> a -> b
$ forall a b. (a -> b) -> [a] -> [b]
map (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] -> 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 forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= \ a
c -> 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 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 = forall (m :: * -> *).
MonadSAT m =>
Bool
-> (Number -> Number -> m Boolean) -> Int -> [Boolean] -> m Boolean
common Bool
True 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 = forall (m :: * -> *).
MonadSAT m =>
Bool
-> (Number -> Number -> m Boolean) -> Int -> [Boolean] -> m Boolean
common Bool
False 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 = forall (m :: * -> *).
MonadSAT m =>
Bool
-> (Number -> Number -> m Boolean) -> Int -> [Boolean] -> m Boolean
common Bool
False 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 forall a b. (a -> b) -> a -> b
$ 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} <-
    forall {m :: * -> *}.
MonadSAT m =>
Int -> [Boolean] -> m NumCarries
count_and_carry (forall (t :: * -> *) a. Foldable t => t a -> Int
length [Bool]
bk) [Boolean]
xs
  Number
goal <- forall (m :: * -> *). MonadSAT m => Integer -> m Number
Satchmo.Binary.constant forall a b. (a -> b) -> a -> b
$ 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 forall (m :: * -> *). MonadSAT m => [Boolean] -> m Boolean
or forall a b. (a -> b) -> a -> b
$ Boolean
ok forall a. a -> [a] -> [a]
: [Boolean]
cs
    else forall (m :: * -> *). MonadSAT m => [Boolean] -> m Boolean
and forall a b. (a -> b) -> a -> b
$ Boolean
ok forall a. a -> [a] -> [a]
: forall a b. (a -> b) -> [a] -> [b]
map Boolean -> Boolean
not [Boolean]
cs