Colorize errors more

* lib/irb/color.rb (IRB::Color.colorize_code): colorize
  `compile_error` part as same as `on_parse_error`.

2 files changed, 16 insertions, 6 deletions

diff --git a/lib/irb/color.rb b/lib/irb/color.rb
index caa7f5032a..0dbb16b60b 100644
--- a/lib/irb/color.rb
+++ b/lib/irb/color.rb
@@ -108,25 +108,34 @@ module IRB # :nodoc:
         symbol_state = SymbolState.new
         colored = +''
         length = 0
+        pos = [1, 0]
 
         scan(code).each do |elem|
           token = elem.event
           str = elem.tok
           expr = elem.state
           in_symbol = symbol_state.scan_token(token)
-          if seq = dispatch_seq(token, expr, str, in_symbol: in_symbol)
-            Reline::Unicode.escape_for_print(str).each_line do |line|
+          next if ([elem.pos[0], elem.pos[1] + str.bytesize] <=> pos) <= 0
+          str.each_line do |line|
+            if line.end_with?("\n")
+              pos[0] += 1
+              pos[1] = 0
+            else
+              pos[1] += line.bytesize
+            end
+            line = Reline::Unicode.escape_for_print(line)
+            if seq = dispatch_seq(token, expr, line, in_symbol: in_symbol)
               colored << seq.map { |s| "\e[#{s}m" }.join('')
               colored << line.sub(/\Z/, clear)
+            else
+              colored << line
             end
-          else
-            colored << Reline::Unicode.escape_for_print(str)
           end
-          length += str.length
+          length += str.bytesize
         end
 
         # give up colorizing incomplete Ripper tokens
-        return code if length != code.length
+        return code if length != code.bytesize
 
         colored
       end
diff --git a/test/irb/test_color.rb b/test/irb/test_color.rb
index edca2d7050..778874aef1 100644
--- a/test/irb/test_color.rb
+++ b/test/irb/test_color.rb
@@ -66,6 +66,7 @@ module TestIRB
         "\t" => "\t", # not ^I
         "foo(*%W(bar))" => "foo(*#{RED}%W(#{CLEAR}#{RED}bar#{CLEAR}#{RED})#{CLEAR})",
         "$stdout" => "#{GREEN}#{BOLD}$stdout#{CLEAR}",
+        "'foo' + 'bar" => "#{RED}'#{CLEAR}#{RED}foo#{CLEAR}#{RED}'#{CLEAR} + #{RED}'#{CLEAR}#{RED}#{REVERSE}bar#{CLEAR}",
       }.each do |code, result|
         actual = with_term { IRB::Color.colorize_code(code) }
         assert_equal(result, actual, "Case: colorize_code(#{code.dump})\nResult: #{humanized_literal(actual)}")