Now injecting the instrumentation into macro, but still in a simple way

spec/run.sh

@@ -1,2 +1,2 @@
 #!/bin/sh
-crystal src/coverage/cli.cr -- spec/main.cr --use-require="./src/coverage/runtime" -p
+crystal src/coverage/cli.cr -- spec/template.cr --use-require="./src/coverage/runtime" -p

spec/template.cr

@@ -0,0 +1,15 @@
+def code
+  {% unless true %}
+    {% if false %}
+      puts "This will not be called"
+    {% else %}
+      puts "This will be called"
+    {% end %}
+  {% end %}
+end
+
+def for_loop
+  {% for x in ["a", "b", "c"] %}
+    puts {{x}}
+  {% end %}
+end

src/coverage/inject/extensions.cr

@@ -0,0 +1,5 @@
+class Crystal::Location
+  def clone(filename = nil, line_number = nil, column_number = nil)
+    Crystal::Location.new(filename || @filename, line_number || @line_number, column_number || @column_number)
+  end
+end

src/coverage/inject/macro_utils.cr

@@ -0,0 +1,52 @@
+module MacroUtils
+  def propagate_location_in_macro(node : Crystal::ASTNode, location : Nil)
+    return nil
+  end
+
+  def propagate_location_in_macro(node : Crystal::Nop, location : Crystal::Location)
+    return location
+  end
+
+  def propagate_location_in_macro(node : Crystal::MacroIf, location : Crystal::Location)
+    location = location.clone
+
+    node.then.location = location
+    location = propagate_location_in_macro(node.then, location)
+
+    node.else.location = location
+    location = propagate_location_in_macro(node.else, location)
+    return location
+  end
+
+  def propagate_location_in_macro(node : Crystal::MacroFor, location : Crystal::Location)
+    location = location.clone
+
+    node.body.location = location
+    location = propagate_location_in_macro(node.body, location)
+
+    return location
+  end
+
+  def propagate_location_in_macro(node : Crystal::MacroLiteral, location : Crystal::Location)
+    node.location = location
+
+    new_loc = location.clone line_number: location.line_number + node.to_s.count('\n')
+
+    return new_loc
+  end
+
+  def propagate_location_in_macro(node : Crystal::Expressions, location)
+    new_loc = location.clone
+
+    node.expressions.each do |e|
+      e.location = new_loc
+      new_loc = propagate_location_in_macro(e, new_loc)
+    end
+
+    return new_loc
+  end
+
+  def propagate_location_in_macro(node : Crystal::ASTNode, location : Crystal::Location)
+    return location
+  end
+end

src/coverage/inject/source_file.cr

@@ -2,6 +2,9 @@ require "compiler/crystal/syntax/*"
 require "digest"
 require "file_utils"
 
+require "./extensions"
+require "./macro_utils"
+
 class Coverage::SourceFile < Crystal::Visitor
   # List of keywords which are trouble with variable
   # name. Some keywoards are not and won't be present in this
@@ -42,6 +45,8 @@ class Coverage::SourceFile < Crystal::Visitor
   getter! enriched_source : String
   getter required_at : Int32
 
+  include MacroUtils
+
   def self.register_file(f)
     @@already_covered_file_name.add(f.path)
     @@file_list << f
@@ -137,11 +142,14 @@ class Coverage::SourceFile < Crystal::Visitor
     "\n::Coverage.get_results(#{@@outputter}.new)"
   end
 
+  # Inject line tracer for easy debugging.
+  # add `;` after the Coverage instrumentation
+  # to avoid some with macros
   private def inject_line_traces(output)
     output.gsub(/\:\:Coverage\[([0-9]+),[ ]*([0-9]+)\](.*)/) do |str, match|
       [
         "::Coverage[", match[1],
-        ", ", match[2], "] ",
+        ", ", match[2], "]; ",
         match[3],
         inject_location(@path, @lines[match[2].to_i] - 1),
       ].join("")
@@ -167,9 +175,10 @@ class Coverage::SourceFile < Crystal::Visitor
     end
   end
 
-  private def force_inject_cover(node : Crystal::ASTNode)
-    return node if @already_covered_locations.includes?(node.location)
-    already_covered_locations << node.location
+  private def force_inject_cover(node : Crystal::ASTNode, location = nil)
+    location ||= node.location
+    return node if @already_covered_locations.includes?(location)
+    already_covered_locations << location
     return Crystal::Expressions.from([inject_coverage_tracker(node), node].unsafe_as(Array(Crystal::ASTNode)))
   end
 
@@ -223,7 +232,7 @@ class Coverage::SourceFile < Crystal::Visitor
       list_of_required_file = [] of Coverage::SourceFile
       Coverage::SourceFile.require_expanders << list_of_required_file
 
-      Dir[files_to_load].each do |file|
+      Dir[files_to_load].sort.each do |file|
         next if file !~ /\.cr$/
 
         Coverage::SourceFile.cover_file(file) do
@@ -297,12 +306,22 @@ class Coverage::SourceFile < Crystal::Visitor
   end
 
   def visit(node : Crystal::MacroIf)
+    # Fix the non-location issue on macro.
+    return false if node.location.nil?
+
+    propagate_location_in_macro(node, node.location.not_nil!)
+
     node.then = force_inject_cover(node.then)
     node.else = force_inject_cover(node.else)
-    false
+    true
   end
 
   def visit(node : Crystal::MacroFor)
+    # Fix the non-location issue on macro.
+    return false if node.location.nil?
+
+    propagate_location_in_macro(node, node.location.not_nil!)
+
     node.body = force_inject_cover(node.body)
     false
   end