module OBDD.Linopt (linopt) where

import OBDD (OBDD, full_fold)

import qualified Data.Map.Strict as M
import Data.Bool (bool)

-- | solve the constrained linear optimisation problem:
-- returns an assignment that is a model of the BDD
-- and maximises the sum of weights of variables.
-- Keys missing from the weight map, but present in the BDD,
-- get weight zero.
linopt :: ( Ord v , Num w, Ord w ) 
       => OBDD v 
       -> M.Map v w 
       -> Maybe (w, M.Map v Bool)
linopt :: OBDD v -> Map v w -> Maybe (w, Map v Bool)
linopt OBDD v
d Map v w
m = Set v
-> (Bool -> Maybe (w, Map v Bool))
-> (v
    -> Maybe (w, Map v Bool)
    -> Maybe (w, Map v Bool)
    -> Maybe (w, Map v Bool))
-> OBDD v
-> Maybe (w, Map v Bool)
forall v a.
Ord v =>
Set v -> (Bool -> a) -> (v -> a -> a -> a) -> OBDD v -> a
full_fold (Map v w -> Set v
forall k a. Map k a -> Set k
M.keysSet Map v w
m) 
   ( Maybe (w, Map v Bool)
-> Maybe (w, Map v Bool) -> Bool -> Maybe (w, Map v Bool)
forall a. a -> a -> Bool -> a
bool Maybe (w, Map v Bool)
forall a. Maybe a
Nothing ( (w, Map v Bool) -> Maybe (w, Map v Bool)
forall a. a -> Maybe a
Just (w
0, Map v Bool
forall k a. Map k a
M.empty) ))
   ( \ v
v Maybe (w, Map v Bool)
ml Maybe (w, Map v Bool)
mr -> case (Maybe (w, Map v Bool)
ml,Maybe (w, Map v Bool)
mr) of
          (Just (w, Map v Bool)
l, Maybe (w, Map v Bool)
Nothing) -> (w, Map v Bool) -> Maybe (w, Map v Bool)
forall a. a -> Maybe a
Just ((w, Map v Bool) -> Maybe (w, Map v Bool))
-> (w, Map v Bool) -> Maybe (w, Map v Bool)
forall a b. (a -> b) -> a -> b
$ Map v w -> v -> (w, Map v Bool) -> (w, Map v Bool)
forall v w. (Ord v, Num w) => Map v w -> v -> Item v w -> Item v w
noadd Map v w
m v
v (w, Map v Bool)
l
          (Maybe (w, Map v Bool)
Nothing, Just (w, Map v Bool)
r) -> (w, Map v Bool) -> Maybe (w, Map v Bool)
forall a. a -> Maybe a
Just ((w, Map v Bool) -> Maybe (w, Map v Bool))
-> (w, Map v Bool) -> Maybe (w, Map v Bool)
forall a b. (a -> b) -> a -> b
$   Map v w -> v -> (w, Map v Bool) -> (w, Map v Bool)
forall v w. (Ord v, Num w) => Map v w -> v -> Item v w -> Item v w
add Map v w
m v
v (w, Map v Bool)
r
          (Just (w, Map v Bool)
l,  Just (w, Map v Bool)
r) -> (w, Map v Bool) -> Maybe (w, Map v Bool)
forall a. a -> Maybe a
Just ((w, Map v Bool) -> Maybe (w, Map v Bool))
-> (w, Map v Bool) -> Maybe (w, Map v Bool)
forall a b. (a -> b) -> a -> b
$
                  let l' :: (w, Map v Bool)
l' = Map v w -> v -> (w, Map v Bool) -> (w, Map v Bool)
forall v w. (Ord v, Num w) => Map v w -> v -> Item v w -> Item v w
noadd Map v w
m v
v (w, Map v Bool)
l
                      r' :: (w, Map v Bool)
r' =   Map v w -> v -> (w, Map v Bool) -> (w, Map v Bool)
forall v w. (Ord v, Num w) => Map v w -> v -> Item v w -> Item v w
add Map v w
m v
v (w, Map v Bool)
r
                  in  if (w, Map v Bool) -> w
forall a b. (a, b) -> a
fst (w, Map v Bool)
l' w -> w -> Bool
forall a. Ord a => a -> a -> Bool
>= (w, Map v Bool) -> w
forall a b. (a, b) -> a
fst (w, Map v Bool)
r' then (w, Map v Bool)
l' else (w, Map v Bool)
r'
          -- the following *can* happen for
          -- interpolated nodes directly above False:
          (Maybe (w, Map v Bool)
Nothing, Maybe (w, Map v Bool)
Nothing) -> Maybe (w, Map v Bool)
forall a. Maybe a
Nothing
       )
       OBDD v
d

type Item v w = (w, M.Map v Bool)

noadd, add :: (Ord v, Num w) 
           => M.Map v w -> v -> Item v w -> Item v w
noadd :: Map v w -> v -> Item v w -> Item v w
noadd Map v w
m v
v (w
w, Map v Bool
b) = 
  (w
w                          , v -> Bool -> Map v Bool -> Map v Bool
forall k a. Ord k => k -> a -> Map k a -> Map k a
M.insert v
v Bool
False Map v Bool
b)
add :: Map v w -> v -> Item v w -> Item v w
add   Map v w
m v
v (w
w, Map v Bool
b) = 
  (w
w w -> w -> w
forall a. Num a => a -> a -> a
+ w -> v -> Map v w -> w
forall k a. Ord k => a -> k -> Map k a -> a
M.findWithDefault w
0 v
v Map v w
m, v -> Bool -> Map v Bool -> Map v Bool
forall k a. Ord k => k -> a -> Map k a -> Map k a
M.insert v
v Bool
True  Map v Bool
b)