[ruby/syntax_suggest] Preserve whitespace in front of comments

When removing comments I previously replaced them with a newline. This loses some context and may affect the order of the indent search which in turn affects the final result. By preserving whitespace in front of the comment, we preserve the "natural" indentation order of the line while also allowing the parser/lexer to see and join naturally consecutive (method chain) lines.

close https://github.com/ruby/syntax_suggest/pull/177

6 files changed, 58 insertions, 28 deletions

diff --git a/lib/syntax_suggest/clean_document.rb b/lib/syntax_suggest/clean_document.rb
index b572189259..08a465dfb0 100644
--- a/lib/syntax_suggest/clean_document.rb
+++ b/lib/syntax_suggest/clean_document.rb
@@ -110,7 +110,7 @@ module SyntaxSuggest
       @document.join
     end
 
-    # Remove comments and whitespace only lines
+    # Remove comments
     #
     # replace with empty newlines
     #
@@ -156,8 +156,9 @@ module SyntaxSuggest
     #
     def clean_sweep(source:)
       source.lines.map do |line|
-        if line.match?(/^\s*(#[^{].*)?$/) # https://rubular.com/r/LLE10D8HKMkJvs
-          $/
+        if line.match?(/^\s*#([^{].*)?$/) # https://rubular.com/r/LLE10D8HKMkJvs
+          whitespace = /^(?<whitespace>\s*)#([^{].*)?$/.match(line).named_captures["whitespace"] || ""
+          whitespace + $/
         else
           line
         end
diff --git a/lib/syntax_suggest/code_line.rb b/lib/syntax_suggest/code_line.rb
index dc738ab128..d771a2c0dd 100644
--- a/lib/syntax_suggest/code_line.rb
+++ b/lib/syntax_suggest/code_line.rb
@@ -48,11 +48,9 @@ module SyntaxSuggest
       strip_line = line.dup
       strip_line.lstrip!
 
-      if strip_line.empty?
-        @empty = true
-        @indent = 0
+      if (@empty = strip_line.empty?)
+        @indent = line.length - 1 # Newline removed from strip_line is not "whitespace"
       else
-        @empty = false
         @indent = line.length - strip_line.length
       end
 
diff --git a/spec/syntax_suggest/integration/syntax_suggest_spec.rb b/spec/syntax_suggest/integration/syntax_suggest_spec.rb
index bb50fafce7..21c02ca287 100644
--- a/spec/syntax_suggest/integration/syntax_suggest_spec.rb
+++ b/spec/syntax_suggest/integration/syntax_suggest_spec.rb
@@ -207,5 +207,32 @@ module SyntaxSuggest
         > 4  end
       EOM
     end
+
+    it "comment inside of a method" do
+      source = <<~'EOM'
+        class Dog
+          def bark
+            # todo
+          end
+
+          def sit
+            print "sit"
+          end
+        end
+        end # extra end
+      EOM
+
+      io = StringIO.new
+      SyntaxSuggest.call(
+        io: io,
+        source: source
+      )
+      out = io.string
+      expect(out).to include(<<~EOM)
+        >  1  class Dog
+        >  9  end
+        > 10  end # extra end
+      EOM
+    end
   end
 end
diff --git a/spec/syntax_suggest/unit/clean_document_spec.rb b/spec/syntax_suggest/unit/clean_document_spec.rb
index 4fb79efd62..bb26e2f392 100644
--- a/spec/syntax_suggest/unit/clean_document_spec.rb
+++ b/spec/syntax_suggest/unit/clean_document_spec.rb
@@ -72,6 +72,25 @@ module SyntaxSuggest
       EOM
     end
 
+
+    it "joins multi-line chained methods when separated by comments" do
+      source = <<~EOM
+        User.
+          # comment
+          where(name: 'schneems').
+          # another comment
+          first
+      EOM
+
+      doc = CleanDocument.new(source: source).join_consecutive!
+      code_lines = doc.lines
+
+      expect(code_lines[0].to_s.count($/)).to eq(5)
+      code_lines[1..-1].each do |line|
+        expect(line.to_s.strip.length).to eq(0)
+      end
+    end
+
     it "helper method: take_while_including" do
       source = <<~EOM
         User
@@ -92,27 +111,10 @@ module SyntaxSuggest
           # yolo
       EOM
 
-      out = CleanDocument.new(source: source).lines.join
-      expect(out.to_s).to eq(<<~EOM)
-
-        puts "what"
-
-      EOM
-    end
-
-    it "whitespace: removes whitespace" do
-      source = "  \n" + <<~EOM
-        puts "what"
-      EOM
-
-      out = CleanDocument.new(source: source).lines.join
-      expect(out.to_s).to eq(<<~EOM)
-
-        puts "what"
-      EOM
-
-      expect(source.lines.first.to_s).to_not eq("\n")
-      expect(out.lines.first.to_s).to eq("\n")
+      lines = CleanDocument.new(source: source).lines
+      expect(lines[0].to_s).to eq($/)
+      expect(lines[1].to_s).to eq('puts "what"' + $/)
+      expect(lines[2].to_s).to eq('  ' + $/)
     end
 
     it "trailing slash: does not join trailing do" do
diff --git a/spec/syntax_suggest/unit/code_line_spec.rb b/spec/syntax_suggest/unit/code_line_spec.rb
index cc4fa48bc9..d5b568fd19 100644
--- a/spec/syntax_suggest/unit/code_line_spec.rb
+++ b/spec/syntax_suggest/unit/code_line_spec.rb
@@ -48,6 +48,7 @@ module SyntaxSuggest
       # Indicates line 1 can join 2, 2 can join 3, but 3 won't join it's next line
       expect(code_lines.map(&:ignore_newline_not_beg?)).to eq([true, true, false, false])
     end
+
     it "trailing if" do
       code_lines = CodeLine.from_source(<<~'EOM')
         puts "lol" if foo
diff --git a/spec/syntax_suggest/unit/code_search_spec.rb b/spec/syntax_suggest/unit/code_search_spec.rb
index f836ba36f3..9a8115c0c8 100644
--- a/spec/syntax_suggest/unit/code_search_spec.rb
+++ b/spec/syntax_suggest/unit/code_search_spec.rb
@@ -338,6 +338,7 @@ module SyntaxSuggest
             end
           EOM
           search.call
+          puts "done"
 
           expect(search.invalid_blocks.join).to eq(<<~'EOM')
             Foo.call do