From 86cf6eabe7d95b13d67a863dd12ade2248d65fc7 Mon Sep 17 00:00:00 2001
From: Nathan Froyd <froydnj@gmail.com>
Date: Thu, 25 Dec 2025 07:08:18 -0500
Subject: [PATCH] [prism] correctly handle implicit `BeginNode` as a wrapper
 node

---
 fixtures/small/begin_ensure_rescue_actual.rb  |  5 ++++
 .../small/begin_ensure_rescue_expected.rb     |  5 ++++
 fixtures/small/block_ensure_actual.rb         | 18 ++++++++++++++
 fixtures/small/block_ensure_expected.rb       | 18 ++++++++++++++
 librubyfmt/src/format_prism.rs                | 24 ++++++++++++++-----
 5 files changed, 64 insertions(+), 6 deletions(-)
 create mode 100644 fixtures/small/block_ensure_actual.rb
 create mode 100644 fixtures/small/block_ensure_expected.rb

diff --git a/fixtures/small/begin_ensure_rescue_actual.rb b/fixtures/small/begin_ensure_rescue_actual.rb
index 2a43c7601..1c20df703 100644
--- a/fixtures/small/begin_ensure_rescue_actual.rb
+++ b/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
 
 
diff --git a/fixtures/small/begin_ensure_rescue_expected.rb b/fixtures/small/begin_ensure_rescue_expected.rb
index 7d4ee2ebf..e23da71c1 100644
--- a/fixtures/small/begin_ensure_rescue_expected.rb
+++ b/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
diff --git a/fixtures/small/block_ensure_actual.rb b/fixtures/small/block_ensure_actual.rb
new file mode 100644
index 000000000..b42ae7ba9
--- /dev/null
+++ b/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
diff --git a/fixtures/small/block_ensure_expected.rb b/fixtures/small/block_ensure_expected.rb
new file mode 100644
index 000000000..b42ae7ba9
--- /dev/null
+++ b/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
diff --git a/librubyfmt/src/format_prism.rs b/librubyfmt/src/format_prism.rs
index 146fd3172..c7755ae40 100644
--- a/librubyfmt/src/format_prism.rs
+++ b/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();
     }
 }