Scan every single characters in IRB::Color.scan

Author: tompng

lib/irb/color.rb

@@ -128,10 +128,14 @@ def colorize_code(code, complete: true, ignore_error: false, colorable: colorabl
 
         symbol_state = SymbolState.new
         colored = +''
-        length = 0
-        end_seen = false
 
         scan(code, allow_last_error: !complete) do |token, str, expr|
+          # handle uncolorable code
+          if token.nil?
+            colored << Reline::Unicode.escape_for_print(str)
+            next
+          end
+
           # IRB::ColorPrinter skips colorizing fragments with any invalid token
           if ignore_error && ERROR_TOKENS.include?(token)
             return Reline::Unicode.escape_for_print(code)
@@ -147,15 +151,7 @@ def colorize_code(code, complete: true, ignore_error: false, colorable: colorabl
               colored << line
             end
           end
-          length += str.bytesize
-          end_seen = true if token == :on___end__
-        end
-
-        # give up colorizing incomplete Ripper tokens
-        unless end_seen or length == code.bytesize
-          return Reline::Unicode.escape_for_print(code)
         end
-
         colored
       end
 
@@ -170,33 +166,38 @@ def without_circular_ref(obj, seen:, &block)
       end
 
       def scan(code, allow_last_error:)
-        pos = [1, 0]
-
         verbose, $VERBOSE = $VERBOSE, nil
         RubyLex.compile_with_errors_suppressed(code) do |inner_code, line_no|
           lexer = Ripper::Lexer.new(inner_code, '(ripper)', line_no)
+          byte_pos = 0
+          line_positions = [0]
+          inner_code.lines.each do |line|
+            line_positions << line_positions.last + line.bytesize
+          end
+
+          on_scan = proc do |elem|
+            str = elem.tok
+            start_pos = line_positions[elem.pos[0] - 1] + elem.pos[1]
+            end_pos = start_pos + str.bytesize
+            next if start_pos < byte_pos
+
+            yield(nil, inner_code.byteslice(byte_pos...start_pos), nil) if byte_pos < start_pos
+            yield(elem.event, str, elem.state)
+            byte_pos = end_pos
+          end
+
           if lexer.respond_to?(:scan) # Ruby 2.7+
             lexer.scan.each do |elem|
-              str = elem.tok
               next if allow_last_error and /meets end of file|unexpected end-of-input/ =~ elem.message
-              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
-              end
-
-              yield(elem.event, str, elem.state)
+              on_scan.call(elem)
             end
           else
-            lexer.parse.each do |elem|
-              yield(elem.event, elem.tok, elem.state)
+            lexer.parse.sort_by(&:pos).each do |elem|
+              on_scan.call(elem)
             end
           end
+          # yield uncolorable DATA section
+          yield(nil, inner_code.byteslice(byte_pos...inner_code.bytesize), nil) if byte_pos < inner_code.bytesize
         end
       ensure
         $VERBOSE = verbose

test/irb/test_color.rb

@@ -88,6 +88,8 @@ def test_colorize_code
         "foo(*%W(bar))" => "foo(*#{RED}#{BOLD}%W(#{CLEAR}#{RED}bar#{CLEAR}#{RED}#{BOLD})#{CLEAR})",
         "$stdout" => "#{GREEN}#{BOLD}$stdout#{CLEAR}",
         "__END__" => "#{GREEN}__END__#{CLEAR}",
+        "foo\n__END__\nbar" => "foo\n#{GREEN}__END__#{CLEAR}\nbar",
+        "foo\n<<A\0\0bar\nA\nbaz" => "foo\n#{RED}<<A#{CLEAR}^@^@bar\n#{RED}A#{CLEAR}\nbaz",
       }
 
       # specific to Ruby 2.7+