JITアセンブラつくりかけ

x86asm.rb

require 'fiddle'

module X86ASM

  # SIB情報を表すクラス
  class SIB
    attr_accessor :base, :index, :scale, :disp

    # SIBバイト生成
    def self.make_sib(src)
      sib = 0

      # espをindexにはできない
      if ESP === src.index
        # scaleが0ならbaseと入れ替えてあげる
        if src.scale == 0
          tmp = src.index
          src.index = src.base
          src.base = tmp
        else
          raise
        end
      end

      # indexの計算
      if src.index
        sib += src.index.bit * 8

        case src.scale
        when 1
          sib += 0x00
        when 2
          sib += 0x40
        when 4
          sib += 0x80
        when 8
          sib += 0xc0
        else
          raise
        end

      # indexなし
      else
        sib += 0x20
      end
      sib += src.base.bit if src.base
      sib.chr
    end
    
    def initialize
      @scale = @disp = 0
    end

    # SIBオブジェクトに何かしら足された場合の処理([eax * 4 + 5]みたいな記述)
    def +(v)
      case v
      # 32bitレジスタが足された
      when Register32
        # 空いてるほうにセットする
        if !@base
          @base = v
        elsif !@index
          @index = v
          @scale = 0
        else
          # baseとindex両方が使用済みの場合は例外とする
          raise
        end

      # 数字が足された場合はdispに追加する
      when Integer
        @disp += v

      # SIBが足された場合
      when SIB
        # 空いてないところにデータが入っていたら例外にする
        raise if @base and v.base
        raise if @index and v.index

        if v.base
          @base = v.base
        end

        if v.index
          @index = v.index
          @scale = v.scale
        end

        @disp += v.disp
      else
        raise
      end
      self
    end

    def -(v)
      self.+(-v)
    end
  end

  # 32bitレジスタ
  class Register32
    # [eax + 10]のような指定がされた場合に呼ばれる+メソッド
    def +(v)
      tmp = SIB.new # SIBオブジェクトを生成してこれを返す
      tmp.base = self # とりあえずSIB情報のbaseにselfを設定しておく

      # 足されるものの種類により分岐する
      case v
      # 32bitレジスタが足された場合はscale0のindexとして扱う
      when Register32
        tmp.index = v
        tmp.scale = 0

      # 数字が足されたらdisp情報として保持する
      when Integer
        tmp.disp = v

      # SIBが足されたらSIBのindexとdispを移植する([eax + ecx * 4]のような指定)
      when SIB
        # baseありのSIBが足された場合はコケておく
        if v.base
          raise
        end
        tmp.index = v.index
        tmp.scale = v.scale
        tmp.disp = v.disp
      else
        raise
      end
      tmp
    end

    def -(v)
      self.+(-v)
    end

    # 掛け算がされた場合(indexのscaleとして処理)
    def *(v)
      tmp = SIB.new
      case v
      # 掛けるものは数字しかありえない
      when Integer
        # indexとして保持する
        tmp.index = self
        tmp.scale = v
      else
        raise
      end
      tmp
    end
  end

  # 32bitレジスタ定義
  class EAX < Register32;def bit;0;end;end
  class ECX < Register32;def bit;1;end;end
  class EDX < Register32;def bit;2;end;end
  class EBX < Register32;def bit;3;end;end
  class ESP < Register32;def bit;4;end;end
  class EBP < Register32;def bit;5;end;end
  class ESI < Register32;def bit;6;end;end
  class EDI < Register32;def bit;7;end;end

  # アセンブラクラス
  # こいつのインスタンスでinstance_evalされることでバイナリが生成される。
  class Assembler
    attr_reader :buf

    # ModR/Mバイトを生成する
    def self.make_modrm(r32, rm32)
      # regビット
      tmp = r32.bit * 8

      sib = ""
      disp = ""

      # rm32の型により分岐
      case rm32
      # reg to reg
      when Register32
        tmp += 0xc0 # mod11
        tmp += rm32.bit # r/mビット

      # アドレス指定
      when Array
        # いまんとこ1要素限定
        raise if rm32.size != 1

        rm = rm32[0]
        case rm
        # 中身がreg
        when Register32
          # espの場合はSIBで指定しなければならない
          if ESP === rm
            # SIBバイト生成
            tmpsib = SIB.new
            tmpsib.base = rm
            sib = SIB.make_sib(tmpsib)

            tmp += 0x04 # r/mビット(SIB指定)
            tmp += 0x00 # mod00

          # ebpの場合はdisp付きじゃないとSIBが作れない
          elsif EBP === rm
            # SIBバイト生成
            tmpsib = SIB.new
            tmpsib.base = rm
            sib = SIB.make_sib(tmpsib)

            tmp += 0x04 # r/mビット(SIB指定)
            tmp += 0x80 # mod10
            
            disp = [0].pack("V")
          else
            tmp += rm.bit # r/mビット
            tmp += 0x00 # mod00
          end

        # 中身がSIBオブジェクト
        when SIB
          # dispが無い
          if rm.disp == 0
            # EBP/ESPの場合もしくはindexがある場合はSIBバイトを生成する
            if EBP === rm.base or ESP === rm.base or rm.index
              # SIBバイト生成
              sib = SIB.make_sib(rm)

              tmp += 0x04 # r/mビット(SIB指定)

            # それ以外はレジスタのビットだけ設定(SIBバイト無し)
            else
              tmp += rm.base.bit # r/mビット
            end

            # EBPでbaseだけの場合にはdisp0で生成する必要がある
            if rm.base and !rm.index and EBP === rm.base
              tmp += 0x80 # mod10
              disp = [0].pack("V")
            else
              tmp += 0x00 # mod00
            end

          # dispある
          else
            # indexがある場合はSIBバイトを生成する
            if rm.index
              # SIBバイト生成
              sib = SIB.make_sib(rm)

              tmp += 0x04 # r/mビット(SIB指定)
              tmp += 0x80 # mod10

            # EBPでindexが無い場合
            elsif EBP === rm.base
              # SIBバイト生成
              sib = SIB.make_sib(rm)

              tmp += 0x04 # r/mビット(SIB指定)
              tmp += 0x80 # mod10

            # それ以外はレジスタのビットだけ設定(SIBバイト無し)
            else
              # SIBバイト生成
              sib = SIB.make_sib(rm)

              tmp += rm.base.bit # r/mビット
              tmp += 0x80 # mod10
            end

            # disp生成
            disp = [rm.disp].pack("V")
          end

        # 中身が数字
        when Integer
          tmp += 0x05 # r/mビット
          tmp += 0x00 # mod00

          # disp生成
          disp = [rm].pack("V")
        else
          raise
        end
      else
        raise
      end
      tmp.chr + sib + disp
    end

    def initialize
      @buf = ""
      @label = {}
      @reserve = {}
    end

    def eax;EAX.new;end
    def ecx;ECX.new;end
    def edx;EDX.new;end
    def ebx;EBX.new;end
    def esp;ESP.new;end
    def ebp;EBP.new;end
    def esi;ESI.new;end
    def edi;EDI.new;end

    def L(name)
      raise if @label.has_key?(name)
      @label[name] = @buf.size

      # このラベル名で既にジャンプ命令などが存在する
      if @reserve.has_key?(name)
        # ジャンプ命令などの相対アドレスを書き換える
        @reserve[name].each do |address|
          @buf[address, 4] = [(@buf.size - (address + 4))].pack("i!")
        end
      end
    end

    def mov(dst, src)
      if Register32 === dst
        case src
        when Integer
          @buf += (0xb8 + dst.bit).chr + [src].pack("V")
        when Register32
