Properly use the disassembler for the cfg

lib/syntax_tree/yarv/control_flow_graph.rb

@@ -33,32 +33,28 @@ def initialize(iseq, insns, blocks)
       end
 
       def disasm
-        fmt = Disassembler.new
-        output = StringIO.new
-        output.puts "== cfg #{iseq.name}"
+        fmt = Disassembler.new(iseq)
+        fmt.output.puts "== cfg #{iseq.name}"
 
         blocks.each do |block|
-          output.print(block.id)
+          fmt.output.print(block.id)
 
           unless block.incoming_blocks.empty?
-            output.print(" # from: #{block.incoming_blocks.map(&:id).join(", ")}")
+            fmt.output.print(" # from: #{block.incoming_blocks.map(&:id).join(", ")}")
           end
 
-          output.puts
+          fmt.output.puts
 
-          block.insns.each do |insn|
-            output.print("    ")
-            output.puts(insn.disasm(fmt))
-          end
+          fmt.with_prefix("    ") { fmt.format_insns!(block.insns) }
 
           dests = block.outgoing_blocks.map(&:id)
           dests << "leaves" if block.insns.last.leaves?
-          output.print("        # to: #{dests.join(", ")}") unless dests.empty?
+          fmt.output.print("        # to: #{dests.join(", ")}") unless dests.empty?
 
-          output.puts
+          fmt.output.puts
         end
 
-        output.string
+        fmt.string
       end
 
       # This method is used to verify that the control flow graph is well

lib/syntax_tree/yarv/disassembler.rb

@@ -4,15 +4,16 @@ module SyntaxTree
   module YARV
     class Disassembler
       attr_reader :output, :queue
+
       attr_reader :current_prefix
       attr_accessor :current_iseq
 
-      def initialize
+      def initialize(current_iseq = nil)
         @output = StringIO.new
         @queue = []
 
         @current_prefix = ""
-        @current_iseq = nil
+        @current_iseq = current_iseq
       end
 
       ########################################################################
@@ -97,16 +98,69 @@ def object(value)
       end
 
       ########################################################################
-      # Main entrypoint
+      # Entrypoints
       ########################################################################
 
+      def string
+        output.string
+      end
+
       def format!
         while (@current_iseq = queue.shift)
           output << "\n" if output.pos > 0
           format_iseq(@current_iseq)
         end
+      end
 
-        output.string
+      def format_insns!(insns, length = 0)
+        events = []
+        lines = []
+
+        insns.each do |insn|
+          case insn
+          when Integer
+            lines << insn
+          when Symbol
+            events << event(insn)
+          when InstructionSequence::Label
+            # skip
+          else
+            output << "#{current_prefix}%04d " % length
+
+            disasm = insn.disasm(self)
+            output << disasm
+
+            if lines.any?
+              output << " " * (65 - disasm.length) if disasm.length < 65
+            elsif events.any?
+              output << " " * (39 - disasm.length) if disasm.length < 39
+            end
+
+            if lines.any?
+              output << "(%4d)" % lines.last
+              lines.clear
+            end
+
+            if events.any?
+              output << "[#{events.join}]"
+              events.clear
+            end
+
+            output << "\n"
+            length += insn.length
+          end
+        end
+      end
+
+      def with_prefix(value)
+        previous = @current_prefix
+
+        begin
+          @current_prefix = value
+          yield
+        ensure
+          @current_prefix = previous
+        end
       end
 
       private
@@ -157,55 +211,7 @@ def format_iseq(iseq)
           output << "#{current_prefix}#{locals.join("    ")}\n"
         end
 
-        length = 0
-        events = []
-        lines = []
-
-        iseq.insns.each do |insn|
-          case insn
-          when Integer
-            lines << insn
-          when Symbol
-            events << event(insn)
-          when InstructionSequence::Label
-            # skip
-          else
-            output << "#{current_prefix}%04d " % length
-
-            disasm = insn.disasm(self)
-            output << disasm
-
-            if lines.any?
-              output << " " * (65 - disasm.length) if disasm.length < 65
-            elsif events.any?
-              output << " " * (39 - disasm.length) if disasm.length < 39
-            end
-
-            if lines.any?
-              output << "(%4d)" % lines.last
-              lines.clear
-            end
-
-            if events.any?
-              output << "[#{events.join}]"
-              events.clear
-            end
-
-            output << "\n"
-            length += insn.length
-          end
-        end
-      end
-
-      def with_prefix(value)
-        previous = @current_prefix
-
-        begin
-          @current_prefix = value
-          yield
-        ensure
-          @current_prefix = previous
-        end
+        format_insns!(iseq.insns)
       end
     end
   end

lib/syntax_tree/yarv/instruction_sequence.rb

@@ -270,9 +270,10 @@ def to_a
       end
 
       def disasm
-        disassembler = Disassembler.new
-        disassembler.enqueue(self)
-        disassembler.format!
+        fmt = Disassembler.new
+        fmt.enqueue(self)
+        fmt.format!
+        fmt.string
       end
 
       # This method converts our linked list of instructions into a final array

test/yarv_test.rb

@@ -305,22 +305,22 @@ def test_cfg
       assert_equal(<<~CFG, cfg.disasm)
         == cfg <compiled>
         block_0
-            putobject                              100
-            putobject                              14
-            putobject_INT2FIX_0_
-            opt_lt                                 <calldata!mid:<, argc:1, ARGS_SIMPLE>
-            branchunless                           13
+            0000 putobject                              100
+            0002 putobject                              14
+            0004 putobject_INT2FIX_0_
+            0005 opt_lt                                 <calldata!mid:<, argc:1, ARGS_SIMPLE>
+            0007 branchunless                           13
                 # to: block_7, block_5
         block_5 # from: block_0
-            putobject                              -1
-            jump                                   14
+            0000 putobject                              -1
+            0002 jump                                   14
                 # to: block_8
         block_7 # from: block_0
-            putobject_INT2FIX_1_
+            0000 putobject_INT2FIX_1_
                 # to: block_8
         block_8 # from: block_5, block_7
-            opt_plus                               <calldata!mid:+, argc:1, ARGS_SIMPLE>
-            leave
+            0000 opt_plus                               <calldata!mid:+, argc:1, ARGS_SIMPLE>
+            0002 leave
                 # to: leaves
       CFG
     end