Merge ada6728 into 2cb4a1b

msakai · Nov 24, 2021 · fce2fad · fce2fad
2 parents 2cb4a1b + ada6728
commit fce2fad
Show file tree

Hide file tree

Showing 3 changed files with 82 additions and 92 deletions.
diff --git a/src/Data/DecisionDiagram/BDD.hs b/src/Data/DecisionDiagram/BDD.hs
@@ -125,7 +125,6 @@ import Control.Monad.ST.Unsafe
 import Data.Bits (Bits (shiftL))
 import qualified Data.Foldable as Foldable
 import Data.Function (on)
-import Data.Functor.Identity
 import Data.Hashable
 import qualified Data.HashMap.Lazy as HashMap
 import qualified Data.HashTable.Class as H
@@ -139,7 +138,6 @@ import Data.Map.Lazy (Map)
 import qualified Data.Map.Lazy as Map
 import Data.Proxy
 import Data.Ratio
-import Data.STRef
 import qualified Data.Vector as V
 import Numeric.Natural
 #if MIN_VERSION_mwc_random(0,15,0)
@@ -152,7 +150,7 @@ import System.Random.MWC.Distributions (bernoulli)
 import Text.Read
 
 import Data.DecisionDiagram.BDD.Internal.ItemOrder
-import Data.DecisionDiagram.BDD.Internal.Node (Sig (..))
+import Data.DecisionDiagram.BDD.Internal.Node (Sig (..), Graph)
 import qualified Data.DecisionDiagram.BDD.Internal.Node as Node
 
 infixr 3 .&&.
@@ -1050,63 +1048,24 @@ outSig (Branch x lo hi) = SBranch x lo hi
 
 -- ------------------------------------------------------------------------
 
-type Graph f = IntMap (f Int)
-
 -- | Convert a BDD into a pointed graph
 --
 -- Nodes @0@ and @1@ are reserved for @SLeaf False@ and @SLeaf True@
 -- even if they are not actually used. Therefore the result may be
 -- larger than 'numNodes' if the leaf nodes are not used.
 toGraph :: BDD a -> (Graph Sig, Int)
-toGraph bdd =
-  case toGraph' (Identity bdd) of
-    (g, Identity v) -> (g, v)
+toGraph (BDD node) = Node.toGraph node
 
 -- | Convert multiple BDDs into a graph
 toGraph' :: Traversable t => t (BDD a) -> (Graph Sig, t Int)
-toGraph' bs = runST $ do
-  h <- C.newSized defaultTableSize
-  H.insert h F 0
-  H.insert h T 1
-  counter <- newSTRef 2
-  ref <- newSTRef $ IntMap.fromList [(0, SLeaf False), (1, SLeaf True)]
-
-  let f F = return 0
-      f T = return 1
-      f p@(Branch x lo hi) = do
-        m <- H.lookup h p
-        case m of
-          Just ret -> return ret
-          Nothing -> do
-            r0 <- f lo
-            r1 <- f hi
-            n <- readSTRef counter
-            writeSTRef counter $! n+1
-            H.insert h p n
-            modifySTRef' ref (IntMap.insert n (SBranch x r0 r1))
-            return n
-
-  vs <- mapM f bs
-  g <- readSTRef ref
-  return (g, vs)
+toGraph' bs = Node.toGraph' (fmap (\(BDD node) -> node) bs)
 
 -- | Convert a pointed graph into a BDD
 fromGraph :: (Graph Sig, Int) -> BDD a
