From 0656a2faf4321cc8a83624273df7bf03beb98608 Mon Sep 17 00:00:00 2001
From: Joscha <joscha@plugh.de>
Date: Thu, 31 Oct 2019 23:50:54 +0000
Subject: [PATCH] Visualize DFAs over a finite set of tokens

---
 src/Rextra/Dfa.hs       | 16 ++++++++--------
 src/Rextra/Visualize.hs | 21 ++++++++++++++++++++-
 2 files changed, 28 insertions(+), 9 deletions(-)

diff --git a/src/Rextra/Dfa.hs b/src/Rextra/Dfa.hs
index 903e9b3..575bf90 100644
--- a/src/Rextra/Dfa.hs
+++ b/src/Rextra/Dfa.hs
@@ -6,6 +6,7 @@ module Rextra.Dfa
   ( Dfa
   , dfa
   , State(..)
+  , stateTransition
   , normalize
   , mapState
   , transitionsByState
@@ -25,6 +26,12 @@ data State s t = State
   , defaultTransition :: s
   } deriving (Show, Eq, Ord)
 
+stateTransition :: (Ord t) => State s t -> t -> s
+stateTransition state token =
+  case transitions state Map.!? token of
+    Nothing -> defaultTransition state
+    Just s  -> s
+
 normalize :: (Eq s) => State s t -> State s t
 normalize State{transitions, defaultTransition} =
   State { transitions = Map.filter (/= defaultTransition) transitions
@@ -50,16 +57,9 @@ type Dfa s t = Fa State s t
 
 instance (Ord s) => Executable (Fa State s) s where
   startState = entryState
-  transition = dfaTransition
+  transition a s = stateTransition (getState a s)
   accepts a s = s `Set.member` exitStates a
 
-dfaTransition :: (Ord s, Ord t) => Dfa s t -> s -> t -> s
-dfaTransition a s t =
-  let state = getState a s
-  in  case transitions state Map.!? t of
-        (Just nextState) -> nextState
-        Nothing          -> defaultTransition state
-
 dfa :: (Ord s, Ord t) => [(s, [(t, s)], s)] -> s -> [s] -> Maybe (Dfa s t)
 dfa stateInfo entryState exitStates =
   let stateList = map (\(s, ts, dt) -> (s, State (Map.fromList ts) dt)) stateInfo
diff --git a/src/Rextra/Visualize.hs b/src/Rextra/Visualize.hs
index f921596..0142850 100644
--- a/src/Rextra/Visualize.hs
+++ b/src/Rextra/Visualize.hs
@@ -1,16 +1,26 @@
 {-# LANGUAGE FlexibleInstances #-}
 
-module Rextra.Visualize where
+module Rextra.Visualize
+  ( showDot
+  , saveDot
+  , saveDotAsPng
+  , dfaToDot
+  , dfaToDotNoDefault
+  , dfaToDotWithTokens
+  , nfaToDot
+  ) where
 
 import           Control.Monad
 import           Data.GraphViz
 import           Data.List
 import qualified Data.Map as Map
 import qualified Data.Set as Set
+import           Data.Tuple
 
 import qualified Rextra.Dfa as Dfa
 import           Rextra.Fa
 import qualified Rextra.Nfa as Nfa
+import           Rextra.Util
 
 -- First, some labelling...
 
@@ -105,6 +115,15 @@ dfaToDot = faToDot dfaShowAll
 dfaToDotNoDefault :: (Ord s, AsLabel s, Ord t, AsLabel t) => Dfa.Dfa s t -> DotGraph String
 dfaToDotNoDefault = faToDot dfaShowNoDefault
 
+dfaShowTokens :: (Ord s, Ord t, AsLabel t) => Set.Set t -> Dfa.State s t -> [(String, s)]
+dfaShowTokens tokens state =
+  let bareEdges    = Map.assocs $ Map.fromSet (Dfa.stateTransition state) tokens
+      groupedEdges = Map.assocs $ groupByFirst $ map swap bareEdges
+  in  map (\(s, t) -> (asLabel t, s)) groupedEdges
+
+dfaToDotWithTokens :: (Ord s, AsLabel s, Ord t, AsLabel t) => [t] -> Dfa.Dfa s t -> DotGraph String
+dfaToDotWithTokens tokenList = faToDot $ dfaShowTokens $ Set.fromList tokenList
+
 {-
  - 'Nfa' stuff
  -}