{-# LANGUAGE DeriveGeneric #-}
{-# LANGUAGE LambdaCase #-}
{-# LANGUAGE RankNTypes #-}

--------------------------------------------------------------------------------
-- |
-- Module      : Data.ZSDD
-- Description : Zero-suppressed decision diagrams
-- Copyright   : (c) Eddie Jones, 2020
-- License     : GPL-3
-- Maintainer  : eddiejones2108@gmail.com
-- Stability   : experimental
-- Portability : POSIX
--
-- Zero suppressed decision diagrams provide an efficient representation of boolean
-- formulas build from atomic propositions.
--
-- == Warning
--
-- The space used by a diagram grows with each operation.
-- Therefore, it is best to use 'runDiagram' as soon as possible (typically when there
-- is no possibility of further overlapping terms).
--
-- == Implementation
-- Internally it is a directed acyclic graph, thereby allowing for sharing
-- of subterms.
-- Operations have a monotonic effect on the diagram and their results are 
-- cached.
-- The implementations was outlined by this paper:
--
--  * Shin-ichi Minato, \"/Zero-suppressed BDDs for set manipulation in combinatorial problems/\",
--    DAC '93: Proceedings of the 30th international Design Automation Conference
--    <https://dl.acm.org/doi/10.1145/157485.164890>
--------------------------------------------------------------------------
module Data.ZSDD
  ( -- * Data types
    Prop (Top, Bot),
    Invertible,
    Diagram,
    runDiagram,

    -- * Construction
    atomic,
    negate,
    assign,
    union,
    intersect,
    difference,

    -- * Query
    isEmpty,
    isNotEmpty,
    count,

    -- * Traversal
    mapAtomic,
    atomicToProp,
  )
where

import Control.Monad
import Control.Monad.State
import qualified Data.HashMap.Lazy as H
import Data.Hashable
import qualified Data.IntMap as I
import GHC.Generics (Generic)
import Prelude hiding (map, negate)

type ID = Int

-- | A proposition parameterised by it's diagram
data Prop d l
  = Top
  | Bot
  | ID ID
  deriving (Eq, Ord, Generic)

instance Hashable (Prop d l)

-- An internal node
data Node d l = Node
  { label :: l,
    lo :: Prop d l,
    hi :: Prop d l
  }
  deriving (Eq, Generic)

instance Hashable l => Hashable (Node d l)

-- | Atomic propositions that have an inverse
class Invertible l where
  neg :: l -> l

-- | The ambient decision diagram of propositions
type Diagram d l =
  StateT
    ( ID, -- Fresh node ID
      H.HashMap (Ops l) (Prop d l), -- Cached operations
      H.HashMap (Node d l) Int, -- node to ID
      I.IntMap (Node d l) -- ID to node
    )

data Ops l
  = Union ID ID
  | Intersect ID ID
  | Difference ID ID
  | DifferenceT ID
  | Negate ID l
  | Assign ID l Bool
  deriving (Eq, Ord, Generic)

instance Hashable l => Hashable (Ops l)

-- | Extract information from a digram without leaking
runDiagram :: Monad m => (forall d. Diagram d l m a) -> m a
runDiagram f = evalStateT f (0, H.empty, H.empty, I.empty)

-- Lookup operation in cache and run otherwise
memo :: (Ord l, Hashable l, Monad m) => Ops l -> Diagram d l m (Prop d l) -> Diagram d l m (Prop d l)
memo o k = do
  (i, mem, h, ns) <- get
  case H.lookup o mem of
    Nothing -> do
      p <- k
      put (i, H.insert o p mem, h, ns)
      return p
    Just p -> return p

-- Get a fresh id
fresh :: Monad m => Diagram d l m ID
fresh = do
  (i, mem, h, ns) <- get
  put (i + 1, mem, h, ns)
  return i

-- Extract internal node information from a proposition
withNode :: Monad m => ID -> (Node d l -> Diagram d l m a) -> Diagram d l m a
withNode i f = do
  (_, _, _, ns) <- get
  case I.lookup i ns of
    Nothing -> error "Node not found!"
    Just n -> f n

-- Insert a new node
insertNode :: (Eq l, Hashable l, Monad m) => ID -> Node d l -> Diagram d l m ()
insertNode i n = do
  (j, m, h, ns) <- get
  put (j, m, H.insert n i h, I.insert i n ns)

-- Node sharing
-- Find the ID of an existing node
lookupNode :: (Eq l, Hashable l, Monad m) => Node d l -> Diagram d l m (Maybe ID)
lookupNode n = do
  (_, _, h, _) <- get
  return $ H.lookup n h

-- Make a new proposition (if not already present) from a node
makeProp :: (Eq l, Hashable l, Monad m) => Node d l -> Diagram d l m (Prop d l)
makeProp n
  | hi n == Bot = return (lo n) -- Node elimination
  | otherwise = lookupNode n >>= \case
    Just i -> return (ID i)
    Nothing -> do
      i <- fresh
      insertNode i n
      return (ID i)

-- | Construct a proposition from an atomic proposition
atomic :: (Eq l, Hashable l, Monad m) => l -> Diagram d l m (Prop d l)
atomic l = makeProp Node {label = l, lo = Bot, hi = Top}

-- | Instantiating an atomic proposition
assign :: (Ord l, Hashable l, Monad m) => Prop d l -> l -> Bool -> Diagram d l m (Prop d l)
assign Top _ _ = return Top
assign Bot _ _ = return Bot
assign (ID i) var b =
  memo (Assign i var b) $
    withNode
      i
      ( \n ->
          case compare (label n) var of
            LT -> return Bot
            EQ -> return (hi n)
            GT -> do
              lo' <- assign (lo n) var b
              hi' <- assign (hi n) var b
              makeProp
                Node
                  { label = label n,
                    lo = lo',
                    hi = hi'
                  }
      )

-- | Negate every occurance of an atomic proposition
negate :: (Ord l, Hashable l, Monad m) => Prop d l -> l -> Diagram d l m (Prop d l)
negate Top _ = return Top
negate Bot _ = return Bot
negate p@(ID i) var =
  memo (Negate i var) $
    withNode
      i
      ( \n ->
          case compare (label n) var of
            LT ->
              makeProp
                Node
                  { label = var,
                    lo = Bot,
                    hi = p
                  }
            EQ ->
              makeProp
                Node
                  { label = var,
                    lo = hi n,
                    hi = lo n
                  }
            GT -> do
              lo' <- negate (lo n) var
              hi' <- negate (hi n) var
              makeProp
                Node
                  { label = label n,
                    lo = lo',
                    hi = hi'
                  }
      )

-- | Union of two propositions
union :: (Ord l, Hashable l, Monad m) => Prop d l -> Prop d l -> Diagram d l m (Prop d l)
union Top _ = return Top
union _ Top = return Top
union Bot q = return q
union p Bot = return p
union p q | p == q = return p
union p@(ID i) (ID j) =
  memo (Union i j) $
    withNode
      i
      ( \n ->
          withNode j $ \m ->
            case compare (label n) (label m) of
              LT -> do
                lo' <- p `union` lo m
                makeProp
                  Node
                    { label = label m,
                      lo = lo',
                      hi = hi m
                    }
              EQ -> do
                lo' <- lo n `union` lo m
                hi' <- hi n `union` hi m
                makeProp
                  Node
                    { label = label n,
                      lo = lo',
                      hi = hi'
                    }
              GT -> do
                lo' <- lo n `union` p
                makeProp
                  Node
                    { label = label n,
                      lo = lo',
                      hi = hi n
                    }
      )

-- | Intersection of two propositions
intersect :: (Ord l, Hashable l, Monad m) => Prop d l -> Prop d l -> Diagram d l m (Prop d l)
intersect Top p = return p
intersect p Top = return p
intersect Bot _ = return Bot
intersect _ Bot = return Bot
intersect p q | p == q = return p
intersect p@(ID i) q@(ID j) =
  memo (Intersect i j) $
    withNode
      i
      ( \n ->
          withNode j $ \m ->
            case compare (label n) (label m) of
              LT -> lo n `intersect` q
              EQ -> do
                lo' <- lo n `intersect` lo m
                hi' <- hi n `intersect` hi m
                makeProp
                  Node
                    { label = label n,
                      lo = lo',
                      hi = hi'
                    }
              GT -> p `intersect` lo m
      )

-- | Difference between two propositions
difference :: (Ord l, Invertible l, Hashable l, Monad m) => Prop d l -> Prop d l -> Diagram d l m (Prop d l)
difference _ Top = return Bot
difference Bot _ = return Bot
difference p Bot = return p
difference p q | p == q = return p
difference Top (ID j) =
  memo (DifferenceT j) $
    withNode
      j
      ( \m ->
          makeProp
            Node
              { label = neg $ label m,
                lo = lo m,
                hi = hi m
              }
      )
difference p@(ID i) q@(ID j) =
  memo (Difference i j) $
    withNode
      i
      ( \n ->
          withNode j $ \m ->
            case compare (label n) (label m) of
              LT -> p `difference` lo m
              EQ -> do
                lo' <- lo n `difference` lo m
                hi' <- hi n `difference` hi m
                makeProp
                  Node
                    { label = label n,
                      lo = lo',
                      hi = hi'
                    }
              GT -> do
                lo' <- lo n `difference` q
                makeProp
                  Node
                    { label = label n,
                      lo = lo',
                      hi = hi n
                    }
      )

-- | Does the proposition have any models?
isEmpty, isNotEmpty :: Monad m => Prop d l -> Diagram d l m Bool
isEmpty Bot = return True
isEmpty Top = return False
isEmpty (ID i) = withNode i $ \n -> do
  p <- isEmpty (lo n)
  q <- isEmpty (hi n)
  return (p && q)
isNotEmpty Bot = return False
isNotEmpty Top = return True
isNotEmpty (ID i) = withNode i $ \n -> do
  p <- isNotEmpty (lo n)
  q <- isNotEmpty (hi n)
  return (p || q)

-- | Count the number of models
count :: Monad m => Prop d l -> Diagram d l m Int
count Bot = return 0
count Top = return 1
count (ID i) = withNode i $ \n -> do
  x <- count (lo n)
  y <- count (hi n)
  return (x + y)

-- | Map between atomic propositions
mapAtomic :: (Eq l, Hashable l, Monad m) => (l -> l) -> Prop d l -> Diagram d l m (Prop d l)
mapAtomic _ Top = return Top
mapAtomic _ Bot = return Bot
mapAtomic f (ID i) = withNode i $ \n -> do
  lo' <- mapAtomic f (lo n)
  hi' <- mapAtomic f (hi n)
  makeProp
    Node
      { label = f (label n),
        lo = lo',
        hi = hi'
      }

-- Loop continuations to another part of the diagram
mapTerminal :: (Eq l, Hashable l, Monad m) => Prop d l -> Prop d l -> Prop d l -> Diagram d l m (Prop d l)
mapTerminal top _ Top = return top
mapTerminal _ bot Bot = return bot
mapTerminal top bot (ID i) = withNode i $ \n -> do
  lo' <- mapTerminal top bot (lo n)
  hi' <- mapTerminal top bot (hi n)
  makeProp
    Node
      { label = label n,
        lo = lo',
        hi = hi'
      }

-- | Instantiate an atomic proposition with a non-atomic proposition
atomicToProp :: (Eq l, Hashable l, Monad m) => (l -> Diagram d l m (Prop d l)) -> Prop d l -> Diagram d l m (Prop d l)
atomicToProp _ Top = return Top
atomicToProp _ Bot = return Bot
atomicToProp f (ID i) = withNode i $ \n -> do
  lo' <- atomicToProp f (lo n)
  hi' <- atomicToProp f (hi n)
  p <- f (label n)
  mapTerminal lo' hi' p