handle drops for tail calls

Author: beepster4096

compiler/rustc_mir_build/src/build/expr/stmt.rs

@@ -101,28 +101,34 @@ impl<'a, 'tcx> Builder<'a, 'tcx> {
             ),
             ExprKind::Become { value } => {
                 let v = &this.thir[value];
-                let ExprKind::Scope { value, .. } = v.kind else { span_bug!(v.span, "what {v:?}") };
+                let ExprKind::Scope { region_scope, lint_level, value, .. } = v.kind else { span_bug!(v.span, "what {v:?}") };
                 let v = &this.thir[value];
                 let ExprKind::Call { ref args, fun, fn_span, .. } = v.kind else { span_bug!(v.span, "what {v:?}") };
 
-                let fun = unpack!(block = this.as_local_operand(block, &this.thir[fun]));
-                let args: Vec<_> = args
-                    .into_iter()
-                    .copied()
-                    .map(|arg| unpack!(block = this.as_local_call_operand(block, &this.thir[arg])))
-                    .collect();
+                this.in_scope((region_scope, source_info), lint_level, |this| {
+                    let fun = unpack!(block = this.as_local_operand(block, &this.thir[fun]));
+                    let args: Vec<_> = args
+                        .into_iter()
+                        .copied()
+                        .map(|arg| {
+                            unpack!(block = this.as_local_call_operand(block, &this.thir[arg]))
+                        })
+                        .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!(

compiler/rustc_mir_build/src/build/scope.rs

@@ -668,8 +668,6 @@ impl<'a, 'tcx> Builder<'a, 'tcx> {
             (None, None) => {}
         }
 
-        // FIXME(explicit_tail_calls): this should drop stuff early for `become`
-
         let region_scope = self.scopes.breakable_scopes[break_index].region_scope;
         let scope_index = self.scopes.scope_index(region_scope, span);
         let drops = if destination.is_some() {
@@ -723,6 +721,42 @@ impl<'a, 'tcx> Builder<'a, 'tcx> {
         self.cfg.terminate(block, source_info, TerminatorKind::Resume);
     }
 
+    /// 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.
+                            unwind: UnwindAction::Terminate,
+                            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()
+    }
+
     // Add a dummy `Assign` statement to the CFG, with the span for the source code's `continue`
     // statement.
     fn add_dummy_assignment(&mut self, span: Span, block: BasicBlock, source_info: SourceInfo) {

tests/ui/explicit-tail-calls/ctfe-arg-bad-borrow.stderr

@@ -4,9 +4,10 @@ error[E0597]: `local` does not live long enough
 LL |     let local = Type;
    |         ----- binding `local` declared here
 LL |     become takes_borrow(&local);
-   |                         ^^^^^^- `local` dropped here while still borrowed
-   |                         |
-   |                         borrowed value does not live long enough
+   |                         ^^^^^^ borrowed value does not live long enough
+LL |
+LL | }
+   | - `local` dropped here while still borrowed
 
 error: aborting due to previous error
 

tests/ui/explicit-tail-calls/ctfe-tmp-arg-drop.rs

@@ -1,4 +1,4 @@
-#![feature(explicit_tail_calls)]
+#![feature(explicit_tail_calls, const_trait_impl, const_mut_refs)]
 
 pub const fn test(_: &View) {
     const fn takes_view(_: &View) {}
@@ -16,7 +16,7 @@ impl HasDrop {
     }
 }
 
-impl Drop for HasDrop {
+impl const Drop for HasDrop {
     fn drop(&mut self) {}
 }
 

tests/ui/explicit-tail-calls/ctfe-tmp-arg-drop.stderr

@@ -2,9 +2,8 @@ error[E0716]: temporary value dropped while borrowed
   --> $DIR/ctfe-tmp-arg-drop.rs:6:23
    |
 LL |     become takes_view(HasDrop.as_view());
-   |     ------------------^^^^^^^-----------
-   |     |                 |                |
-   |     |                 |                temporary value is freed at the end of this statement
+   |     ------------------^^^^^^^------------ temporary value is freed at the end of this statement
+   |     |                 |
    |     |                 creates a temporary value which is freed while still in use
    |     borrow later used here
    |

tests/ui/explicit-tail-calls/drop-order.rs

@@ -0,0 +1,67 @@
+// run-pass
+#![feature(explicit_tail_calls)]
+use std::cell::RefCell;
+
+fn main() {
+    let tail_counter = Default::default();
+    tail_recursive(0, &tail_counter);
+    assert_eq!(tail_counter.into_inner(), (0..128).collect::<Vec<u8>>());
+
+    let simply_counter = Default::default();
+    simply_recursive(0, &simply_counter);
+    assert_eq!(simply_counter.into_inner(), (0..128).rev().collect::<Vec<u8>>());
+
+    let scope_counter = Default::default();
+    out_of_inner_scope(&scope_counter);
+    assert_eq!(scope_counter.into_inner(), (0..8).collect::<Vec<u8>>());
+}
+
+fn tail_recursive(n: u8, order: &RefCell<Vec<u8>>) {
+    if n >= 128 {
+        return;
+    }
+
+    let _local = DropCounter(n, order);
+
+    become tail_recursive(n + 1, order)
+}
+
+fn simply_recursive(n: u8, order: &RefCell<Vec<u8>>) {
+    if n >= 128 {
+        return;
+    }
+
+    let _local = DropCounter(n, order);
+
+    return simply_recursive(n + 1, order)
+}
+
+fn out_of_inner_scope(order: &RefCell<Vec<u8>>) {
+    fn inner(order: &RefCell<Vec<u8>>) {
+        let _7 = DropCounter(7, order);
+        let _6 = DropCounter(6, order);
+    }
+
+    let _5 = DropCounter(5, order);
+    let _4 = DropCounter(4, order);
+
+    if true {
+        let _3 = DropCounter(3, order);
+        let _2 = DropCounter(2, order);
+        loop {
+            let _1 = DropCounter(1, order);
+            let _0 = DropCounter(0, order);
+
+            become inner(order);
+        }
+    }
+}
+
+struct DropCounter<'a>(u8, &'a RefCell<Vec<u8>>);
+
+impl Drop for DropCounter<'_> {
+    #[track_caller]
+    fn drop(&mut self) {
+        self.1.borrow_mut().push(self.0);
+    }
+}