Render CFG using new mermaid code

Author: kddnewton

lib/syntax_tree.rb

@@ -23,6 +23,7 @@
 require_relative "syntax_tree/visitor/environment"
 require_relative "syntax_tree/visitor/with_environment"
 
+require_relative "syntax_tree/mermaid"
 require_relative "syntax_tree/parser"
 require_relative "syntax_tree/pattern"
 require_relative "syntax_tree/search"

lib/syntax_tree/mermaid.rb

@@ -6,7 +6,7 @@ module SyntaxTree
   # This module is responsible for rendering mermaid flow charts.
   module Mermaid
     class Node
-      SHAPES = %i[circle rectangle stadium].freeze
+      SHAPES = %i[circle rectangle rounded stadium].freeze
 
       attr_reader :id, :label, :shape
 
@@ -19,17 +19,23 @@ def initialize(id, label, shape)
       end
 
       def render
-        left_bound, right_bound =
-          case shape
-          when :circle
-            ["((", "))"]
-          when :rectangle
-            ["[", "]"]
-          when :stadium
-            ["([", "])"]
-          end
+        left_bound, right_bound = bounds
+        "#{id}#{left_bound}\"#{CGI.escapeHTML(label)}\"#{right_bound}"
+      end
 
-        "  #{id}#{left_bound}\"#{CGI.escapeHTML(label)}\"#{right_bound}"
+      private
+
+      def bounds
+        case shape
+        when :circle
+          ["((", "))"]
+        when :rectangle
+          ["[", "]"]
+        when :rounded
+          ["(", ")"]
+        when :stadium
+          ["([", "])"]
+        end
       end
     end
 
@@ -50,34 +56,57 @@ def initialize(from, to, label, type)
       def render
         case type
         when :directed
-          "  #{from.id} -- \"#{CGI.escapeHTML(label)}\" --> #{to.id}"
+          if label
+            "#{from.id} -- \"#{CGI.escapeHTML(label)}\" --> #{to.id}"
+          else
+            "#{from.id} --> #{to.id}"
+          end
         end
       end
     end
 
     class FlowChart
-      attr_reader :nodes, :edges
+      attr_reader :output, :prefix, :nodes
 
       def initialize
+        @output = StringIO.new
+        @output.puts("flowchart TD")
+        @prefix = "  "
         @nodes = {}
-        @edges = []
       end
 
-      def edge(from, to, label, type = :directed)
-        edges << Edge.new(from, to, label, type)
+      def edge(from, to, label = nil, type: :directed)
+        edge = Edge.new(from, to, label, type)
+        output.puts("#{prefix}#{edge.render}")
       end
 
-      def node(id, label, shape = :rectangle)
-        nodes[id] = Node.new(id, label, shape)
+      def fetch(id)
+        nodes.fetch(id)
       end
 
-      def render
-        output = StringIO.new
-        output.puts("flowchart TD")
+      def node(id, label, shape: :rectangle)
+        node = Node.new(id, label, shape)
+        nodes[id] = node
+
+        output.puts("#{prefix}#{nodes[id].render}")
+        node
+      end
+
+      def subgraph(id)
+        output.puts("#{prefix}subgraph #{id}")
+
+        previous = prefix
+        @prefix = "#{prefix}  "
 
-        nodes.each_value { |node| output.puts(node.render) }
-        edges.each { |edge| output.puts(edge.render) }
+        begin
+          yield
+        ensure
+          @prefix = previous
+          output.puts("#{prefix}end")
+        end
+      end
 
+      def render
         output.string
       end
     end

lib/syntax_tree/visitor/mermaid_visitor.rb

@@ -29,7 +29,7 @@ def field(name, value)
         when Node
           flowchart.edge(target, visit(value), name)
         else
-          to = flowchart.node("#{target.id}_#{name}", value.inspect, :stadium)
+          to = flowchart.node("#{target.id}_#{name}", value.inspect, shape: :stadium)
           flowchart.edge(target, to, name)
         end
       end
@@ -54,7 +54,7 @@ def node(node, type)
 
       def pairs(name, values)
         values.each_with_index do |(key, value), index|
-          to = flowchart.node("#{target.id}_#{name}_#{index}", " ", :circle)
+          to = flowchart.node("#{target.id}_#{name}_#{index}", " ", shape: :circle)
 
           flowchart.edge(target, to, "#{name}[#{index}]")
           flowchart.edge(to, visit(key), "[0]")

lib/syntax_tree/yarv/control_flow_graph.rb

@@ -208,25 +208,24 @@ def to_son
       end
 
       def to_mermaid
-        output = StringIO.new
-        output.puts("flowchart TD")
+        flowchart = Mermaid::FlowChart.new
+        disasm = Disassembler::Mermaid.new
 
-        fmt = Disassembler::Mermaid.new
         blocks.each do |block|
-          output.puts("  subgraph #{block.id}")
-          previous = nil
-
-          block.each_with_length do |insn, length|
-            node_id = "node_#{length}"
-            label = "%04d %s" % [length, insn.disasm(fmt)]
-
-            output.puts("    #{node_id}(\"#{CGI.escapeHTML(label)}\")")
-            output.puts("    #{previous} --> #{node_id}") if previous
-
-            previous = node_id
+          flowchart.subgraph(block.id) do
+            previous = nil
+
+            block.each_with_length do |insn, length|
+              node =
+                flowchart.node(
+                  "node_#{length}",
+                  "%04d %s" % [length, insn.disasm(disasm)]
+                )
+
+              flowchart.edge(previous, node) if previous
+              previous = node
+            end
           end
-
-          output.puts("  end")
         end
 
         blocks.each do |block|
@@ -235,11 +234,13 @@ def to_mermaid
               block.block_start + block.insns.sum(&:length) -
                 block.insns.last.length
 
-            output.puts("  node_#{offset} --> node_#{outgoing.block_start}")
+            from = flowchart.fetch("node_#{offset}")
+            to = flowchart.fetch("node_#{outgoing.block_start}")
+            flowchart.edge(from, to)
           end
         end
 
-        output.string
+        flowchart.render
       end
 
       # This method is used to verify that the control flow graph is well