From a3cfe2fd083408b53fa02a90af718b49d2bbe83a Mon Sep 17 00:00:00 2001
From: David Tolnay <dtolnay@gmail.com>
Date: Mon, 2 Dec 2024 18:38:45 -0800
Subject: Visit Stmt span in MutVisitor::flat_map_stmt

---
 compiler/rustc_ast/src/mut_visit.rs | 13 +++++++------
 1 file changed, 7 insertions(+), 6 deletions(-)

(limited to 'compiler/rustc_ast')

diff --git a/compiler/rustc_ast/src/mut_visit.rs b/compiler/rustc_ast/src/mut_visit.rs
index 1e02dc8fb24..0c2fe61a698 100644
--- a/compiler/rustc_ast/src/mut_visit.rs
+++ b/compiler/rustc_ast/src/mut_visit.rs
@@ -1789,20 +1789,21 @@ pub fn noop_filter_map_expr<T: MutVisitor>(vis: &mut T, mut e: P<Expr>) -> Optio
 
 pub fn walk_flat_map_stmt<T: MutVisitor>(
     vis: &mut T,
-    Stmt { kind, mut span, mut id }: Stmt,
+    Stmt { kind, span, mut id }: Stmt,
 ) -> SmallVec<[Stmt; 1]> {
     vis.visit_id(&mut id);
-    let stmts: SmallVec<_> = walk_flat_map_stmt_kind(vis, kind)
+    let mut stmts: SmallVec<[Stmt; 1]> = walk_flat_map_stmt_kind(vis, kind)
         .into_iter()
         .map(|kind| Stmt { id, kind, span })
         .collect();
-    if stmts.len() > 1 {
-        panic!(
+    match stmts.len() {
+        0 => {}
+        1 => vis.visit_span(&mut stmts[0].span),
+        2.. => panic!(
             "cloning statement `NodeId`s is prohibited by default, \
              the visitor should implement custom statement visiting"
-        );
+        ),
     }
-    vis.visit_span(&mut span);
     stmts
 }
 
-- 
cgit 1.4.1-3-g733a5