Fix optimizer to adjust jumps after instruction deletion.

src/chalk/compiler.cr

@@ -52,13 +52,16 @@ module Chalk
       return table
     end
 
-    private def create_code(tree : TreeFunction, table)
+    private def create_code(tree : TreeFunction, table, instruction = ReturnInstruction.new)
+      optimizer = Optimizer.new
       generator = CodeGenerator.new table, tree
       @logger.debug("Generating code for #{tree.name}")
-      return generator.generate!
+      code = generator.generate!
+      code << instruction
+      return optimizer.optimize(code)
     end
 
-    private def create_code(tree : BuiltinFunction, table)
+    private def create_code(tree : BuiltinFunction, table, instruction = nil)
       instructions = [] of Instruction
       tree.generate!(instructions)
       return instructions
@@ -132,9 +135,9 @@ module Chalk
       names.delete "main"
 
       main_entry = table["main"]?.as(FunctionEntry)
-      all_instructions.concat create_code(main_entry.function.as(TreeFunction), table)
+      all_instructions.concat create_code(main_entry.function.as(TreeFunction),
+        table, JumpRelativeInstruction.new 0)
       main_entry.addr = 0
-      all_instructions << JumpRelativeInstruction.new 0
 
       names.each do |name|
         entry = table[name]?.as(FunctionEntry)

src/chalk/optimizer.cr

@@ -1,21 +1,19 @@
 module Chalk
   class Optimizer
-     private def check_dead(inst)
+    private def check_dead(inst)
       if inst.is_a?(LoadRegInstruction)
         return inst.from == inst.into
       end
       return false
     end
 
-    private def optimize!(instructions, range)
-      offset = 0
+    private def optimize!(instructions, range, deletions, deletions_at)
       range.each do |index|
-        if check_dead(instructions[index + offset])
-          instructions.delete_at(index + offset)
-          offset -= 1
+        if check_dead(instructions[index])
+          deletions << index
         end
+        deletions_at[index] = deletions.size
       end
-      return offset
     end
 
     def optimize(instructions)
@@ -23,18 +21,34 @@ module Chalk
       block_boundaries = [instructions.size]
       instructions.each_with_index do |inst, i|
         if inst.is_a?(JumpRelativeInstruction)
+          block_boundaries << i
           block_boundaries << (inst.offset + i)
         end
       end
-      block_boundaries.sort!
+      block_boundaries.uniq!.sort!
 
       previous = 0
-      offset = 0
+      deletions = [] of Int32
+      deletions_at = {} of Int32 => Int32
       block_boundaries.each do |boundary|
-        range = (previous + offset)...(boundary + offset)
-        offset += optimize!(instructions, range)
+        range = previous...boundary
+        optimize!(instructions, range, deletions, deletions_at)
         previous = boundary
       end
+
+      instructions.each_with_index do |inst, i|
+        next if !inst.is_a?(JumpRelativeInstruction)
+        jump_to = inst.offset + i
+        next unless deletions_at[jump_to]?
+        deletions_offset = deletions_at[i] - deletions_at[jump_to]
+        inst.offset += deletions_offset
+      end
+
+      deletions.reverse!
+      deletions.each do |i|
+        instructions.delete_at i
+      end
+
       return instructions
     end
   end