Re-implement prev_spaces feature for pasted code

Author: tompng

lib/irb/ruby-lex.rb

@@ -367,7 +367,7 @@ def calc_indent_level(opens)
             indent_level = 0
           end
         end
-      when :on_tstring_beg, :on_regexp_beg, :on_symbeg
+      when :on_tstring_beg, :on_regexp_beg, :on_symbeg, :on_backtick
         # can be indented if t.tok starts with `%`
       when :on_words_beg, :on_qwords_beg, :on_symbols_beg, :on_qsymbols_beg, :on_embexpr_beg
         # can be indented but not indented in current implementation
@@ -386,6 +386,26 @@ def free_indent_token?(token)
     FREE_INDENT_TOKENS.include?(token&.event)
   end
 
+  # Calculates the difference of pasted code's indent and indent calculated from tokens
+  def indent_difference(lines, line_results, line_index)
+    loop do
+      _tokens, prev_opens, _next_opens, min_depth = line_results[line_index]
+      open_token = prev_opens.last
+      if !open_token || (open_token.event != :on_heredoc_beg && !free_indent_token?(open_token))
+        # If the leading whitespace is an indent, return the difference
+        indent_level = calc_indent_level(prev_opens.take(min_depth))
+        calculated_indent = 2 * indent_level
+        actual_indent = lines[line_index][/^ */].size
+        return actual_indent - calculated_indent
+      elsif open_token.event == :on_heredoc_beg && open_token.tok.match?(/^<<[^-~]/)
+        return 0
+      end
+      # If the leading whitespace is not an indent but part of a multiline token
+      # Calculate base_indent of the multiline token's beginning line
+      line_index = open_token.pos[0] - 1
+    end
+  end
+
   def process_indent_level(tokens, lines, line_index, is_newline)
     line_results = IRB::NestingParser.parse_by_line(tokens)
     result = line_results[line_index]
@@ -406,10 +426,20 @@ def process_indent_level(tokens, lines, line_index, is_newline)
     prev_open_token = prev_opens.last
     next_open_token = next_opens.last
 
+    # Calculates base indent for pasted code on the line where prev_open_token is located
+    # irb(main):001:1*   if a # base_indent is 2, indent calculated from tokens is 0
+    # irb(main):002:1*         if b # base_indent is 6, indent calculated from tokens is 2
+    # irb(main):003:0>           c # base_indent is 6, indent calculated from tokens is 4
+    if prev_open_token
+      base_indent = [0, indent_difference(lines, line_results, prev_open_token.pos[0] - 1)].max
+    else
+      base_indent = 0
+    end
+
     if free_indent_token?(prev_open_token)
       if is_newline && prev_open_token.pos[0] == line_index
         # First newline inside free-indent token
-        indent
+        base_indent + indent
       else
         # Accept any number of indent inside free-indent token
         preserve_indent
@@ -427,21 +457,21 @@ def process_indent_level(tokens, lines, line_index, is_newline)
       if prev_opens.size <= next_opens.size
         if is_newline && lines[line_index].empty? && line_results[line_index - 1][1].last != next_open_token
           # First line in heredoc
-          indent
+          tok.match?(/^<<[-~]/) ? base_indent + indent : indent
         elsif tok.match?(/^<<~/)
           # Accept extra indent spaces inside `<<~` heredoc
-          [indent, preserve_indent].max
+          [base_indent + indent, preserve_indent].max
         else
           # Accept any number of indent inside other heredoc
           preserve_indent
         end
       else
         # Heredoc close
         prev_line_indent_level = calc_indent_level(prev_opens)
-        tok.match?(/^<<[~-]/) ? 2 * (prev_line_indent_level - 1) : 0
+        tok.match?(/^<<[~-]/) ? base_indent + 2 * (prev_line_indent_level - 1) : 0
       end
     else
-      indent
+      base_indent + indent
     end
   end
 

test/irb/test_ruby_lex.rb

@@ -706,6 +706,63 @@ def dynamic_prompt(&block)
       end
     end
 
+    def test_pasted_code_keep_base_indent_spaces
+      input_with_correct_indents = [
+        Row.new(%q(    def foo), 0, 6, 1),
+        Row.new(%q(        if bar), 6, 10, 2),
+        Row.new(%q(          [1), 10, 12, 3),
+        Row.new(%q(          ]+[["a), 10, 14, 4),
+        Row.new(%q(b" + `c), 0, 14, 4),
+        Row.new(%q(d` + /e), 0, 14, 4),
+        Row.new(%q(f/ + :"g), 0, 14, 4),
+        Row.new(%q(h".tap do), 0, 16, 5),
+        Row.new(%q(                1), 16, 16, 5),
+        Row.new(%q(              end), 14, 14, 4),
+        Row.new(%q(            ]), 12, 12, 3),
+        Row.new(%q(          ]), 10, 10, 2),
+        Row.new(%q(        end), 8, 6, 1),
+        Row.new(%q(    end), 4, 0, 0),
+      ]
+      lines = []
+      input_with_correct_indents.each do |row|
+        lines << row.content
+        assert_row_indenting(lines, row)
+        assert_indent_level(lines, row.indent_level)
+      end
+    end
+
+    def test_pasted_code_keep_base_indent_spaces_with_heredoc
+      input_with_correct_indents = [
+        Row.new(%q(    def foo), 0, 6, 1),
+        Row.new(%q(        if bar), 6, 10, 2),
+        Row.new(%q(          [1), 10, 12, 3),
+        Row.new(%q(          ]+[["a), 10, 14, 4),
+        Row.new(%q(b" + <<~A + <<-B + <<C), 0, 16, 5),
+        Row.new(%q(                a#{), 16, 16, 5),
+        Row.new(%q(                1), 16, 16, 5),
+        Row.new(%q(                }), 16, 16, 5),
+        Row.new(%q(              A), 14, 16, 5),
+        Row.new(%q(                b#{), 16, 16, 5),
+        Row.new(%q(                1), 16, 16, 5),
+        Row.new(%q(                }), 16, 16, 5),
+        Row.new(%q(              B), 14, 0, 0),
+        Row.new(%q(c#{), 0, 0, 0),
+        Row.new(%q(1), 0, 0, 0),
+        Row.new(%q(}), 0, 0, 0),
+        Row.new(%q(C), 0, 14, 4),
+        Row.new(%q(            ]), 12, 12, 3),
+        Row.new(%q(          ]), 10, 10, 2),
+        Row.new(%q(        end), 8, 6, 1),
+        Row.new(%q(    end), 4, 0, 0),
+      ]
+      lines = []
+      input_with_correct_indents.each do |row|
+        lines << row.content
+        assert_row_indenting(lines, row)
+        assert_indent_level(lines, row.indent_level)
+      end
+    end
+
     def assert_dynamic_prompt(lines, expected_prompt_list)
       pend if RUBY_ENGINE == 'truffleruby'
       context = build_context