Render sea of nodes to mermaid using new API

Author: kddnewton

.rubocop.yml

@@ -117,6 +117,9 @@ Style/FormatStringToken:
 Style/GuardClause:
   Enabled: false
 
+Style/HashLikeCase:
+  Enabled: false
+
 Style/IdenticalConditionalBranches:
   Enabled: false
 

lib/syntax_tree/mermaid.rb

@@ -3,20 +3,85 @@
 require "cgi"
 
 module SyntaxTree
-  # This module is responsible for rendering mermaid flow charts.
+  # This module is responsible for rendering mermaid (https://mermaid.js.org/)
+  # flow charts.
   module Mermaid
-    def self.escape(label)
-      "\"#{CGI.escapeHTML(label)}\""
+    # This is the main class that handles rendering a flowchart. It keeps track
+    # of its nodes and links and renders them according to the mermaid syntax.
+    class FlowChart
+      attr_reader :output, :prefix, :nodes, :links
+
+      def initialize
+        @output = StringIO.new
+        @output.puts("flowchart TD")
+        @prefix = "  "
+
+        @nodes = {}
+        @links = []
+      end
+
+      # Retrieve a node that has already been added to the flowchart by its id.
+      def fetch(id)
+        nodes.fetch(id)
+      end
+
+      # Add a link to the flowchart between two nodes with an optional label.
+      def link(from, to, label = nil, type: :directed, color: nil)
+        link = Link.new(from, to, label, type, color)
+        links << link
+
+        output.puts("#{prefix}#{link.render}")
+        link
+      end
+
+      # Add a node to the flowchart with an optional label.
+      def node(id, label = " ", shape: :rectangle)
+        node = Node.new(id, label, shape)
+        nodes[id] = node
+
+        output.puts("#{prefix}#{nodes[id].render}")
+        node
+      end
+
+      # Add a subgraph to the flowchart. Within the given block, all of the
+      # nodes will be rendered within the subgraph.
+      def subgraph(label)
+        output.puts("#{prefix}subgraph #{Mermaid.escape(label)}")
+
+        previous = prefix
+        @prefix = "#{prefix}  "
+
+        begin
+          yield
+        ensure
+          @prefix = previous
+          output.puts("#{prefix}end")
+        end
+      end
+
+      # Return the rendered flowchart.
+      def render
+        links.each_with_index do |link, index|
+          if link.color
+            output.puts("#{prefix}linkStyle #{index} stroke:#{link.color}")
+          end
+        end
+
+        output.string
+      end
     end
 
+    # This class represents a link between two nodes in a flowchart. It is not
+    # meant to be interacted with directly, but rather used as a data structure
+    # by the FlowChart class.
     class Link
-      TYPES = %i[directed].freeze
+      TYPES = %i[directed dotted].freeze
       COLORS = %i[green red].freeze
 
       attr_reader :from, :to, :label, :type, :color
 
       def initialize(from, to, label, type, color)
-        raise if !TYPES.include?(type)
+        raise unless TYPES.include?(type)
         raise if color && !COLORS.include?(color)
 
         @from = from
@@ -27,17 +92,31 @@ def initialize(from, to, label, type, color)
       end
 
       def render
+        left_side, right_side, full_side = sides
+
+        if label
+          escaped = Mermaid.escape(label)
+          "#{from.id} #{left_side} #{escaped} #{right_side} #{to.id}"
+        else
+          "#{from.id} #{full_side} #{to.id}"
+        end
+      end
+
+      private
+
+      def sides
         case type
         when :directed
-          if label
-            "#{from.id} -- #{Mermaid.escape(label)} --> #{to.id}"
-          else
-            "#{from.id} --> #{to.id}"
-          end
+          %w[-- --> -->]
+        when :dotted
+          %w[-. .-> -.->]
         end
       end
     end
 
+    # This class represents a node in a flowchart. Unlike the Link class, it can
+    # be used directly. It is the return value of the #node method, and is meant
+    # to be passed around to #link methods to create links between nodes.
     class Node
       SHAPES = %i[circle rectangle rounded stadium].freeze
 
@@ -61,72 +140,37 @@ def render
       def bounds
         case shape
         when :circle
-          ["((", "))"]
+          %w[(( ))]
         when :rectangle
           ["[", "]"]
         when :rounded
