Properly handle trailing comments after a then operator on a When or In clause.

Author: kddnewton

CHANGELOG.md

@@ -54,6 +54,7 @@ The format is based on [Keep a Changelog](http://keepachangelog.com/en/1.0.0/) a
 - Properly handle breaking parameters when there are no parentheses.
 - Properly handle trailing operators in call chains with attached comments.
 - Force using braces if the block is within the predicate of a ternary.
+- Properly handle trailing comments after a `then` operator on a `When` or `In` clause.
 
 ### Removed
 

lib/syntax_tree.rb

@@ -7166,8 +7166,9 @@ def on_in(pattern, statements, consequent)
     beginning = find_token(Kw, "in")
     ending = consequent || find_token(Kw, "end")
 
+    statements_start = find_token(Kw, "then", consume: false) || pattern
     statements.bind(
-      find_next_statement_start(pattern.location.end_char),
+      find_next_statement_start(statements_start.location.end_char),
       ending.location.start_char
     )
 
@@ -12758,8 +12759,9 @@ def on_when(arguments, statements, consequent)
     beginning = find_token(Kw, "when")
     ending = consequent || find_token(Kw, "end")
 
+    statements_start = find_token(Kw, "then", consume: false) || arguments
     statements.bind(
-      find_next_statement_start(arguments.location.end_char),
+      find_next_statement_start(statements_start.location.end_char),
       ending.location.start_char
     )