1 files changed, 52 insertions, 55 deletions

diff --git a/chapter02/src/main/java/com/seaofnodes/simple/GraphVisualizer.java b/chapter02/src/main/java/com/seaofnodes/simple/GraphVisualizer.java
index d2aa8179..1af19d80 100644
--- a/chapter02/src/main/java/com/seaofnodes/simple/GraphVisualizer.java
+++ b/chapter02/src/main/java/com/seaofnodes/simple/GraphVisualizer.java
@@ -5,9 +5,7 @@ import com.seaofnodes.simple.node.Control;
 import com.seaofnodes.simple.node.Node;
 import com.seaofnodes.simple.node.StartNode;
 
-import java.util.Collection;
-import java.util.HashMap;
-import java.util.Map;
+import java.util.*;
 
 /**
  * Simple visualizer that outputs GraphViz dot format.
@@ -16,59 +14,59 @@ import java.util.Map;
  */
 public class GraphVisualizer {
 
-    public String generateDotOutput(StartNode start) {
+    public String generateDotOutput(Parser parser) {
+
         // Since the graph has cycles, we need to create a flat list of all the
         // nodes in the graph.
-        Collection<Node> all = findAll(start);
+        Collection<Node> all = findAll(parser);
         StringBuilder sb = new StringBuilder();
-        sb.append("digraph chapter02\n" +
-                "{\n");
+        sb.append("digraph chapter02 {\n");
+
+        // To keep the Scopes below the graph and pointing up into the graph we
+        // need to group the Nodes in a subgraph cluster, and the scopes into a
+        // different subgraph cluster.  THEN we can draw edges between the
+        // scopes and nodes.  If we try to cross subgraph cluster borders while
+        // still making the subgraphs DOT gets confused.
+        sb.append("\trankdir=BT;\n"); // Force Nodes before Scopes
+
+        // Just the Nodes first, in a cluster no edges
+        nodes(sb, all);
+
+        // Walk the Node edges
+        nodeEdges(sb, all);
+
+        sb.append("}\n");
+        return sb.toString();
+    }
+
+
+    private void nodes(StringBuilder sb, Collection<Node> all) {
+        // Just the Nodes first, in a cluster no edges
+        sb.append("\tsubgraph cluster_Nodes {\n"); // Magic "cluster_" in the subgraph name
         for (Node n : all) {
-            sb.append('\t').append(n.uniqueName());
+            sb.append("\t\t").append(n.uniqueName()).append(" [ ");
             // control nodes have box shape
             // other nodes are ellipses, i.e. default shape
             if (n instanceof Control)
-                sb.append(" [shape=box, ");
-            else
-                sb.append(" [");
-            sb.append(" label=\"")
-                    .append(n.label())
-                    .append("\"];\n");
-            // Output the non control edges
-            walkIns(sb, n, false);
+                sb.append("shape=box style=filled fillcolor=yellow ");
+            sb.append("label=\"").append(n.label()).append("\" ");
+            sb.append("];\n");
         }
-        // Now output the control edges which must be in red
-        sb.append("\tedge [color=red];\n");
-        for (Node n : all) {
-            walkIns(sb, n, true);
-        }
-        sb.append("}\n");
-        return sb.toString();
+        sb.append("\t}\n");     // End Node cluster
     }
 
-    /**
-     * Outputs edges. If doControlEdges is true then
-     * edges between control nodes are output. Else other edges
-     * output.
-     */
-    private void walkIns(StringBuilder sb, Node in, boolean doControlEdges) {
-        for (int i = 0; i < in.nIns(); i++) {
-            Node out = in.in(i);
-            if (out == null)
-                continue;
-            boolean isControlEdge = (out instanceof Control) && (in instanceof Control);
-            if (doControlEdges && !isControlEdge)
-                continue;
-            if (!doControlEdges && isControlEdge)
-                continue;
-            sb.append('\t')
-                    .append(out.uniqueName())
-                    .append(" -> ")
-                    .append(in.uniqueName());
-            // the edge from start node to constants is just for convenience so
-            // show it differently
-            if ((out instanceof StartNode) && (in instanceof ConstantNode))
+
+    // Walk the node edges
+    private void nodeEdges(StringBuilder sb, Collection<Node> all) {
+        for( Node n : all )
+            for( Node out : n._inputs )
+                if( out != null ) {
+                    sb.append('\t').append(n.uniqueName()).append(" -> ").append(out.uniqueName());
+                    // Control edges are colored red
+                    if( n instanceof ConstantNode && out instanceof StartNode )
                 sb.append(" [style=dotted]");
+                    else if( n instanceof Control && out instanceof Control )
+                      sb.append(" [color=red]");
             sb.append(";\n");
         }
     }
@@ -77,26 +75,25 @@ public class GraphVisualizer {
      * Walks the whole graph, starting from Start.
      * Since Start is the input to all constants - we look at the outputs for
      * Start, but for then subsequently we look at the inputs of each node.
+     * During graph construction not all nodes are reachable this way, so we
+     * also scan the symbol tables.
      */
-    private Collection<Node> findAll(StartNode start) {
-        Map<Integer, Node> all = new HashMap<>();
-        for (int i = 0; i < start.nOuts(); i++) {
-            Node n = start.out(i);
+    private Collection<Node> findAll(Parser parser) {
+        final StartNode start = Parser.START;
+        final HashMap<Integer, Node> all = new HashMap<>();
+        for( Node n : start._outputs )
             walk(all, n);
-        }
         return all.values();
     }
 
     /**
-     * Walk a subgraph and populate distinct nodes in the _all list.
+     * Walk a subgraph and populate distinct nodes in the all list.
      */
-    private void walk(Map<Integer, Node> all, Node n) {
-        if (all.get(n._nid) == null) {
-            // Not yet seen
+    private void walk(HashMap<Integer, Node> all, Node n) {
+        if (all.get(n._nid) != null) return; // Been there, done that
             all.put(n._nid, n);
             for (Node c : n._inputs)
                 if (c != null)
                     walk(all, c);
-        }
     }
 }