-          ["(", ")"]
+          %w[( )]
         when :stadium
           ["([", "])"]
         end
       end
     end
 
-    class FlowChart
-      attr_reader :output, :prefix, :nodes, :links
-
-      def initialize
-        @output = StringIO.new
-        @output.puts("flowchart TD")
-        @prefix = "  "
-
-        @nodes = {}
-        @links = []
-      end
-
-      def fetch(id)
-        nodes.fetch(id)
-      end
-
-      def link(from, to, label = nil, type: :directed, color: nil)
-        link = Link.new(from, to, label, type, color)
-        links << link
-
-        output.puts("#{prefix}#{link.render}")
-        link
+    class << self
+      # Escape a label to be used in the mermaid syntax. This is used to escape
+      # HTML entities such that they render properly within the quotes.
+      def escape(label)
+        "\"#{CGI.escapeHTML(label)}\""
       end
 
-      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(label)
-        output.puts("#{prefix}subgraph #{Mermaid.escape(label)}")
-
-        previous = prefix
-        @prefix = "#{prefix}  "
-
-        begin
-          yield
-        ensure
-          @prefix = previous
-          output.puts("#{prefix}end")
+      # Create a new flowchart. If a block is given, it will be yielded to and
+      # the flowchart will be rendered. Otherwise, the flowchart will be
+      # returned.
+      def flowchart
+        flowchart = FlowChart.new
+
+        if block_given?
+          yield flowchart
+          flowchart.render
+        else
+          flowchart
         end
       end
-
-      def render
-        links.each_with_index do |link, index|
-          if link.color
-            output.puts("#{prefix}linkStyle #{index} stroke:#{link.color}")
-          end
-        end
-
-        output.string
-      end
     end
   end
 end

lib/syntax_tree/visitor/mermaid_visitor.rb

@@ -7,7 +7,7 @@ class MermaidVisitor < FieldVisitor
       attr_reader :flowchart, :target
 
       def initialize
-        @flowchart = Mermaid::FlowChart.new
+        @flowchart = Mermaid.flowchart
         @target = nil
       end
 
@@ -29,7 +29,12 @@ def field(name, value)
         when Node
           flowchart.link(target, visit(value), name)
         else
-          to = flowchart.node("#{target.id}_#{name}", value.inspect, shape: :stadium)
+          to =
+            flowchart.node(
+              "#{target.id}_#{name}",
+              value.inspect,
+              shape: :stadium
+            )
           flowchart.link(target, to, name)
         end
       end
@@ -54,7 +59,7 @@ def node(node, type)
 
       def pairs(name, values)
         values.each_with_index do |(key, value), index|
-          to = flowchart.node("#{target.id}_#{name}_#{index}", " ", shape: :circle)
+          to = flowchart.node("#{target.id}_#{name}_#{index}", shape: :circle)
 
           flowchart.link(target, to, "#{name}[#{index}]")
           flowchart.link(to, visit(key), "[0]")

lib/syntax_tree/yarv/control_flow_graph.rb

@@ -208,39 +208,38 @@ def to_son
       end
 
       def to_mermaid
-        flowchart = Mermaid::FlowChart.new
-        disasm = Disassembler::Mermaid.new
-
-        blocks.each do |block|
-          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.link(previous, node) if previous
-              previous = node
+        Mermaid.flowchart do |flowchart|
+          disasm = Disassembler::Squished.new
+
+          blocks.each do |block|
+            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.link(previous, node) if previous
+                previous = node
+              end
             end
           end
-        end
 
-        blocks.each do |block|
-          block.outgoing_blocks.each do |outgoing|
-            offset =
-              block.block_start + block.insns.sum(&:length) -
-                block.insns.last.length
-
-            from = flowchart.fetch("node_#{offset}")
-            to = flowchart.fetch("node_#{outgoing.block_start}")
-            flowchart.link(from, to)
+          blocks.each do |block|
+            block.outgoing_blocks.each do |outgoing|
+              offset =
+                block.block_start + block.insns.sum(&:length) -
+                  block.insns.last.length
+
+              from = flowchart.fetch("node_#{offset}")
+              to = flowchart.fetch("node_#{outgoing.block_start}")
+              flowchart.link(from, to)
+            end
           end
         end
-
-        flowchart.render
       end
 
       # This method is used to verify that the control flow graph is well