Ruby 1.9.0のバイトコードをいじり倒す(その2)

VM::InstructionSequenceの簡単な使い方です。

require 'pp'
iseq = VM::InstructionSequence.compile("print 'Hell world\n'")
PP.pp iseq.to_a
VM::InstructionSequence.load(iseq.to_a).eval

出力結果は次のようになります。

["YARVInstructionSequence/SimpleDataFormat",
 1,
 1,
 1,
 {:arg_size=>0, :local_size=>1, :stack_max=>2},
 "<compiled>",
 "<compiled>",
 0,
 :top,
 [],
 0,
 [],
 [2,
  [:putnil],
  [:putstring, "Hell world\n"],
  [:send, :print, 1, nil, 8, nil],
  [:leave]]]
Hell world

結局、VM::InstructionSequence.loadには "YARVInstructionSequence/SimpleDataFormat"で始まる配列を渡せばいいことになります。配列はRubyではバリバリにいじれますので、VM::InstructionSequence.compileで得たRubyプログラムのバイトコード表現に色々バッチを当てることが出来ます。しかし、loadに渡す配列は結構複雑な構造をしています。この配列を配列のまま扱うと結構煩雑です。
そこで、バイトコード列をオブジェクトとして扱うクラスを作りました。続きは前の日記なので注意してください

module VMLib
  class InstSeqTree
    Headers = %w(magic major_version minor_version format_type
                 misc name filename line type locals args exception_table)
    def initialize(iseq = nil)
      @klasses = {}
      @methodes = {}
      @blockes = []
      @line = {}
      @line[nil] = []
      @line_list = []
      
      @header = {}
      Headers.each do |name|
        @header[name] = nil
      end
      
      if iseq then
        init_from_ary(iseq.to_a)
      end
    end

    attr :klasses
    attr :methodes
    attr :blockes
    attr :line
    attr :line_list
    attr :header
    
    def init_from_ary(ary)
      i = 0
      Headers.each do |name|
        @header[name] = ary[i]
        i = i + 1
      end
      body = ary[i]


      curlinno = nil
      body.each do |inst|
        if inst.is_a? Integer then
          # Line number
          curlinno = "#{inst}"
          i = 1
          while @line_list.include?(curlinno) do
            curlinno = "#{inst}-#{i}"
            i = i + 1
          end
          @line_list.push curlinno
          @line[curlinno] = []

        elsif inst.is_a? Array
          case inst[0]
          when :defineclass
            if inst[2] then
              obj = InstSeqTree.new
              obj.init_from_ary(inst[2])
              @klasses[inst[1]] = obj
            end
            
          when :definemethod
            if inst[2] then
              obj = InstSeqTree.new
              obj.init_from_ary(inst[2])
              @methodes[inst[1]] = obj
            end
            
          when :send
            if inst[3] then
              obj = InstSeqTree.new
              obj.init_from_ary(inst[3])
              @blockes.push obj
            end
            
          when :invokesuper
            if inst[2] then
              obj = InstSeqTree.new
              obj.init_from_ary(inst[2])
              @blockes.push obj
            end
          end

          @line[curlinno].push inst
          
        elsif inst.is_a? Symbol
          # Label
          if !@line_list.include?(curlinno) then
            @line_list.push curlinno
          end
          @line[curlinno].push inst
            
        else
          raise inst
        end
      end
    end

    def to_a
      res = []
      Headers.each do |name|
        res.push @header[name]
      end
      
      body = []
      blno = 0
      @line_list.each do |ln|
        body.push ln.to_i
        @line[ln].each do |inst|
          if inst.is_a? Array then
            case inst[0]
            when :defineclass
              if inst[2] then
                inst[2] = @klasses[inst[1]].to_a
              end
              
            when :definemethod
              if inst[2] then
                inst[2] = @methodes[inst[1]].to_a
              end
              
            when :send
              if inst[3] then
                inst[3] = @blockes[blno].to_a
                blno = blno + 1
              end
              
            when :invokesuper
              if inst[2] then
                inst[2] = @blockes[blno].to_a
                blno = blno + 1
              end
            end
          end
          body.push inst
        end
      end

      res.push body
      res
    end
  end

  def add_code_before_block(&action)
    add_code([nil, nil, nil]) do |code, info|
      if code.line_list[0] != :before then
        code.line_list.unshift :before
        code.line[:before] = []
      end

      curlno = code.line_list[1]
      insco = action.call(header['filename'], info, curlno)
      code.line[:before] =  insco + code.line[:before]
    end
  end

  def add_code_after_block(&action)
    add_code([nil, nil, nil]) do |code, info|
      code.line.each do |no, cont|
        code.line[no] = 
          cont.inject([]) do |res, inst|
            if inst.is_a? Array then
              if inst[0] == :leave then
                curlno = code.line_list.last
                insco = action.call(header['filename'], info, curlno)
                res = res + insco
              end
            end
            res.push inst

            res
          end
      end
    end
  end

  def add_code_before_line(&action)
    add_code([nil, nil, nil]) do |code, info|
      code.line.each do |no, cont|
        pre = []
        if cont[0].is_a? Array and cont[0][0] == :getinlinecache then
          pre = [cont.shift]
        end
        code.line[no] = pre + action.call(header['filename'], info, no) + cont
      end
    end
  end

  def add_code(info, &action)
    action.call(self, info)
    
    klasses.each do |name, cont|
      cont.add_code([name, nil, nil], &action)
    end
    
    methodes.each do |name, cont|
      cont.add_code([info[0], name, nil], &action)
    end
    
    blockes.each_with_index do |cont, idx|
      cont.add_code([info[0], info[1], idx], &action)
    end
  end

end