-Zfused-futures

Author: GrigorenkoPV

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_yield_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))))
@@ -1252,13 +1262,18 @@ fn create_coroutine_resume_function<'tcx>(
     }
 
     if can_return {
+        // TODO: simplify this match
         let block = match transform.coroutine_kind {
             CoroutineKind::Desugared(CoroutineDesugaring::Async, _)
             | CoroutineKind::Coroutine(_) => {
                 // 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 if tcx.sess.opts.unstable_opts.fused_futures
+                    && !matches!(transform.coroutine_kind, CoroutineKind::Coroutine(_))
+                {
+                    transform.insert_none_ret_block(body)
                 } else {
                     insert_panic_block(tcx, body, ResumedAfterReturn(transform.coroutine_kind))
                 }

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/ui/coroutine/fused-futures.rs

@@ -0,0 +1,19 @@
+//@ run-pass
+//@ edition: 2024
+//@ compile-flags: -Zfused-futures
+
+use std::pin::pin;
+use std::task::{Context, Poll, Waker};
+
+async fn foo() {}
+
+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(()));
+    for _ in 0..N {
+        assert_eq!(x.as_mut().poll(cx), Poll::Pending);
+    }
+}