-Zfused-futures

compiler/rustc_abi/src/layout/ty.rs

@@ -65,6 +65,14 @@ rustc_index::newtype_index! {
         const FIRST_VARIANT = 0;
     }
 }
+
+impl VariantIdx {
+    /// The second variant, at index 1.
+    ///
+    /// For use alongside [`VariantIdx::ZERO`].
+    pub const ONE: VariantIdx = VariantIdx::from_u32(1);
+}
+
 #[derive(Copy, Clone, PartialEq, Eq, Hash, HashStable_Generic)]
 #[rustc_pass_by_value]
 pub struct Layout<'a>(pub Interned<'a, LayoutData<FieldIdx, VariantIdx>>);

compiler/rustc_mir_transform/src/coroutine.rs

@@ -45,6 +45,9 @@
 //! For coroutines with state 1 (returned) and state 2 (poisoned) it panics.
 //! Otherwise it continues the execution from the last suspension point.
 //!
+//! If -Zfused-futures is given however, then `Future::poll` from the state 1 (returned)
+//! will not panic and will instead return `Poll::Pending`.
+//!
 //! The other function is the drop glue for the coroutine.
 //! For coroutines with state 0 (unresumed) it drops the upvars of the coroutine.
 //! For coroutines with state 1 (returned) and state 2 (poisoned) it does nothing.