-fromGraph (g, v) =
-  case IntMap.lookup v (fromGraph' g) of
-    Nothing -> error ("Data.DecisionDiagram.BDD.fromGraph: invalid node id " ++ show v)
-    Just bdd -> bdd
+fromGraph = Node.foldGraph inSig
 
 -- | Convert nodes of a graph into BDDs
 fromGraph' :: Graph Sig -> IntMap (BDD a)
-fromGraph' g = ret
-  where
-    ret = IntMap.map f g
-    f (SLeaf b) = Leaf b
-    f (SBranch x lo hi) =
-      case (IntMap.lookup lo ret, IntMap.lookup hi ret) of
-        (Nothing, _) -> error ("Data.DecisionDiagram.BDD.fromGraph': invalid node id " ++ show lo)
-        (_, Nothing) -> error ("Data.DecisionDiagram.BDD.fromGraph': invalid node id " ++ show hi)
-        (Just lo', Just hi') -> Branch x lo' hi'
+fromGraph' = Node.foldGraphNodes inSig
 
 -- ------------------------------------------------------------------------
diff --git a/src/Data/DecisionDiagram/BDD/Internal/Node.hs b/src/Data/DecisionDiagram/BDD/Internal/Node.hs
@@ -35,15 +35,26 @@ module Data.DecisionDiagram.BDD.Internal.Node
 
   -- * (Co)algebraic structure
   , Sig (..)
+
+  -- * Graph
+  , Graph
+  , toGraph
+  , toGraph'
+  , foldGraph
+  , foldGraphNodes
   ) where
 
 import Control.Monad
 import Control.Monad.ST
 import Control.Monad.ST.Unsafe
+import Data.Functor.Identity
 import Data.Hashable
 import qualified Data.HashTable.Class as H
 import qualified Data.HashTable.ST.Cuckoo as C
 import Data.Interned
+import Data.IntMap.Lazy (IntMap)
+import qualified Data.IntMap.Lazy as IntMap
+import Data.STRef
 import GHC.Generics
 
 -- ------------------------------------------------------------------------
@@ -192,3 +203,64 @@ mkFold'Op br lf = do
 
 -- ------------------------------------------------------------------------
 
+-- | Graph where nodes are decorated using a functor @f@.
+--
+-- The occurrences of the parameter of @f@ represent out-going edges.
+type Graph f = IntMap (f Int)
+
+-- | Convert a node into a pointed graph
+--
+-- Nodes @0@ and @1@ are reserved for @SLeaf False@ and @SLeaf True@ even if
+-- they are not actually used. Therefore the result may be larger than
+-- 'numNodes' if the leaf nodes are not used.
+toGraph :: Node -> (Graph Sig, Int)
+toGraph bdd =
+  case toGraph' (Identity bdd) of
+    (g, Identity v) -> (g, v)
+
+-- | Convert multiple nodes into a graph
+toGraph' :: Traversable t => t Node -> (Graph Sig, t Int)
+toGraph' bs = runST $ do
+  h <- C.newSized defaultTableSize
+  H.insert h (Leaf False) 0
+  H.insert h (Leaf True) 1
+  counter <- newSTRef 2
+  ref <- newSTRef $ IntMap.fromList [(0, SLeaf False), (1, SLeaf True)]
+
+  let f (Leaf False) = return 0
+      f (Leaf True) = return 1
+      f p@(Branch x lo hi) = do
+        m <- H.lookup h p
+        case m of
+          Just ret -> return ret
+          Nothing -> do
+            r0 <- f lo
+            r1 <- f hi
+            n <- readSTRef counter
+            writeSTRef counter $! n+1
+            H.insert h p n
+            modifySTRef' ref (IntMap.insert n (SBranch x r0 r1))
+            return n
+
+  vs <- mapM f bs
+  g <- readSTRef ref
+  return (g, vs)
+
+-- | Fold over pointed graph
+foldGraph :: Functor f => (f a -> a) -> (Graph f, Int) -> a
+foldGraph f (g, v) =
+  case IntMap.lookup v (foldGraphNodes f g) of
+    Just x -> x
+    Nothing -> error ("foldGraphNodes: invalid node id " ++ show v)
+
+-- | Fold over graph nodes
+foldGraphNodes :: Functor f => (f a -> a) -> Graph f -> IntMap a
+foldGraphNodes f m = ret
+  where
+    ret = IntMap.map (f . fmap h) m
+    h v =
+      case IntMap.lookup v ret of
+        Just x -> x
+        Nothing -> error ("foldGraphNodes: invalid node id " ++ show v)
+
+-- ------------------------------------------------------------------------
diff --git a/src/Data/DecisionDiagram/ZDD.hs b/src/Data/DecisionDiagram/ZDD.hs
@@ -143,7 +143,6 @@ import Control.Monad.Primitive
 import Control.Monad.ST
 import qualified Data.Foldable as Foldable
 import Data.Function (on)