#          @buf += 0x8b.chr + Assembler.make_modrm(dst, src)
          @buf += 0x89.chr + (0xc0 + src.bit * 8 + dst.bit).chr # reg to reg 特化命令
        when Array
          @buf += 0x8b.chr + Assembler.make_modrm(dst, src)
        else
          raise
        end
      else
        raise
      end
    end

    def push(src)
      if Register32 === src
        @buf += (0x50 + src.bit).chr
      else
        raise
      end
    end

    def pop(dst)
      if Register32 === dst
        @buf += (0x58 + dst.bit).chr
      else
        raise
      end
    end

    def inc(dst)
      if Register32 === dst
        @buf += (0x40 + dst.bit).chr
      else
        raise
      end
    end

    def dec(dst)
      if Register32 === dst
        @buf += (0x48 + dst.bit).chr
      else
        raise
      end
    end

    def add(dst, src)
      case dst
      when Register32
        case src
        when Integer
          @buf += 0x81.chr + (0xc0 + dst.bit).chr + [src].pack("V")
        when Register32
          @buf += 0x03.chr + Assembler.make_modrm(dst, src)
        when Array
          @buf += 0x03.chr + Assembler.make_modrm(dst, src)
        else
          raise
        end
      else
        raise
      end
    end

    def sub(dst, src)
      if Register32 === dst
        if Integer === src
          @buf += 0x81.chr + (0xe8 + dst.bit).chr + [src].pack("V")
        else
          raise
        end
      else
        raise
      end
    end

    def cmp(dst, src)
      case dst
      when Register32
        case src
        when Integer
          @buf += 0x81.chr + (0xc0 + dst.bit + 0x07 * 8).chr + [src].pack("V")
        when Register32
          @buf += 0x3b.chr + Assembler.make_modrm(dst, src)
        when Array
          @buf += 0x3b.chr + Assembler.make_modrm(dst, src)
        else
          raise
        end
      else
        raise
      end
    end

    def make_jmp_rel_address(code, name)
      if @label.has_key?(name)
        @buf += code + [@label[name] - @buf.size - code.size - 4].pack("i!")
      else
        @buf += code + [0].pack("i!")
        if @reserve.has_key?(name)
          @reserve[name] << (@buf.size - 4)
        else
          @reserve[name] = [@buf.size - 4]
        end
      end
    end

    def jc(name)
      make_jmp_rel_address(0x0f.chr + 0x82.chr, name)
    end

    def jnc(name)
      make_jmp_rel_address(0x0f.chr + 0x83.chr, name)
    end

    def jz(name)
      make_jmp_rel_address(0x0f.chr + 0x84.chr, name)
    end

    def jnz(name)
      make_jmp_rel_address(0x0f.chr + 0x85.chr, name)
    end

    def jmp(name)
      make_jmp_rel_address(0xe9.chr, name)
    end

    def call(name)
      make_jmp_rel_address(0xe8.chr, name)
    end

    def ret
      @buf += 0xc3.chr
    end
  end

  def self.asm(param=[], result=Fiddle::TYPE_INT, &b)
    tmp = Assembler.new
    tmp.instance_eval &b
#    p tmp.buf
    pointer = [tmp.buf].pack("P").unpack("L!")[0]
    Fiddle::Function.new(pointer, param, result)
  end
end