Clean graph formatting

1 files changed, 36 insertions, 10 deletions

diff --git a/yjit.rb b/yjit.rb
index 442f9ef1b6..0758aa068b 100644
--- a/yjit.rb
+++ b/yjit.rb
@@ -58,28 +58,54 @@ module YJIT
 
     def self.graphviz_for(iseq)
       iseq = RubyVM::InstructionSequence.of(iseq)
+      cs = YJIT::Disasm.new
+
+      highlight = ->(comment) { "<b>#{comment}</b>" }
+      linebreak = "<br align=\"left\"/>\n"
+
       buff = ''
-      blocks = blocks_for(iseq)
-      buff << "digraph g {"
-      buff << "  node [shape=record];"
+      blocks = blocks_for(iseq).sort_by(&:id)
+      buff << "digraph g {\n"
+
+      # Write the iseq info as a legend
+      buff << "  legend [shape=record fontsize=\"30\" fillcolor=\"lightgrey\" style=\"filled\"];\n"
+      buff << "  legend [label=\"{ Instruction Disassembly For: | {#{iseq.base_label}@#{iseq.absolute_path}:#{iseq.first_lineno}}}\"];\n"
+
+      # Subgraph contains disassembly
+      buff << "  subgraph disasm {\n"
+      buff << "  node [shape=record fontname=\"courier\"];\n"
       blocks.each do |block|
-        buff << "b#{block.id} [label=\"#{disasm_block(block).gsub(/\n/, "\\l\n")}\"];"
-        block.outgoing_ids.each do |id|
-          buff << "b#{block.id} -> b#{id};"
-        end
+        disasm = disasm_block(cs, block, highlight)
+
+        # convert newlines to breaks that graphviz understands
+        disasm.gsub!(/\n/, linebreak)
+
+        # strip leading whitespace
+        disasm.gsub!(/^\s+/, '')
+
+        buff << "b#{block.id} [label=<#{disasm}>];\n"
+        buff << block.outgoing_ids.map { |id|
+          next_block = blocks.bsearch { |nb| id <=> nb.id }
+          if next_block.address == (block.address + block.code.length)
+            "b#{block.id} -> b#{id}[color=\"green\"];"
+          else
+            "b#{block.id} -> b#{id};"
+          end
+        }.join("\n")
+        buff << "\n"
       end
+      buff << "  }"
       buff << "}"
       buff
     end
 
-    def self.disasm_block(block)
-      cs = YJIT::Disasm.new
+    def self.disasm_block(cs, block, highlight)
       comments = comments_for(block.address, block.address + block.code.length)
       comment_idx = 0
       str = ''
       cs.disasm(block.code, block.address).each do |i|
         while (comment = comments[comment_idx]) && comment.address <= i.address
-          str << "  ; #{comment.comment}\n"
+          str << "  ; #{highlight.call(comment.comment)}\n"
           comment_idx += 1
         end