{-# LANGUAGE CPP #-}
module BDD(
        BDD, bddFalse, bddTrue, bddVar,
        bddAnd, bddOr, bddNot, bddImplies, bddEquiv,
        bddRestrict, bddAllSat,
        bddIsTrue, bddIsFalse, bddIsIf) where

import Control.Monad(liftM, ap)

import Data.List(sort)
import qualified Data.Map as M
import qualified Data.IntMap as I


type Unique = Int
--type UniquePair = (Int, Int)
type UniquePair = Int
mkPair :: Unique -> Unique -> UniquePair
--mkPair x y = (x, y)
mkPair x y = x * 65536 + y

data BDD a = N !Unique a (BDD a) (BDD a) | L !Bool
        deriving (Show)

bddFalse, bddTrue :: BDD a
bddFalse = L False
bddTrue  = L True

bddAnd, bddOr :: Ord a => BDD a -> BDD a -> BDD a
bddAnd a b = apply (&&) a b
bddOr  a b = apply (||) a b

bddNot :: BDD a -> BDD a
--bddNot e = bddOp (\ x y -> not x) e bddFalse                -- inefficient
bddNot e = fst (flip I.empty e)
  where flip m (L b) = (L (not b), m)
        flip m (N u x t f) =
            case I.lookup u m of
            Just b -> (b, m)
            Nothing ->
                case flip m t of
                (t', m') ->
                    case flip m' f of
                    (f', m'') ->
                        let n' = N u x t' f'
                        in  (n', I.insert u n' m'')

bddImplies, bddEquiv :: (Ord a) => BDD a -> BDD a -> BDD a
bddImplies a b = apply (\ x y -> not x || y) a b
bddEquiv a b = apply (==) a b

bddVar :: a -> BDD a
bddVar v = N 0 v bddFalse bddTrue

--bddOp :: (Ord a) => (Bool -> Bool -> Bool) -> BDD a -> BDD a -> BDD a
--bddOp op e1 e2 = run $ app op e1 e2

-- Check
bddIsTrue :: BDD a -> Bool
bddIsTrue (L b) = b
bddIsTrue _ = False

bddIsFalse :: BDD a -> Bool
bddIsFalse (L b) = not b
bddIsFalse _ = False

bddIsIf :: BDD a -> Maybe (a, BDD a, BDD a)
bddIsIf (L _) = Nothing
bddIsIf (N _ v f t) = Just (v, t, f)

bddAllSat :: BDD a -> [[(a, Bool)]]
bddAllSat (L True) = [ [] ]
bddAllSat (L False) = []
bddAllSat (N _ v l r) = map ((v,False):) (bddAllSat l) ++ map ((v,True):) (bddAllSat r)

bddRestrict :: (Ord a) => [(a, Bool)] -> BDD a -> BDD a
bddRestrict vs b = run $ restrict (sort vs) b

-- bddSubstitute
-- bddIsIso

----

nodeNo :: BDD a -> Unique

nodeNo (L False) = 0
nodeNo (L True)  = 1
nodeNo (N u _ _ _) = u

----

-- The state contains 3 parts:
--   the next unique number to use
--   a cache of N nodes to avoid recreating an existing N node
data S a = S !Unique (M.Map (UniquePair, a) (BDD a))

newtype M v a = M (S v -> (S v, a))

instance Monad (M v) where
    return = pure
    M a >>= f = M $ \ s ->
        case a s of
        (s', b) ->
            let M f' = f b
            in  f' s'

instance Functor (M v) where
  fmap = liftM

instance Applicative (M v) where
  pure a = M $ \ s -> (s, a)
  (<*>) = ap

run :: M v a -> a
run (M m) =
    case m (S 2 M.empty) of
    (_, a) -> a

genSym :: M a Unique
genSym = M $ \ (S u h) -> (S (u+1) h, u)

find :: (Ord a) => a -> BDD a -> BDD a -> M a (Maybe (BDD a))
find v b1 b2 = M $ \ s@(S _ h) -> (s, M.lookup (mkPair (nodeNo b1) (nodeNo b2), v) h)

insert :: (Ord a) => a -> BDD a -> BDD a -> BDD a -> M a ()
insert v b1 b2 n = M $ \ s@(S u h) -> (S u (M.insert (mkPair (nodeNo b1) (nodeNo b2), v) n h), ())

----

apply :: (Ord a) => (Bool -> Bool -> Bool) -> BDD a -> BDD a -> BDD a
apply op e1 e2 = run $ app op e1 e2

-- XXX should pre-compute what to do when constants are encountered
app :: (Ord a) => (Bool -> Bool -> Bool) -> BDD a -> BDD a -> M a (BDD a)
app op e1 e2 = doApp op e1 e2

doApp :: (Ord a) => (Bool -> Bool -> Bool) -> BDD a -> BDD a -> M a (BDD a)
doApp op (L b1) (L b2) = return (L (op b1 b2))
doApp op e@(L b) (N _ v f t) = do
    f' <- doApp op e f
    t' <- doApp op e t
    node v f' t'
doApp op (N _ v f t) e@(L b) = do
    f' <- doApp op f e
    t' <- doApp op t e
    node v f' t'
doApp op e1@(N _ v1 f1 t1) e2@(N _ v2 f2 t2) =
    case v1 `compare` v2 of
    LT -> do
        f' <- doApp op f1 e2
        t' <- doApp op t1 e2
        node v1 f' t'
    EQ -> do
        f' <- doApp op f1 f2
        t' <- doApp op t1 t2
        node v1 f' t'
    GT -> do
        f' <- doApp op e1 f2
        t' <- doApp op e1 t2
        node v2 f' t'

node :: (Ord a) => a -> BDD a -> BDD a -> M a (BDD a)
node v a b =
    if nodeNo a == nodeNo b then
        return a
    else do
        mv <- find v a b
        case mv of
         Just n -> return n
         Nothing -> do
            u <- genSym
            let n = N u v a b
            insert v a b n
            return n

restrict :: (Ord a) => [(a, Bool)] -> BDD a -> M a (BDD a)
restrict [] b = return b
restrict vs b@(L _) = return b
restrict ((v,c):vs) b@(N _ v' f t) =
    case v `compare` v' of
    LT -> restrict vs b
    EQ -> if c then restrict vs t else restrict vs f
    GT -> do
        -- XXX exponential
        f' <- restrict vs f
        t' <- restrict vs t
        node v' f' t'

{-

-- Some BDD test code to compute the solution to the N queen problem.

data V = V Int Int
        deriving (Eq, Ord)

instance Show V where
        show (V i j) = "v" ++ itos i ++ "_" ++ itos j

queens :: Int -> [[V]]
queens = map (map fst . filter snd) . bddAllSat . genBoard

genBoard :: Int -> BDD V
genBoard n =
    bddAnds
      [
        checkRow n i j `bddAnd`
        checkCol n i j `bddAnd`
        checkDiagL n i j `bddAnd`
        checkDiagR n i j
      | i <- [1..n], j <- [1..n]
      ]
  `bddAnd`
    -- make sure there is (at least) one queen on each row
    bddAnds [ bddOrs [ bddVar (V i j) | j <- [1..n] ] | i <- [1..n] ]

checkRow n i j = var i j `implies` bddAnds [ notVar i l | l <- [1..n], l /= j ]
checkCol n i j = var i j `implies` bddAnds [ notVar k j | k <- [1..n], k /= i ]
checkDiagL n i j = var i j `implies` bddAnds [ notVar k m | k <- [1..n], let m = j+k-i, 1 <= m && m <= n, k /= i ]
checkDiagR n i j = var i j `implies` bddAnds [ notVar k m | k <- [1..n], let m = j+i-k, 1 <= m && m <= n, k /= i ]

implies x y = bddNot x `bddOr` y

notVar i j = bddNot (var i j)

var i j = bddVar (V i j)

bddAnds xs = foldr bddAnd bddTrue  xs
bddOrs  xs = foldr bddOr  bddFalse xs

main = print (queens 5)

-}