[prism] correctly handle implicit BeginNode as a wrapper node

fixtures/small/begin_ensure_rescue_actual.rb

@@ -1,7 +1,9 @@
 class ForIndents
   def func
+    ants
   rescue Bees => e
     a
+    # too tired to do more
   end
 
   def func2
@@ -12,8 +14,11 @@ def func2
   end
 
   def func3
+    bees
   ensure
+    # TODO
     a
+    # more TODO
   end
 
 

fixtures/small/begin_ensure_rescue_expected.rb

@@ -1,7 +1,9 @@
 class ForIndents
   def func
+    ants
   rescue Bees => e
     a
+    # too tired to do more
   end
 
   def func2
@@ -12,8 +14,11 @@ def func2
   end
 
   def func3
+    bees
   ensure
+    # TODO
     a
+    # more TODO
   end
 
   def func4

fixtures/small/block_ensure_actual.rb

@@ -0,0 +1,18 @@
+it "calls the function" do
+  # temporarily set this
+  ENV["CALL"] = "the function"
+  the_function.call
+ensure
+  # reset
+  ENV["CALL"] = nil
+  # maybe something rubocop disable
+end
+
+it "calls the function again" do
+  ENV["CALL"] = "the function again"
+  the_function.call
+ensure
+  # reset
+  ENV["CALL"] = nil
+  # maybe something rubocop disable
+end

fixtures/small/block_ensure_expected.rb

@@ -0,0 +1,18 @@
+it "calls the function" do
+  # temporarily set this
+  ENV["CALL"] = "the function"
+  the_function.call
+ensure
+  # reset
+  ENV["CALL"] = nil
+  # maybe something rubocop disable
+end
+
+it "calls the function again" do
+  ENV["CALL"] = "the function again"
+  the_function.call
+ensure
+  # reset
+  ENV["CALL"] = nil
+  # maybe something rubocop disable
+end

librubyfmt/src/format_prism.rs

@@ -15,11 +15,23 @@ pub fn format_node(ps: &mut ParserState, node: prism::Node) {
     use prism::Node;
 
     ps.at_offset(node.location().start_offset());
-    // StatementsNode is the only real "wrapper" node, meaning it purely contains
-    // other statements which themselves would be at the start of a line.
-    // We just ignore it here -- the alternative would be callers might need to have
-    // `ps.with_start_of_line(false, ...` for statements, which is semantically confusing
-    if ps.at_start_of_line() && !matches!(node, Node::StatementsNode { .. }) {
+
+    // StatementsNode is a "wrapper" node, meaning it purely contains other statements
+    // which themselves would be at the start of a line.  We just ignore it here -- the
+    // alternative would be callers might need to have `ps.with_start_of_line(false, ...`
+    // for statements, which is semantically confusing.
+    //
+    // BeginNode without a location for `begin` occurs in `def f; ... rescue; ... end` or
+    // similar and also -- at least for its `statements` field -- serves a similar
+    // wrapping function.
+    let needs_handling = ps.at_start_of_line()
+        && !(matches!(node, Node::StatementsNode { .. })
+            || node
+                .as_begin_node()
+                .map(|b| b.begin_keyword_loc().is_none())
+                .unwrap_or(false));
+
+    if needs_handling {
         ps.emit_indent();
     }
 
@@ -409,7 +421,7 @@ pub fn format_node(ps: &mut ParserState, node: prism::Node) {
     }
 
     ps.at_offset(node.location().end_offset());
-    if ps.at_start_of_line() && !matches!(node, Node::StatementsNode { .. }) {
+    if needs_handling {
         ps.emit_newline();
     }
 }