resolve label jumps in text section

exe/vaporware

@@ -20,4 +20,4 @@ rescue => e
   exit 1
 end
 
-Vaporware::Compiler.compile(ARGV.shift, **options)
+Vaporware::Compiler.compile!(input: ARGV.shift, **options)

lib/vaporware/assembler/elf.rb

@@ -67,7 +67,6 @@ def read!(input: @input, text: @sections.text.body)
       r.each_line do |line|
         read[:main] = line.match(/main:/) unless read[:main]
         next unless read[:main] && !/main:/.match(line)
-        next if /\.L.+:/.match(line)
         text.assemble!(line)
       end
       end

lib/vaporware/assembler/elf/section/text.rb

@@ -21,20 +21,75 @@ class Vaporware::Assembler::ELF::Section::Text
   }.freeze
   HEX_PATTERN = /\A0x[0-9a-fA-F]+\z/.freeze
 
-  def initialize(**opts) = @bytes = []
+  def initialize(**opts)
+    @bytes = []
+    @lines = []
+    @label_positions = {}
+  end
 
   def assemble!(line)
-    op, *operands = line.split(/\s+/).reject { |o| o.empty? }.map { |op| op.gsub(/,/, "") }
-    @bytes << opecode(op, *operands)
+    line = line.strip
+    return if line.empty?
+    @lines << line
+  end
+
+  def build
+    @label_positions.clear
+    offset = 0
+    @lines.each do |line|
+      label = label_name(line)
+      if label
+        @label_positions[label] = offset
+        next
+      end
+      offset += instruction_size(line)
+    end
+
+    @bytes = []
+    offset = 0
+    @lines.each do |line|
+      next if label_name(line)
+      bytes = encode(line, offset)
+      @bytes << bytes
+      offset += bytes.size
+    end
+
+    @bytes.flatten.pack("C*")
   end
 
-  def build = @bytes.flatten.pack("C*")
   def size = build.bytesize
   def align(val, bytes) = (val << [0] until build.bytesize % bytes == 0)
 
   private
 
-  def opecode(op, *operands)
+  def encode(line, offset)
+    op, *operands = parse_line(line)
+    opecode(op, offset, *operands)
+  end
+
+  def parse_line(line)
+    line.split(/\s+/).reject { |o| o.empty? }.map { |op| op.gsub(/,/, "") }
+  end
+
+  def label_name(line)
+    return nil unless line.end_with?(":")
+
+    line.delete_suffix(":")
+  end
+
+  def instruction_size(line)
+    op, *operands = parse_line(line)
+    case op
+    when "je"
+      6
+    when "jmp"
+      5
+    else
+      opecode(op, 0, *operands).size
+    end
+  end
+
+  def opecode(op, offset, *operands)
     case op
     when "push"
       push(*operands)
@@ -49,7 +104,7 @@ def opecode(op, *operands)
     when "sete", "setl"
       sete(op, *operands)
     when "je", "jmp"
-      jump(op, *operands)
+      jump(op, offset, *operands)
     when "syscall"
       [0x0f, 0x05]
     when "ret"
@@ -59,23 +114,20 @@ def opecode(op, *operands)
     end
   end
 
-  def jump(op, *operands)
-    opecode = case op
-              when "je"
-                [0x74]
-              when "jmp"
-                [0xeb]
-              end
-    addr = case operands
-           in [".Lend0"]
-             [0x08]
-           in [".Lelse0"]
-             [0x0a]
-           in [".Lbegin0"]
-             opecode = [0xe9]
-             [0x48, 0xff, 0xff, 0xff]
-           end # steep:ignore
-    [opecode, addr].flatten
+  def jump(op, offset, *operands)
+    label = operands.first
+    target = @label_positions.fetch(label) do
+      raise Vaporware::Compiler::Assembler::ELF::Error, "unknown label: #{label}"
+    end
+    size = op == "je" ? 6 : 5
+    rel = target - (offset + size)
+    displacement = [rel].pack("l<").unpack("C*")
+    case op
+    when "je"
+      [0x0f, 0x84, *displacement]
+    when "jmp"
+      [0xe9, *displacement]
+    end
   end
 
   def mov(op, *operands)

lib/vaporware/compiler/generator.rb

@@ -33,7 +33,6 @@ def compile
           prologue_methods(output)
           output.puts "  .globl main" unless @shared
           to_asm(@ast, output)
-          epilogue(output)
         end
         output.close
       end
@@ -76,10 +75,12 @@ def prologue_methods(output)
       def define_method_prologue(node, output)
         output.puts "  push rbp"
         output.puts "  mov rbp, rsp"
-        output.puts "  sub rsp, #{lvar_offset(nil) * 8}"
-        _name, args, _block = node.children
-        args.children.each_with_index do |_, i|
-          output.puts "  mov [rbp-#{(i + 1) * 8}], #{REGISTER[i]}"
+        unless @defined_variables.empty?
+          output.puts "  sub rsp, #{lvar_offset(nil) * 8}"
+          _name, args, _block = node.children
+          args.children.each_with_index do |_, i|
+            output.puts "  mov [rbp-#{(i + 1) * 8}], #{REGISTER[i]}"
+          end
         end
         nil
       end
@@ -91,6 +92,7 @@ def method(method, node, output)
           next unless child.kind_of?(RubyVM::AbstractSyntaxTree::Node)
           to_asm(child, output, true)
         end