-import Data.Functor.Identity
 import Data.Hashable
 import Data.HashMap.Lazy (HashMap)
 import qualified Data.HashMap.Lazy as HashMap
@@ -161,7 +160,6 @@ import Data.Proxy
 import Data.Ratio
 import Data.Set (Set)
 import qualified Data.Set as Set
-import Data.STRef
 import qualified Data.Vector as V
 import qualified GHC.Exts as Exts
 import Numeric.Natural
@@ -174,7 +172,7 @@ import System.Random.MWC.Distributions (bernoulli)
 import Text.Read
 
 import Data.DecisionDiagram.BDD.Internal.ItemOrder
-import Data.DecisionDiagram.BDD.Internal.Node (Sig (..))
+import Data.DecisionDiagram.BDD.Internal.Node (Sig (..), Graph)
 import qualified Data.DecisionDiagram.BDD.Internal.Node as Node
 import qualified Data.DecisionDiagram.BDD as BDD
 
@@ -1002,63 +1000,24 @@ outSig (Branch x lo hi) = SBranch x lo hi
 
 -- ------------------------------------------------------------------------
 
-type Graph f = IntMap (f Int)
-
 -- | Convert a ZDD into a pointed graph
 --
 -- Nodes @0@ and @1@ are reserved for @SLeaf False@ and @SLeaf True@ even if
 -- they are not actually used. Therefore the result may be larger than
 -- 'numNodes' if the leaf nodes are not used.
 toGraph :: ZDD a -> (Graph Sig, Int)
-toGraph bdd =
-  case toGraph' (Identity bdd) of
-    (g, Identity v) -> (g, v)
+toGraph (ZDD node) = Node.toGraph node
 
 -- | Convert multiple ZDDs into a graph
 toGraph' :: Traversable t => t (ZDD a) -> (Graph Sig, t Int)
-toGraph' bs = runST $ do
-  h <- C.newSized defaultTableSize
-  H.insert h Empty 0
-  H.insert h Base 1
-  counter <- newSTRef 2
-  ref <- newSTRef $ IntMap.fromList [(0, SLeaf False), (1, SLeaf True)]
-
-  let f Empty = return 0
-      f Base = return 1
-      f p@(Branch x lo hi) = do
-        m <- H.lookup h p
-        case m of
-          Just ret -> return ret
-          Nothing -> do
-            r0 <- f lo
-            r1 <- f hi
-            n <- readSTRef counter
-            writeSTRef counter $! n+1
-            H.insert h p n
-            modifySTRef' ref (IntMap.insert n (SBranch x r0 r1))
-            return n
-
-  vs <- mapM f bs
-  g <- readSTRef ref
-  return (g, vs)
+toGraph' bs = Node.toGraph' (fmap (\(ZDD node) -> node) bs)
 
 -- | Convert a pointed graph into a ZDD
 fromGraph :: (Graph Sig, Int) -> ZDD a
-fromGraph (g, v) =
-  case IntMap.lookup v (fromGraph' g) of
-    Nothing -> error ("Data.DecisionDiagram.ZDD.fromGraph: invalid node id " ++ show v)
-    Just bdd -> bdd
+fromGraph = Node.foldGraph inSig
 
 -- | Convert nodes of a graph into ZDDs
 fromGraph' :: Graph Sig -> IntMap (ZDD a)
-fromGraph' g = ret
-  where
-    ret = IntMap.map f g
-    f (SLeaf b) = Leaf b
-    f (SBranch x lo hi) =
-      case (IntMap.lookup lo ret, IntMap.lookup hi ret) of
-        (Nothing, _) -> error ("Data.DecisionDiagram.ZDD.fromGraph': invalid node id " ++ show lo)
-        (_, Nothing) -> error ("Data.DecisionDiagram.ZDD.fromGraph': invalid node id " ++ show hi)
-        (Just lo', Just hi') -> Branch x lo' hi'
+fromGraph' = Node.foldGraphNodes inSig
 
 -- ------------------------------------------------------------------------