Builder.hs

-- | Hash-consed BDD manager and the core BDD operations.
module Imp.BDD.Builder
  ( BDDManager
  , nodeTable
  , BDDM
    -- Manager operations.
  , emptyManager
  , newVar
  , lookupNode
    -- BDD operations.
  , bddNot
  , bddAnd
  , bddOr
  , bddAny
  , bddIte
  , bddRestrict
  ) where

import Control.Monad.State.Strict
import Data.IntMap.Strict (IntMap)
import qualified Data.IntMap.Strict as IntMap
import Data.Map.Strict (Map)
import qualified Data.Map.Strict as Map

import Imp.BDD

-- | Hash-consed node store with memo caches.
data BDDManager = BDDManager
  { nextNodeId    :: !Int
    -- | The hash-consed node store, keyed by t'NodeId'.
  , nodeTable     :: !(IntMap BDDNode)
  , uniqueTable   :: !(Map BDDNode NodeId)
  , iteCache      :: !(Map (BDD, BDD, BDD) BDD)
  , restrictCache :: !(Map (BDD, VarLabel, Bool) BDD)
  , nextVarId     :: !Int
  } deriving stock (Show)

-- | A manager with no nodes or variables.
emptyManager :: BDDManager
emptyManager = BDDManager 0 IntMap.empty Map.empty Map.empty Map.empty 0

-- | BDD-manager state monad.
type BDDM = State BDDManager

-- | Allocate a fresh BDD variable and return its positive literal
--   together with its t'VarLabel'.
newVar :: BDDM (BDD, VarLabel)
newVar = do
  v <- VarLabel <$> gets nextVarId
  modify' (\mgr -> mgr { nextVarId = nextVarId mgr + 1 })
  bdd <- mkNode (BDDNode v BDDFalse BDDTrue)
  return (bdd, v)

-- | Negate a BDD in O(1) via complemented edges.
bddNot :: BDD -> BDD
bddNot BDDTrue     = BDDFalse
bddNot BDDFalse    = BDDTrue
bddNot (BDDRef n)  = BDDComp n
bddNot (BDDComp n) = BDDRef n

-- | Conjunction of two BDDs.
bddAnd :: BDD -> BDD -> BDDM BDD
bddAnd a b = bddIte a b BDDFalse

-- | Disjunction of two BDDs.
bddOr :: BDD -> BDD -> BDDM BDD
bddOr a b = bddIte a BDDTrue b

-- | OR together a list of BDDs using a balanced tree fold.
bddAny :: [BDD] -> BDDM BDD
bddAny []  = return BDDFalse
bddAny [g] = return g
bddAny xs  = do
  let (l, r) = splitAt (length xs `div` 2) xs
  lv <- bddAny l
  rv <- bddAny r
  bddOr lv rv

-- | BDD if-then-else terminal cases and memoisation.
bddIte :: BDD -> BDD -> BDD -> BDDM BDD
bddIte f g h
  | f == BDDTrue                   = return g
  | f == BDDFalse                  = return h
  | g == BDDTrue  && h == BDDFalse = return f
  | g == BDDFalse && h == BDDTrue  = return (bddNot f)
  | g == h                         = return g
  | otherwise = do
      cached <- gets (Map.lookup (f, g, h) . iteCache)
      case cached of
        Just result -> return result
        Nothing -> do
          result <- bddIteExpand f g h
          modify' $ \s -> s { iteCache = Map.insert (f, g, h) result (iteCache s) }
          return result

-- | Shannon expansion on the topmost variable.
bddIteExpand :: BDD -> BDD -> BDD -> BDDM BDD
bddIteExpand f g h = do
  nf <- lookupNode f
  ng <- lookupNode g
  nh <- lookupNode h
  -- f is never constant here, so at least one node exists.
  let topV = minimum [bddVar n | Just n <- [nf, ng, nh]]
      cofactor branch bdd node = case node of
        Just n | bddVar n == topV -> if branch then bddHigh n else bddLow n
        _                         -> bdd
  lo <- bddIte (cofactor False f nf) (cofactor False g ng) (cofactor False h nh)
  hi <- bddIte (cofactor True  f nf) (cofactor True  g ng) (cofactor True  h nh)
  mkNode (BDDNode topV lo hi)

-- | Restrict a BDD by fixing a variable to True or False, with memoisation.
bddRestrict :: BDD -> VarLabel -> Bool -> BDDM BDD
bddRestrict BDDTrue  _ _ = return BDDTrue
bddRestrict BDDFalse _ _ = return BDDFalse
bddRestrict bdd var val = do
  cached <- gets (Map.lookup (bdd, var, val) . restrictCache)
  case cached of
    Just result -> return result
    Nothing -> do
      result <- bddRestrictExpand bdd var val
      modify' $ \s ->
        s { restrictCache = Map.insert (bdd, var, val) result (restrictCache s) }
      return result

-- | Restrict one level and recurse.
bddRestrictExpand :: BDD -> VarLabel -> Bool -> BDDM BDD
bddRestrictExpand bdd var val = do
  mnode <- lookupNode bdd
  case mnode of
    Nothing -> return bdd
    Just (BDDNode v lo hi)
      | v == var  -> return $ if val then hi else lo
      | v > var   -> return bdd  -- var not in this sub-BDD
      | otherwise -> do
          lo' <- bddRestrict lo var val
          hi' <- bddRestrict hi var val
          mkNode (BDDNode v lo' hi')

-- | Look up a BDD reference.
lookupNode :: BDD -> BDDM (Maybe BDDNode)
lookupNode BDDTrue              = return Nothing
lookupNode BDDFalse             = return Nothing
lookupNode (BDDRef (NodeId n))  = IntMap.lookup n <$> gets nodeTable
lookupNode (BDDComp (NodeId n)) = do
  node <- IntMap.lookup n <$> gets nodeTable
  return $ (\(BDDNode v lo hi) -> BDDNode v (bddNot lo) (bddNot hi)) <$> node

-- | Hash-consed node creation with complemented-edge normalization.
mkNode :: BDDNode -> BDDM BDD
mkNode node@(BDDNode v lo hi)
  | lo == hi = return lo
  -- Normalize nodes to never have a complemented low edge.
  | BDDTrue <- lo = bddNot <$> mkNode (BDDNode v (bddNot lo) (bddNot hi))
  | BDDComp _ <- lo = bddNot <$> mkNode (BDDNode v (bddNot lo) (bddNot hi))
  | otherwise = do
      mgr <- get
      case Map.lookup node (uniqueTable mgr) of
        Just nid -> return (BDDRef nid)
        Nothing  -> do
          let nidInt = nextNodeId mgr
              nid    = NodeId nidInt
          put mgr { nextNodeId  = nidInt + 1
                  , nodeTable   = IntMap.insert nidInt node (nodeTable mgr)
                  , uniqueTable = Map.insert node nid (uniqueTable mgr)
                  }
          return (BDDRef nid)