Fix inline comments on loop and conditional predicates

Author: kddnewton

lib/syntax_tree/node.rb

@@ -6004,6 +6004,8 @@ def format(q)
       q.group(0, "->") do
         if params.is_a?(Paren)
           q.format(params) unless params.contents.empty?
+        elsif params.empty? && params.comments.any?
+          q.format(params)
         elsif !params.empty?
           q.group do
             q.text("(")

lib/syntax_tree/parser.rb

@@ -302,8 +302,8 @@ def find_next_statement_start(position)
     def on_BEGIN(statements)
       lbrace = find_token(LBrace)
       rbrace = find_token(RBrace)
-      start_char = find_next_statement_start(lbrace.location.end_char)
 
+      start_char = find_next_statement_start(lbrace.location.end_char)
       statements.bind(
         start_char,
         start_char - line_counts[lbrace.location.start_line - 1].start,
@@ -340,8 +340,8 @@ def on_CHAR(value)
     def on_END(statements)
       lbrace = find_token(LBrace)
       rbrace = find_token(RBrace)
-      start_char = find_next_statement_start(lbrace.location.end_char)
 
+      start_char = find_next_statement_start(lbrace.location.end_char)
       statements.bind(
         start_char,
         start_char - line_counts[lbrace.location.start_line - 1].start,
@@ -831,8 +831,8 @@ def on_brace_block(block_var, statements)
       lbrace = find_token(LBrace)
       rbrace = find_token(RBrace)
       location = (block_var || lbrace).location
-      start_char = find_next_statement_start(location.end_char)
 
+      start_char = find_next_statement_start(location.end_char)
       statements.bind(
         start_char,
         start_char - line_counts[location.start_line - 1].start,
@@ -1329,8 +1329,8 @@ def on_else(statements)
 
       node = tokens[index]
       ending = node.value == "end" ? tokens.delete_at(index) : node
-      start_char = find_next_statement_start(keyword.location.end_char)
 
+      start_char = find_next_statement_start(keyword.location.end_char)
       statements.bind(
         start_char,
         start_char - line_counts[keyword.location.start_line - 1].start,
@@ -1355,9 +1355,10 @@ def on_elsif(predicate, statements, consequent)
       beginning = find_token(Kw, "elsif")
       ending = consequent || find_token(Kw, "end")
 
+      start_char = find_next_statement_start(predicate.location.end_char)
       statements.bind(
-        predicate.location.end_char,
-        predicate.location.end_column,
+        start_char,
+        start_char - line_counts[predicate.location.start_line - 1].start,
         ending.location.start_char,
         ending.location.start_column
       )
@@ -1598,9 +1599,12 @@ def on_for(index, collection, statements)
         tokens.delete(keyword)
       end
 
+      start_char =
+        find_next_statement_start((keyword || collection).location.end_char)
       statements.bind(
-        (keyword || collection).location.end_char,
-        (keyword || collection).location.end_column,
+        start_char,
+        start_char -
+          line_counts[(keyword || collection).location.end_line - 1].start,
         ending.location.start_char,
         ending.location.start_column
       )
@@ -1778,9 +1782,10 @@ def on_if(predicate, statements, consequent)
       beginning = find_token(Kw, "if")
       ending = consequent || find_token(Kw, "end")
 
+      start_char = find_next_statement_start(predicate.location.end_char)
       statements.bind(
-        predicate.location.end_char,
-        predicate.location.end_column,
+        start_char,
+        start_char - line_counts[predicate.location.end_line - 1].start,
         ending.location.start_char,
         ending.location.start_column
       )
@@ -2024,9 +2029,10 @@ def on_lambda(params, statements)
         closing = find_token(Kw, "end")
       end
 
+      start_char = find_next_statement_start(opening.location.end_char)
       statements.bind(
-        opening.location.end_char,
-        opening.location.end_column,
+        start_char,
+        start_char - line_counts[opening.location.end_line - 1].start,
         closing.location.start_char,
         closing.location.start_column
       )
@@ -3456,9 +3462,10 @@ def on_unless(predicate, statements, consequent)
       beginning = find_token(Kw, "unless")
       ending = consequent || find_token(Kw, "end")
 
+      start_char = find_next_statement_start(predicate.location.end_char)
       statements.bind(
-        predicate.location.end_char,
-        predicate.location.end_column,
+        start_char,
+        start_char - line_counts[predicate.location.end_line - 1].start,
         ending.location.start_char,
         ending.location.start_column
       )
@@ -3498,9 +3505,10 @@ def on_until(predicate, statements)
       end
 
       # Update the Statements location information
+      start_char = find_next_statement_start(predicate.location.end_char)
       statements.bind(
-        predicate.location.end_char,
-        predicate.location.end_column,
+        start_char,
+        start_char - line_counts[predicate.location.end_line - 1].start,
         ending.location.start_char,
         ending.location.start_column
       )
@@ -3633,9 +3641,10 @@ def on_while(predicate, statements)
       end
 
       # Update the Statements location information
+      start_char = find_next_statement_start(predicate.location.end_char)
       statements.bind(
-        predicate.location.end_char,
-        predicate.location.end_column,
+        start_char,
+        start_char - line_counts[predicate.location.end_line - 1].start,
         ending.location.start_char,
         ending.location.start_column
       )

test/fixtures/elsif.rb

@@ -17,3 +17,8 @@
 else
   qyz
 end
+%
+if true
+elsif false # comment1
+  # comment2
+end

test/fixtures/for.rb

@@ -38,3 +38,7 @@
 for foo, in [[foo, bar]]
   foo
 end
+%
+for foo in bar # comment1
+  # comment2
+end

test/fixtures/if.rb

@@ -63,3 +63,7 @@
 if (x = x + 1).to_i
   x
 end
+%
+if true # comment1
+  # comment2
+end

test/fixtures/lambda.rb

@@ -76,3 +76,7 @@
   )
 ) do
 end
+%
+-> do # comment1
+  # comment2
+end

test/fixtures/unless.rb

@@ -32,3 +32,7 @@
 unless foo
   a ? b : c
 end
+%
+unless true # comment1
+  # comment2
+end

test/fixtures/until.rb

@@ -23,3 +23,7 @@
 until (foo += 1)
   foo
 end
+%
+until true # comment1
+  # comment2
+end

test/fixtures/while.rb

@@ -23,3 +23,7 @@
 while (foo += 1)
   foo
 end
+%
+while true # comment1
+  # comment2
+end