[ruby/irb] Heredoc may contain multiple newlines in a single token

Use the start token as the indentation criteria so that it works properly in
heredoc.

ref. https://github.com/ruby/reline/pull/242

9704808dfd

lib/irb/ruby-lex.rb

@@ -116,13 +116,32 @@ class RubyLex
     tokens
   end
 
+  def find_prev_spaces(line_index)
+    return 0 if @tokens.size == 0
+    md = @tokens[0][2].match(/(\A +)/)
+    prev_spaces = md.nil? ? 0 : md[1].count(' ')
+    line_count = 0
+    @tokens.each_with_index do |t, i|
+      if t[2].include?("\n")
+        line_count += t[2].count("\n")
+        if line_count >= line_index
+          return prev_spaces
+        end
+        if (@tokens.size - 1) > i
+          md = @tokens[i + 1][2].match(/(\A +)/)
+          prev_spaces = md.nil? ? 0 : md[1].count(' ')
+        end
+      end
+    end
+    prev_spaces
+  end
+
   def set_auto_indent(context)
     if @io.respond_to?(:auto_indent) and context.auto_indent_mode
       @io.auto_indent do |lines, line_index, byte_pointer, is_newline|
         if is_newline
-          md = lines[line_index - 1].match(/(\A +)/)
-          prev_spaces = md.nil? ? 0 : md[1].count(' ')
           @tokens = ripper_lex_without_warning(lines[0..line_index].join("\n"))
+          prev_spaces = find_prev_spaces(line_index)
           depth_difference = check_newline_depth_difference
           depth_difference = 0 if depth_difference < 0
           prev_spaces + depth_difference * 2

test/irb/test_ruby_lex.rb

@@ -333,6 +333,23 @@ module TestIRB
       end
     end
 
+    def test_heredoc_with_indent
+      input_with_correct_indents = [
+        Row.new(%q(<<~Q), nil, 0, 0),
+        Row.new(%q({), nil, 0, 0),
+        Row.new(%q(  #), nil, 0, 0),
+        Row.new(%q(}), nil, 0, 0),
+      ]
+
+      lines = []
+      input_with_correct_indents.each do |row|
+        lines << row.content
+        assert_indenting(lines, row.current_line_spaces, false)
+        assert_indenting(lines, row.new_line_spaces, true)
+        assert_nesting_level(lines, row.nesting_level)
+      end
+    end
+
     def test_oneliner_def_in_multiple_lines
       input_with_correct_indents = [
         Row.new(%q(def a()=[), nil, 4, 2),