BDD.hs

-- | Tests the BDD manager and WMC algorithm.
module Test.BDD (tests) where

import Control.Monad.State.Strict (gets, runState)
import Data.IntMap.Strict (IntMap)
import qualified Data.IntMap.Strict as IntMap
import Test.Tasty
import Test.Tasty.HUnit

import Imp.BDD
import Imp.BDD.Builder
import Imp.BDD.WMC
import Imp.Semiring

probsOf :: [(VarLabel, Double)] -> IntMap Weight
probsOf probs = IntMap.fromList [ (unVarLabel vl, Prob p) | (vl, p) <- probs ]

probF :: Weight -> (ProbS, ProbS)
probF w = case w of
  Prob p   -> (ProbS (1 - p), ProbS p)
  Knight _ -> error "probF: not used in these tests"

-- | Run a manager action from an empty manager.
runBDD :: BDDM a -> a
runBDD m = fst (runState m emptyManager)

nodeCount :: BDDM Int
nodeCount = gets (IntMap.size . nodeTable)

tests :: TestTree
tests = testGroup "BDD"
  [ testGroup "Variables and negation"
    [ testCase "newVar creates distinct variables" $
        runBDD (do
          (v1, l1) <- newVar
          (v2, l2) <- newVar
          return (v1 /= v2 && l1 /= l2)) @? "variables should be distinct"

    , testCase "bddNot on terminals" $ do
        bddNot BDDTrue  @?= BDDFalse
        bddNot BDDFalse @?= BDDTrue

    , testCase "bddNot is an involution" $
        runBDD (do
          (v, _) <- newVar
          return (bddNot (bddNot v) == v)) @? "not . not = id"
    ]

  , testGroup "ITE terminal cases"
    [ testCase "bddAnd/bddOr with constants" $
        runBDD (do
          (v, _) <- newVar
          at <- bddAnd v BDDTrue
          af <- bddAnd v BDDFalse
          of' <- bddOr v BDDFalse
          ot <- bddOr v BDDTrue
          return [at == v, af == BDDFalse, of' == v, ot == BDDTrue])
          @?= [True, True, True, True]

    , testCase "bddIte f False True = bddNot f" $
        runBDD (do
          (v, _) <- newVar
          r <- bddIte v BDDFalse BDDTrue
          return (r == bddNot v)) @? "ite f 0 1 should be the complemented edge"

    , testCase "bddIte f g g = g" $
        runBDD (do
          (a, _) <- newVar
          (b, _) <- newVar
          r <- bddIte a b b
          return (r == b)) @? "both branches equal"

    , testCase "bddAny of no disjuncts is False" $
        runBDD (bddAny []) @?= BDDFalse

    , testCase "bddAny of one disjunct is itself" $
        runBDD (do
          (v, _) <- newVar
          r <- bddAny [v]
          return (r == v)) @? "singleton disjunction"

    , testCase "bddAny agrees with a right fold of bddOr" $
        runBDD (do
          (a, _) <- newVar
          (b, _) <- newVar
          (c, _) <- newVar
          balanced <- bddAny [a, b, c]
          folded <- bddOr a =<< bddOr b c
          return (balanced == folded)) @? "balanced fold = linear fold"
    ]

  , testGroup "Canonicalisation and sharing"
    [ testCase "De Morgan: a AND b and NOT (NOT a OR NOT b) are the same node" $
        runBDD (do
          (a, _) <- newVar
          (b, _) <- newVar
          conj <- bddAnd a b
          disj <- bddOr (bddNot a) (bddNot b)
          return (conj == bddNot disj)) @? "complement pair must not be interned twice"

    , testCase "De Morgan interns 3 nodes, not 4" $
        runBDD (do
          (a, _) <- newVar
          (b, _) <- newVar
          _ <- bddAnd a b
          _ <- bddOr (bddNot a) (bddNot b)
          nodeCount) @?= 3

    , testCase "bddAnd is commutative and shares the node" $
        runBDD (do
          (a, _) <- newVar
          (b, _) <- newVar
          ab <- bddAnd a b
          ba <- bddAnd b a
          n <- nodeCount
          return (ab == ba, n)) @?= (True, 3)
    ]

  , testGroup "Restrict"
    [ testCase "restrict a AND b on the top variable" $
        runBDD (do
          (a, la) <- newVar
          (b, _) <- newVar
          ab <- bddAnd a b
          t <- bddRestrict ab la True
          f <- bddRestrict ab la False
          return (t == b, f)) @?= (True, BDDFalse)

    , testCase "restrict a AND b on the lower variable" $
        runBDD (do
          (a, _) <- newVar
          (b, lb) <- newVar
          ab <- bddAnd a b
          t <- bddRestrict ab lb True
          f <- bddRestrict ab lb False
          return (t == a, f)) @?= (True, BDDFalse)

    , testCase "restrict on a variable above the root is a no-op" $
        runBDD (do
          (_, la) <- newVar
          (b, _) <- newVar
          r <- bddRestrict b la True
          return (r == b)) @? "variable not in this sub-BDD"

    , testCase "restrict terminals" $
        runBDD (do
          (_, la) <- newVar
          t <- bddRestrict BDDTrue la True
          f <- bddRestrict BDDFalse la True
          return (t, f)) @?= (BDDTrue, BDDFalse)
    ]

  , testGroup "lookupNode"
    [ testCase "terminals have no node" $
        runBDD (do
          t <- lookupNode BDDTrue
          f <- lookupNode BDDFalse
          return (t, f)) @?= (Nothing, Nothing)

    , testCase "a complemented reference negates both children" $
        runBDD (do
          (v, _) <- newVar
          pos <- lookupNode v
          neg <- lookupNode (bddNot v)
          return (pos, neg))
          @?= ( Just (BDDNode (VarLabel 0) BDDFalse BDDTrue)
              , Just (BDDNode (VarLabel 0) BDDTrue BDDFalse) )
    ]

  , testGroup "WMC"
    [ testCase "constants" $ do
        unProb (wmc probF emptyManager IntMap.empty BDDTrue) @?= 1.0
        unProb (wmc probF emptyManager IntMap.empty BDDFalse) @?= 0.0

    , testCase "single variable" $ do
        let ((var, vl), mgr) = runState newVar emptyManager
        unProb (wmc probF mgr (probsOf [(vl, 0.7)]) var) @?= 0.7

    , testCase "AND of two variables" $ do
        let ((r, l1, l2), mgr) = runState (do
              (a, la) <- newVar
              (b, lb) <- newVar
              ab <- bddAnd a b
              return (ab, la, lb)) emptyManager
        unProb (wmc probF mgr (probsOf [(l1, 0.5), (l2, 0.5)]) r) @?= 0.25

    , testCase "OR of two variables (complemented edge under a weight)" $ do
        let ((r, l1, l2), mgr) = runState (do
              (a, la) <- newVar
              (b, lb) <- newVar
              ab <- bddOr a b
              return (ab, la, lb)) emptyManager
        unProb (wmc probF mgr (probsOf [(l1, 0.5), (l2, 0.5)]) r) @?= 0.75

    , testCase "wmcBatch over a shared manager" $ do
        let ((bdds, l1, l2), mgr) = runState (do
              (a, la) <- newVar
              (b, lb) <- newVar
              conj <- bddAnd a b
              disj <- bddOr a b
              return ([a, b, conj, disj], la, lb)) emptyManager
        map unProb (wmcBatch probF mgr (probsOf [(l1, 0.5), (l2, 0.25)]) bdds)
          @?= [0.5, 0.25, 0.125, 0.625]
    ]
  ]