Fix comment placement at the end of lambda bodies (#846)

* Fixtures

* Fix

* do/end fixtures

Author: reese

fixtures/small/lambda_trailing_comment_actual.rb

@@ -0,0 +1,23 @@
+->(local) {
+  cached_connections = applications.map {|x| x}
+  return cached_connections.size == applications.size
+  # for now assume cache miss
+}
+
+-> {
+  y = x + 1
+  y * 2
+  # trailing comment
+}
+
+->(local) do
+  cached_connections = applications.map {|x| x}
+  return cached_connections.size == applications.size
+  # for now assume cache miss
+end
+
+-> do
+  y = x + 1
+  y * 2
+  # trailing comment
+end

fixtures/small/lambda_trailing_comment_expected.rb

@@ -0,0 +1,23 @@
+-> (local) {
+  cached_connections = applications.map { |x| x }
+  return cached_connections.size == applications.size
+  # for now assume cache miss
+}
+
+-> {
+  y = x + 1
+  y * 2
+  # trailing comment
+}
+
+-> (local) do
+  cached_connections = applications.map { |x| x }
+  return cached_connections.size == applications.size
+  # for now assume cache miss
+end
+
+-> do
+  y = x + 1
+  y * 2
+  # trailing comment
+end

librubyfmt/src/format_prism.rs

@@ -4284,25 +4284,18 @@ fn format_lambda_node<'src>(ps: &mut ParserState<'src>, lambda_node: prism::Lamb
             ps.emit_space();
             ps.inline_breakable_of(BreakableDelims::for_brace_block(), |ps| {
                 if let Some(body) = lambda_node.body() {
-                    let has_multiple_statements = body
-                        .as_statements_node()
-                        .map(|statements_node| statements_node.body().len() > 1)
-                        .unwrap_or(false);
-                    if has_multiple_statements {
-                        ps.emit_soft_newline();
-                        ps.with_start_of_line(true, |ps| {
-                            format_node(ps, body);
-                        });
-                    } else {
-                        ps.with_start_of_line(false, |ps| {
-                            if let Some(node) = body.as_statements_node().unwrap().body().first() {
-                                ps.emit_soft_newline();
+                    ps.with_start_of_line(false, |ps| {
+                        let statements = body.as_statements_node().unwrap().body();
+                        if !statements.is_empty() {
+                            ps.emit_soft_newline();
+                            for node in statements.iter() {
                                 ps.emit_soft_indent();
                                 format_node(ps, node);
                                 ps.emit_soft_newline();
                             }
-                        });
-                    }
+                            ps.shift_comments();
+                        }
+                    });
                 } else if ps.has_comment_in_offset_span(
                     lambda_node.opening_loc().start_offset(),
                     lambda_node.closing_loc().end_offset(),