+        output.puts "  pop rax"
         ret(output)
         @doned << method
         nil
@@ -149,7 +151,6 @@ def lvar_offset(var)
       end
 
       def ret(output)
-        output.puts "  pop rax"
         output.puts "  mov rsp, rbp"
         output.puts "  pop rbp"
         output.puts "  ret"
@@ -214,17 +215,21 @@ def to_asm(node, output, method_tree = false)
           if fblock
             output.puts "  je .Lelse#{@seq}"
             to_asm(tblock, output, method_tree)
-            ret(output)
+            output.puts "  pop rax"
             output.puts "  jmp .Lend#{@seq}"
             output.puts ".Lelse#{@seq}:"
             to_asm(fblock, output, method_tree)
-            ret(output)
+            output.puts "  pop rax"
             output.puts ".Lend#{@seq}:"
           else
-            output.puts "  je .Lend#{@seq}"
-            to_asm(tblock, output, method_tree)
-            ret(output)
-            output.puts ".Lend#{@seq}:"
+            if method_tree
+              to_asm(tblock, output, method_tree)
+              ret(output)
+            else
+              output.puts "  je .Lend#{@seq}"
+              to_asm(tblock, output, method_tree)
+              output.puts ".Lend#{@seq}:"
+            end
           end
           @seq += 1
           return

sample/assembler/jmp_backward.s

@@ -0,0 +1,13 @@
+  .intel_syntax noprefix
+  .globl main
+main:
+  push rbp
+  mov rbp, rsp
+  push 1
+.Ltarget0:
+  push 2
+  jmp .Ltarget0
+  pop rax
+  mov rsp, rbp
+  pop rbp
+  ret

sample/assembler/jmp_forward.s

@@ -0,0 +1,13 @@
+  .intel_syntax noprefix
+  .globl main
+main:
+  push rbp
+  mov rbp, rsp
+  jmp .Ltarget0
+  push 1
+.Ltarget0:
+  push 2
+  pop rax
+  mov rsp, rbp
+  pop rbp
+  ret

sig/vaporware/assembler/elf/section/text.rbs

@@ -4,20 +4,22 @@ class Vaporware::Assembler::ELF::Section::Text
   OPECODE: Hash[Symbol, Array[Integer]]
 
   @bytes: Array[untyped]
+  @lines: Array[String]
+  @label_positions: Hash[String, Integer]
 
   attr_reader offset: Integer
 
   def initialize: () -> void
   def assemble!: (String) -> void
-  def align!: (Integer) -> void
+  def align: (untyped, Integer) -> untyped
   def build: () -> String
 
   private
 
-  def opecode: ((String | Symbol)?, *String) -> Array[Integer]
+  def opecode: ((String | Symbol)?, Integer, *String) -> Array[Integer]
   def mov: ((String | Symbol), *String) -> Array[Integer]
   def calc: ((String | Symbol), *String) -> Array[Integer]
-  def jump: ((String | Symbol), *String) -> Array[Integer]
+  def jump: ((String | Symbol), Integer, *String) -> Array[Integer]
   def push: (*String) -> Array[Integer]
   def pop: (*String) -> Array[Integer]
   def cmp: ((String | Symbol), *String) -> Array[Integer]

test/vaporware/assembler/binary/elf/jump_test.rb

@@ -0,0 +1,53 @@
+require "vaporware"
+require "test/unit"
+require "pathname"
+
+class ELFJumpTest < Test::Unit::TestCase
+  def test_forward_jump
+    input = Pathname.pwd.join("sample", "assembler", "jmp_forward.s").to_s
+    output = "jmp_forward.o"
+
+    assembler = Vaporware::Assembler.new(input:, output:)
+    _header, _null, text, = assembler.to_elf
+
+    expected = [
+      0x55, # push rbp
+      0x48, 0x89, 0xe5, # mov rbp, rsp
+      0xe9, 0x02, 0x00, 0x00, 0x00, # jmp .Ltarget0
+      0x6a, 0x01, # push 1
+      0x6a, 0x02, # push 2
+      0x58, # pop rax
+      0x48, 0x89, 0xec, # mov rsp, rbp
+      0x5d, # pop rbp
+      0xc3, # ret
+    ]
+
+    assert_equal(expected, text.unpack("C*")[...expected.size])
+  ensure
+    File.delete(output) if File.exist?(output)
+  end
+
+  def test_backward_jump
+    input = Pathname.pwd.join("sample", "assembler", "jmp_backward.s").to_s
+    output = "jmp_backward.o"
+
+    assembler = Vaporware::Assembler.new(input:, output:)
+    _header, _null, text, = assembler.to_elf
+
+    expected = [
+      0x55, # push rbp
+      0x48, 0x89, 0xe5, # mov rbp, rsp
+      0x6a, 0x01, # push 1
+      0x6a, 0x02, # push 2
+      0xe9, 0xf9, 0xff, 0xff, 0xff, # jmp .Ltarget0
+      0x58, # pop rax
+      0x48, 0x89, 0xec, # mov rsp, rbp
+      0x5d, # pop rbp
+      0xc3, # ret
+    ]
+
+    assert_equal(expected, text.unpack("C*")[...expected.size])
+  ensure
+    File.delete(output) if File.exist?(output)
+  end
+end