Make on type formatting end completions more robust

lib/ruby_lsp/requests/on_type_formatting.rb

@@ -30,13 +30,11 @@ class OnTypeFormatting < BaseRequest
       def initialize(document, position, trigger_character)
         super(document)
 
-        scanner = document.create_scanner
-        line_begin = position[:line] == 0 ? 0 : scanner.find_char_position({ line: position[:line] - 1, character: 0 })
-        @line_end = T.let(scanner.find_char_position(position), Integer)
-        line = T.must(@document.source[line_begin..@line_end])
+        @lines = T.let(@document.source.lines, T::Array[String])
+        line = @lines[[position[:line] - 1, 0].max]
 
-        @indentation = T.let(find_indentation(line), Integer)
-        @previous_line = T.let(line.strip.chomp, String)
+        @indentation = T.let(line ? find_indentation(line) : 0, Integer)
+        @previous_line = T.let(line ? line.strip.chomp : "", String)
         @position = position
         @edits = T.let([], T::Array[Interface::TextEdit])
         @trigger_character = trigger_character
@@ -87,30 +85,15 @@ def handle_statement_end
         return unless END_REGEXES.any? { |regex| regex.match?(@previous_line) }
 
         indents = " " * @indentation
+        current_line = @lines[@position[:line]]
+        next_line = @lines[@position[:line] + 1]
 
-        if @previous_line.include?("\n")
-          # If the previous line has a line break, then it means there's content after the line break that triggered
-          # this completion. For these cases, we want to add the `end` after the content and move the cursor back to the
-          # keyword that triggered the completion
-
-          line = @position[:line]
-
-          # If there are enough lines in the document, we want to add the `end` token on the line below the extra
-          # content. Otherwise, we want to insert and extra line break ourselves
-          correction = if T.must(@document.source[@line_end..-1]).count("\n") >= 2
-            line -= 1
-            "#{indents}end"
-          else
-            "#{indents}\nend"
-          end
-
-          add_edit_with_text(correction, { line: @position[:line] + 1, character: @position[:character] })
-          move_cursor_to(line, @indentation + 3)
-        else
-          # If there's nothing after the new line break that triggered the completion, then we want to add the `end` and
-          # move the cursor to the body of the statement
+        if current_line.nil? || current_line.blank?
           add_edit_with_text(" \n#{indents}end")
           move_cursor_to(@position[:line], @indentation + 2)
+        elsif next_line.nil? || next_line.blank?
+          add_edit_with_text("#{indents}end", { line: @position[:line] + 1, character: @position[:character] })
+          move_cursor_to(@position[:line], @indentation + 3)
         end
       end
 

test/requests/on_type_formatting_test.rb

@@ -10,20 +10,20 @@ def test_adding_missing_ends
     document.push_edits(
       [{
         range: { start: { line: 0, character: 0 }, end: { line: 0, character: 0 } },
-        text: "class Foo\n",
+        text: "class Foo",
       }],
       version: 2,
     )
     document.parse
 
-    edits = RubyLsp::Requests::OnTypeFormatting.new(document, { line: 0, character: 8 }, "\n").run
+    edits = RubyLsp::Requests::OnTypeFormatting.new(document, { line: 1, character: 2 }, "\n").run
     expected_edits = [
       {
-        range: { start: { line: 0, character: 8 }, end: { line: 0, character: 8 } },
+        range: { start: { line: 1, character: 2 }, end: { line: 1, character: 2 } },
         newText: " \nend",
       },
       {
-        range: { start: { line: 0, character: 2 }, end: { line: 0, character: 2 } },
+        range: { start: { line: 1, character: 2 }, end: { line: 1, character: 2 } },
         newText: "$0",
       },
     ]
@@ -223,4 +223,20 @@ def test_breaking_line_between_keyword_and_more_content
 
     assert_equal(expected_edits.to_json, T.must(edits).to_json)
   end
+
+  def test_breaking_line_between_keyword_when_there_is_content_on_the_next_line
+    document = RubyLsp::Document.new(source: +"", version: 1, uri: "file:///fake.rb")
+
+    document.push_edits(
+      [{
+        range: { start: { line: 0, character: 0 }, end: { line: 0, character: 0 } },
+        text: "if something\n  other_thing",
+      }],
+      version: 2,
+    )
+    document.parse
+
+    edits = RubyLsp::Requests::OnTypeFormatting.new(document, { line: 0, character: 2 }, "\n").run
+    assert_empty(edits)
+  end
 end