@@ -218,10 +221,17 @@ impl<'tcx> TransformVisitor<'tcx> {
         let source_info = SourceInfo::outermost(body.span);
 
         let none_value = match self.coroutine_kind {
+            CoroutineKind::Coroutine(_) => span_bug!(body.span, "`Coroutine`s cannot be fused"),
+            // Fused futures continue to return `Poll::Pending`.
             CoroutineKind::Desugared(CoroutineDesugaring::Async, _) => {
-                span_bug!(body.span, "`Future`s are not fused inherently")
+                let poll_def_id = self.tcx.require_lang_item(LangItem::Poll, body.span);
+                make_aggregate_adt(
+                    poll_def_id,
+                    VariantIdx::ONE,
+                    self.tcx.mk_args(&[self.old_ret_ty.into()]),
+                    IndexVec::new(),
+                )
             }
-            CoroutineKind::Coroutine(_) => span_bug!(body.span, "`Coroutine`s cannot be fused"),
             // `gen` continues return `None`
             CoroutineKind::Desugared(CoroutineDesugaring::Gen, _) => {
                 let option_def_id = self.tcx.require_lang_item(LangItem::Option, body.span);
@@ -278,7 +288,7 @@ impl<'tcx> TransformVisitor<'tcx> {
         statements: &mut Vec<Statement<'tcx>>,
     ) {
         const ZERO: VariantIdx = VariantIdx::ZERO;
-        const ONE: VariantIdx = VariantIdx::from_usize(1);
+        const ONE: VariantIdx = VariantIdx::ONE;
         let rvalue = match self.coroutine_kind {
             CoroutineKind::Desugared(CoroutineDesugaring::Async, _) => {
                 let poll_def_id = self.tcx.require_lang_item(LangItem::Poll, source_info.span);
@@ -1099,7 +1109,7 @@ fn return_poll_ready_assign<'tcx>(tcx: TyCtxt<'tcx>, source_info: SourceInfo) ->
         const_: Const::zero_sized(tcx.types.unit),
     }));
     let ready_val = Rvalue::Aggregate(
-        Box::new(AggregateKind::Adt(poll_def_id, VariantIdx::from_usize(0), args, None, None)),
+        Box::new(AggregateKind::Adt(poll_def_id, VariantIdx::ZERO, args, None, None)),
         IndexVec::from_raw(vec![val]),
     );
     Statement::new(source_info, StatementKind::Assign(Box::new((Place::return_place(), ready_val))))
@@ -1253,17 +1263,23 @@ fn create_coroutine_resume_function<'tcx>(
 
     if can_return {
         let block = match transform.coroutine_kind {
+            CoroutineKind::Coroutine(_) => {
+                insert_panic_block(tcx, body, ResumedAfterReturn(transform.coroutine_kind))
+            }
             CoroutineKind::Desugared(CoroutineDesugaring::Async, _)
-            | CoroutineKind::Coroutine(_) => {
+                if tcx.is_async_drop_in_place_coroutine(body.source.def_id()) =>
+            {
                 // For `async_drop_in_place<T>::{closure}` we just keep return Poll::Ready,
                 // because async drop of such coroutine keeps polling original coroutine
-                if tcx.is_async_drop_in_place_coroutine(body.source.def_id()) {
-                    insert_poll_ready_block(tcx, body)
-                } else {
-                    insert_panic_block(tcx, body, ResumedAfterReturn(transform.coroutine_kind))
-                }
+                insert_poll_ready_block(tcx, body)
             }
-            CoroutineKind::Desugared(CoroutineDesugaring::AsyncGen, _)
+            CoroutineKind::Desugared(CoroutineDesugaring::Async, _)
+                if !tcx.sess.opts.unstable_opts.fused_futures =>
+            {
+                insert_panic_block(tcx, body, ResumedAfterReturn(transform.coroutine_kind))
+            }
+            CoroutineKind::Desugared(CoroutineDesugaring::Async, _)
+            | CoroutineKind::Desugared(CoroutineDesugaring::AsyncGen, _)
             | CoroutineKind::Desugared(CoroutineDesugaring::Gen, _) => {
                 transform.insert_none_ret_block(body)
             }

compiler/rustc_session/src/options.rs

@@ -2336,6 +2336,8 @@ options! {
         "replace returns with jumps to `__x86_return_thunk` (default: `keep`)"),
     function_sections: Option<bool> = (None, parse_opt_bool, [TRACKED],
         "whether each function should go in its own section"),
+    fused_futures: bool = (false, parse_bool, [TRACKED],
+        "make compiler-generated futures return `Poll::Pending` and not `panic!` when polled after completion"),
     future_incompat_test: bool = (false, parse_bool, [UNTRACKED],
         "forces all lints to be future incompatible, used for internal testing (default: no)"),
     graphviz_dark_mode: bool = (false, parse_bool, [UNTRACKED],

tests/mir-opt/fused_futures.future-{closure#0}.coroutine_resume.0.panic-abort.mir

@@ -0,0 +1,46 @@
+// MIR for `future::{closure#0}` 0 coroutine_resume
+/* coroutine_layout = CoroutineLayout {
+    field_tys: {},
+    variant_fields: {
+        Unresumed(0): [],
+        Returned (1): [],
+        Panicked (2): [],
+    },
+    storage_conflicts: BitMatrix(0x0) {},
+} */
+
+fn future::{closure#0}(_1: Pin<&mut {async fn body of future()}>, _2: &mut Context<'_>) -> Poll<u32> {
+    debug _task_context => _2;
+    let mut _0: std::task::Poll<u32>;
+    let mut _3: u32;
+    let mut _4: u32;
+
+    bb0: {
+        _4 = discriminant((*(_1.0: &mut {async fn body of future()})));
+        switchInt(move _4) -> [0: bb1, 1: bb4, otherwise: bb5];
+    }
+
+    bb1: {
+        _3 = const 42_u32;
+        goto -> bb3;
+    }
+
+    bb2: {
+        _0 = Poll::<u32>::Ready(move _3);
+        discriminant((*(_1.0: &mut {async fn body of future()}))) = 1;
+        return;
+    }
+
+    bb3: {
+        goto -> bb2;
+    }
+
+    bb4: {
+        _0 = Poll::<u32>::Pending;
+        return;
+    }
+
+    bb5: {
+        unreachable;
+    }
+}

tests/mir-opt/fused_futures.future-{closure#0}.coroutine_resume.0.panic-unwind.mir

@@ -0,0 +1,46 @@
+// MIR for `future::{closure#0}` 0 coroutine_resume
+/* coroutine_layout = CoroutineLayout {
+    field_tys: {},
+    variant_fields: {
+        Unresumed(0): [],
+        Returned (1): [],
+        Panicked (2): [],
+    },
+    storage_conflicts: BitMatrix(0x0) {},
+} */
+
+fn future::{closure#0}(_1: Pin<&mut {async fn body of future()}>, _2: &mut Context<'_>) -> Poll<u32> {
+    debug _task_context => _2;
+    let mut _0: std::task::Poll<u32>;
+    let mut _3: u32;
+    let mut _4: u32;
+
+    bb0: {
+        _4 = discriminant((*(_1.0: &mut {async fn body of future()})));
+        switchInt(move _4) -> [0: bb1, 1: bb4, otherwise: bb5];
+    }
+
+    bb1: {
+        _3 = const 42_u32;
+        goto -> bb3;
+    }
+
+    bb2: {
+        _0 = Poll::<u32>::Ready(move _3);
+        discriminant((*(_1.0: &mut {async fn body of future()}))) = 1;
+        return;
+    }
+
+    bb3: {
+        goto -> bb2;
+    }
+
+    bb4: {
+        _0 = Poll::<u32>::Pending;
+        return;
+    }
+
+    bb5: {
+        unreachable;
+    }
+}

tests/mir-opt/fused_futures.main-{closure#0}.coroutine_resume.0.panic-abort.mir

@@ -0,0 +1,66 @@
+// MIR for `main::{closure#0}` 0 coroutine_resume
+/* coroutine_layout = CoroutineLayout {
+    field_tys: {},
+    variant_fields: {
+        Unresumed(0): [],
+        Returned (1): [],
+        Panicked (2): [],
+        Suspend0 (3): [],
+    },
+    storage_conflicts: BitMatrix(0x0) {},
+} */
+
+fn main::{closure#0}(_1: Pin<&mut {coroutine@$DIR/fused_futures.rs:17:5: 17:7}>, _2: ()) -> CoroutineState<i32, &str> {
+    let mut _0: std::ops::CoroutineState<i32, &str>;
+    let mut _3: !;
+    let _4: ();
+    let mut _5: &str;
+    let mut _6: u32;
+
+    bb0: {
+        _6 = discriminant((*(_1.0: &mut {coroutine@$DIR/fused_futures.rs:17:5: 17:7})));
+        switchInt(move _6) -> [0: bb1, 1: bb6, 3: bb5, otherwise: bb7];
+    }
+
+    bb1: {
+        StorageLive(_4);
+        _0 = CoroutineState::<i32, &str>::Yielded(const 1_i32);
+        StorageDead(_4);
+        discriminant((*(_1.0: &mut {coroutine@$DIR/fused_futures.rs:17:5: 17:7}))) = 3;
+        return;
+    }
+
+    bb2: {
+        StorageDead(_4);
+        _5 = const "foo";
+        goto -> bb4;
+    }
+
+    bb3: {
+        _0 = CoroutineState::<i32, &str>::Complete(move _5);
+        discriminant((*(_1.0: &mut {coroutine@$DIR/fused_futures.rs:17:5: 17:7}))) = 1;
+        return;
+    }
+
+    bb4: {
+        goto -> bb3;
+    }
+
+    bb5: {
+        StorageLive(_4);
+        _4 = move _2;
+        goto -> bb2;
+    }
+
+    bb6: {
+        assert(const false, "coroutine resumed after completion") -> [success: bb6, unwind unreachable];
+    }
+
+    bb7: {
+        unreachable;
+    }
+}
+
+ALLOC0 (size: 3, align: 1) {
+    66 6f 6f                                        │ foo
+}

tests/mir-opt/fused_futures.main-{closure#0}.coroutine_resume.0.panic-unwind.mir

@@ -0,0 +1,66 @@
+// MIR for `main::{closure#0}` 0 coroutine_resume
+/* coroutine_layout = CoroutineLayout {
+    field_tys: {},
+    variant_fields: {
+        Unresumed(0): [],
+        Returned (1): [],
+        Panicked (2): [],
+        Suspend0 (3): [],
+    },
+    storage_conflicts: BitMatrix(0x0) {},
+} */
+
+fn main::{closure#0}(_1: Pin<&mut {coroutine@$DIR/fused_futures.rs:17:5: 17:7}>, _2: ()) -> CoroutineState<i32, &str> {
+    let mut _0: std::ops::CoroutineState<i32, &str>;
+    let mut _3: !;
+    let _4: ();
+    let mut _5: &str;
+    let mut _6: u32;
+
+    bb0: {
+        _6 = discriminant((*(_1.0: &mut {coroutine@$DIR/fused_futures.rs:17:5: 17:7})));
+        switchInt(move _6) -> [0: bb1, 1: bb6, 3: bb5, otherwise: bb7];
+    }
+
+    bb1: {
+        StorageLive(_4);
+        _0 = CoroutineState::<i32, &str>::Yielded(const 1_i32);
+        StorageDead(_4);
+        discriminant((*(_1.0: &mut {coroutine@$DIR/fused_futures.rs:17:5: 17:7}))) = 3;
+        return;
+    }
+
+    bb2: {
+        StorageDead(_4);
+        _5 = const "foo";
+        goto -> bb4;
+    }
+
+    bb3: {
+        _0 = CoroutineState::<i32, &str>::Complete(move _5);
+        discriminant((*(_1.0: &mut {coroutine@$DIR/fused_futures.rs:17:5: 17:7}))) = 1;
+        return;
+    }
+
+    bb4: {
+        goto -> bb3;
+    }
+
+    bb5: {
+        StorageLive(_4);
+        _4 = move _2;
+        goto -> bb2;
+    }
+
+    bb6: {
+        assert(const false, "coroutine resumed after completion") -> [success: bb6, unwind continue];
+    }
+
+    bb7: {
+        unreachable;
+    }
+}
+
+ALLOC0 (size: 3, align: 1) {
+    66 6f 6f                                        │ foo
+}

tests/mir-opt/fused_futures.rs

@@ -0,0 +1,21 @@
+//@ edition:2024
+//@ compile-flags: -Zfused-futures
+// skip-filecheck
+// EMIT_MIR_FOR_EACH_PANIC_STRATEGY
+
+#![feature(coroutines, stmt_expr_attributes)]
+#![allow(unused)]
+
+// EMIT_MIR fused_futures.future-{closure#0}.coroutine_resume.0.mir
+pub async fn future() -> u32 {
+    42
+}
+
+// EMIT_MIR fused_futures.main-{closure#0}.coroutine_resume.0.mir
+fn main() {
+    let mut coroutine = #[coroutine]
+    || {
+        yield 1;
+        return "foo";
+    };
+}

tests/ui/coroutine/fused-futures.rs

@@ -0,0 +1,21 @@
+//@ run-pass
+//@ edition: 2024
+//@ compile-flags: -Zfused-futures
+
+use std::pin::pin;
+use std::task::{Context, Poll, Waker};
+
+async fn foo() -> u8 {
+    12
+}
+
+const N: usize = 10;
+
+fn main() {
+    let cx = &mut Context::from_waker(Waker::noop());
+    let mut x = pin!(foo());
+    assert_eq!(x.as_mut().poll(cx), Poll::Ready(12));
+    for _ in 0..N {
+        assert_eq!(x.as_mut().poll(cx), Poll::Pending);
+    }
+}