Properly handle drops for tail calls

beepster4096 · WaffleLapkin · commit 9ee491b9052a · 2024-02-15T02:50:02.000Z
diff --git a/compiler/rustc_mir_build/src/build/expr/stmt.rs b/compiler/rustc_mir_build/src/build/expr/stmt.rs
@@ -93,7 +93,7 @@ impl<'a, 'tcx> Builder<'a, 'tcx> {
             }
             ExprKind::Become { value } => {
                 let v = &this.thir[value];
-                let ExprKind::Scope { value, .. } = v.kind else {
+                let ExprKind::Scope { value, lint_level, region_scope } = v.kind else {
                     span_bug!(v.span, "`thir_check_tail_calls` should have disallowed this {v:?}")
                 };
 
@@ -102,27 +102,31 @@ impl<'a, 'tcx> Builder<'a, 'tcx> {
                     span_bug!(v.span, "`thir_check_tail_calls` should have disallowed this {v:?}")
                 };
 
-                let fun = unpack!(block = this.as_local_operand(block, fun));
-                let args: Vec<_> = args
-                    .into_iter()
-                    .copied()
-                    .map(|arg| Spanned {
-                        node: unpack!(block = this.as_local_call_operand(block, arg)),
-                        span: this.thir.exprs[arg].span,
-                    })
-                    .collect();
+                this.in_scope((region_scope, source_info), lint_level, |this| {
+                    let fun = unpack!(block = this.as_local_operand(block, fun));
+                    let args: Vec<_> = args
+                        .into_iter()
+                        .copied()
+                        .map(|arg| Spanned {
+                            node: unpack!(block = this.as_local_call_operand(block, arg)),
+                            span: this.thir.exprs[arg].span,
+                        })
+                        .collect();
 
-                this.record_operands_moved(&args);
+                    this.record_operands_moved(&args);
 
-                debug!("expr_into_dest: fn_span={:?}", fn_span);
+                    debug!("expr_into_dest: fn_span={:?}", fn_span);
 
-                this.cfg.terminate(
-                    block,
-                    source_info,
-                    TerminatorKind::TailCall { func: fun, args, fn_span },
-                );
+                    unpack!(block = this.break_for_tail_call(block));
 
-                this.cfg.start_new_block().unit()
+                    this.cfg.terminate(
+                        block,
+                        source_info,
+                        TerminatorKind::TailCall { func: fun, args, fn_span },
+                    );
+
+                    this.cfg.start_new_block().unit()
+                })
             }
             _ => {
                 assert!(
diff --git a/compiler/rustc_mir_build/src/build/scope.rs b/compiler/rustc_mir_build/src/build/scope.rs
@@ -713,6 +713,43 @@ impl<'a, 'tcx> Builder<'a, 'tcx> {
         self.cfg.terminate(block, source_info, TerminatorKind::UnwindResume);
     }
 
+    /// Sets up the drops for explict tail calls.
+    ///
+    /// Unlike other kinds of early exits, tail calls do not go through the drop tree.
+    /// Instead, all scheduled drops are immediately added to the CFG.
+    pub(crate) fn break_for_tail_call(&mut self, mut block: BasicBlock) -> BlockAnd<()> {
+        // the innermost scope contains only the destructors for the tail call arguments
+        // we only want to drop these in case of a panic, so we skip it
+        for scope in self.scopes.scopes[1..].iter().rev().skip(1) {
+            for drop in scope.drops.iter().rev() {
+                match drop.kind {
+                    DropKind::Value => {
+                        let target = self.cfg.start_new_block();
+                        let terminator = TerminatorKind::Drop {
+                            target,
+                            // The caller will handle this if needed.
+                            // FIXME(explicit_tail_calls): add new reason
+                            unwind: UnwindAction::Terminate(UnwindTerminateReason::Abi),
+                            place: drop.local.into(),
+                            replace: false,
+                        };
+                        self.cfg.terminate(block, drop.source_info, terminator);
+                        block = target;
+                    }
+                    DropKind::Storage => {
+                        let stmt = Statement {
+                            source_info: drop.source_info,
+                            kind: StatementKind::StorageDead(drop.local),
+                        };
+                        self.cfg.push(block, stmt);
+                    }
+                }
+            }
+        }
+
+        block.unit()
+    }
+
     fn leave_top_scope(&mut self, block: BasicBlock) -> BasicBlock {
         // If we are emitting a `drop` statement, we need to have the cached
         // diverge cleanup pads ready in case that drop panics.
diff --git a/tests/mir-opt/tail_call_drops.f.ElaborateDrops.diff b/tests/mir-opt/tail_call_drops.f.ElaborateDrops.diff
@@ -0,0 +1,114 @@
+- // MIR for `f` before ElaborateDrops
++ // MIR for `f` after ElaborateDrops
+  
+  fn f() -> () {
+      let mut _0: ();
+      let mut _1: !;
+      let _2: std::string::String;
+      let _6: ();
+      let mut _7: std::string::String;
++     let mut _8: bool;
+      scope 1 {
+          debug _a => _2;
+          let _3: i32;
+          scope 2 {
+              debug _b => _3;
+              let _4: std::string::String;
+              scope 3 {
+                  debug _c => _4;
+                  let _5: std::string::String;
+                  scope 4 {
+                      debug _d => _5;
+                  }
+              }
+          }
+      }
+  
+      bb0: {
++         _8 = const false;
+          StorageLive(_2);
+          _2 = String::new() -> [return: bb1, unwind continue];
+      }
+  
+      bb1: {
+          StorageLive(_3);
+          _3 = const 12_i32;
+          StorageLive(_4);
+          _4 = String::new() -> [return: bb2, unwind: bb11];
+      }
+  
+      bb2: {
++         _8 = const true;
+          StorageLive(_5);
+          _5 = String::new() -> [return: bb3, unwind: bb10];
+      }
+  
+      bb3: {
+          StorageLive(_6);
+          StorageLive(_7);
++         _8 = const false;
+          _7 = move _4;
+          _6 = std::mem::drop::<String>(move _7) -> [return: bb4, unwind: bb8];
+      }
+  
+      bb4: {
+          StorageDead(_7);
+          StorageDead(_6);
+-         drop(_5) -> [return: bb5, unwind terminate];
++         drop(_5) -> [return: bb5, unwind: bb13];
+      }
+  
+      bb5: {
+          StorageDead(_5);
+-         drop(_4) -> [return: bb6, unwind terminate];
++         goto -> bb6;
+      }
+  
+      bb6: {
++         _8 = const false;
+          StorageDead(_4);
+          StorageDead(_3);
+-         drop(_2) -> [return: bb7, unwind terminate];
++         drop(_2) -> [return: bb7, unwind: bb13];
+      }
+  
+      bb7: {
+          StorageDead(_2);
+          tailcall g();
+      }
+  
+      bb8 (cleanup): {
+-         drop(_7) -> [return: bb9, unwind terminate];
++         goto -> bb9;
+      }
+  
+      bb9 (cleanup): {
+          drop(_5) -> [return: bb10, unwind terminate];
+      }
+  
+      bb10 (cleanup): {
+-         drop(_4) -> [return: bb11, unwind terminate];
++         goto -> bb15;
+      }
+  
+      bb11 (cleanup): {
+          drop(_2) -> [return: bb12, unwind terminate];
+      }
+  
+      bb12 (cleanup): {
+          resume;
++     }
++ 
++     bb13 (cleanup): {
++         abort;
++     }
++ 
++     bb14 (cleanup): {
++         drop(_4) -> [return: bb11, unwind terminate];
++     }
++ 
++     bb15 (cleanup): {
++         switchInt(_8) -> [0: bb11, otherwise: bb14];
+      }
+  }
+  
diff --git a/tests/mir-opt/tail_call_drops.f.built.after.mir b/tests/mir-opt/tail_call_drops.f.built.after.mir
@@ -0,0 +1,118 @@
+// MIR for `f` after built
+
+fn f() -> () {
+    let mut _0: ();
+    let mut _1: !;
+    let _2: std::string::String;
+    let _6: ();
+    let mut _7: std::string::String;
+    scope 1 {
+        debug _a => _2;
+        let _3: i32;
+        scope 2 {
+            debug _b => _3;
+            let _4: std::string::String;
+            scope 3 {
+                debug _c => _4;
+                let _5: std::string::String;
+                scope 4 {
+                    debug _d => _5;
+                }
+            }
+        }
+    }
+
+    bb0: {
+        StorageLive(_2);
+        _2 = String::new() -> [return: bb1, unwind: bb17];
+    }
+
+    bb1: {
+        FakeRead(ForLet(None), _2);
+        StorageLive(_3);
+        _3 = const 12_i32;
+        FakeRead(ForLet(None), _3);
+        StorageLive(_4);
+        _4 = String::new() -> [return: bb2, unwind: bb16];
+    }
+
+    bb2: {
+        FakeRead(ForLet(None), _4);
+        StorageLive(_5);
+        _5 = String::new() -> [return: bb3, unwind: bb15];
+    }
+
+    bb3: {
+        FakeRead(ForLet(None), _5);
+        StorageLive(_6);
+        StorageLive(_7);
+        _7 = move _4;
+        _6 = std::mem::drop::<String>(move _7) -> [return: bb4, unwind: bb13];
+    }
+
+    bb4: {
+        StorageDead(_7);
+        StorageDead(_6);
+        drop(_5) -> [return: bb5, unwind terminate];
+    }
+
+    bb5: {
+        StorageDead(_5);
+        drop(_4) -> [return: bb6, unwind terminate];
+    }
+
+    bb6: {
+        StorageDead(_4);
+        StorageDead(_3);
+        drop(_2) -> [return: bb7, unwind terminate];
+    }
+
+    bb7: {
+        StorageDead(_2);
+        tailcall g();
+    }
+
+    bb8: {
+        drop(_5) -> [return: bb9, unwind: bb15];
+    }
+
+    bb9: {
+        StorageDead(_5);
+        drop(_4) -> [return: bb10, unwind: bb16];
+    }
+
+    bb10: {
+        StorageDead(_4);
+        StorageDead(_3);
+        drop(_2) -> [return: bb11, unwind: bb17];
+    }
+
+    bb11: {
+        StorageDead(_2);
+        unreachable;
+    }
+
+    bb12: {
+        return;
+    }
+
+    bb13 (cleanup): {
+        drop(_7) -> [return: bb14, unwind terminate];
+    }
+
+    bb14 (cleanup): {
+        drop(_5) -> [return: bb15, unwind terminate];
+    }
+
+    bb15 (cleanup): {
+        drop(_4) -> [return: bb16, unwind terminate];
+    }
+
+    bb16 (cleanup): {
+        drop(_2) -> [return: bb17, unwind terminate];
+    }
+
+    bb17 (cleanup): {
+        resume;
+    }
+}
diff --git a/tests/mir-opt/tail_call_drops.rs b/tests/mir-opt/tail_call_drops.rs
@@ -0,0 +1,25 @@
+#![allow(incomplete_features)]
+#![feature(explicit_tail_calls)]
+
+// EMIT_MIR tail_call_drops.f.built.after.mir
+//   Expected result:
+//   drop(_d) -> drop(_c) -> drop(_a) -> tailcall g()
+//
+// EMIT_MIR tail_call_drops.f.ElaborateDrops.diff
+//   Expected result:
+//   drop(_d) ->             drop(_a) -> tailcall g()
+fn f() {
+
+    let _a = String::new();
+    let _b = 12;
+    let _c = String::new();
+    let _d = String::new();
+
+    drop(_c);
+
+    become g();
+}
+
+fn g() {}
+
+fn main() {}
diff --git a/tests/ui/explicit-tail-calls/drop-order.rs b/tests/ui/explicit-tail-calls/drop-order.rs
