Rollup merge of #120950 - compiler-errors:miri-async-closurs, r=RalfJung,oli-obk

Fix async closures in CTFE

First commit renames `is_coroutine_or_closure` into `is_closure_like`, because `is_coroutine_or_closure_or_coroutine_closure` seems confusing and long.

Second commit fixes some forgotten cases where we want to handle `TyKind::CoroutineClosure` the same as closures and coroutines.

The test exercises the change to `ValidityVisitor::aggregate_field_path_elem` which is the source of #120946, but not the change to `UsedParamsNeedSubstVisitor`, though I feel like it's not that big of a deal. Let me know if you'd like for me to look into constructing a test for the latter, though I have no idea what it'd look like (we can't assert against `TooGeneric` anywhere?).

Fixes #120946

r? oli-obk cc ``@RalfJung``

compiler/rustc_borrowck/src/diagnostics/conflict_errors.rs

@@ -3158,7 +3158,7 @@ impl<'cx, 'tcx> MirBorrowckCtxt<'cx, 'tcx> {
     ) -> Option<AnnotatedBorrowFnSignature<'tcx>> {
         // Define a fallback for when we can't match a closure.
         let fallback = || {
-            let is_closure = self.infcx.tcx.is_closure_or_coroutine(self.mir_def_id().to_def_id());
+            let is_closure = self.infcx.tcx.is_closure_like(self.mir_def_id().to_def_id());
             if is_closure {
                 None
             } else {
@@ -3369,7 +3369,7 @@ impl<'cx, 'tcx> MirBorrowckCtxt<'cx, 'tcx> {
         sig: ty::PolyFnSig<'tcx>,
     ) -> Option<AnnotatedBorrowFnSignature<'tcx>> {
         debug!("annotate_fn_sig: did={:?} sig={:?}", did, sig);
-        let is_closure = self.infcx.tcx.is_closure_or_coroutine(did.to_def_id());
+        let is_closure = self.infcx.tcx.is_closure_like(did.to_def_id());
         let fn_hir_id = self.infcx.tcx.local_def_id_to_hir_id(did);
         let fn_decl = self.infcx.tcx.hir().fn_decl_by_hir_id(fn_hir_id)?;
 

compiler/rustc_borrowck/src/diagnostics/mutability_errors.rs

@@ -67,7 +67,7 @@ impl<'a, 'tcx> MirBorrowckCtxt<'a, 'tcx> {
                 local,
                 projection: [proj_base @ .., ProjectionElem::Field(upvar_index, _)],
             } => {
-                debug_assert!(is_closure_or_coroutine(
+                debug_assert!(is_closure_like(
                     Place::ty_from(local, proj_base, self.body, self.infcx.tcx).ty
                 ));
 
@@ -126,9 +126,7 @@ impl<'a, 'tcx> MirBorrowckCtxt<'a, 'tcx> {
                 {
                     item_msg = access_place_desc;
                     debug_assert!(self.body.local_decls[ty::CAPTURE_STRUCT_LOCAL].ty.is_ref());
-                    debug_assert!(is_closure_or_coroutine(
-                        the_place_err.ty(self.body, self.infcx.tcx).ty
-                    ));
+                    debug_assert!(is_closure_like(the_place_err.ty(self.body, self.infcx.tcx).ty));
 
                     reason = if self.is_upvar_field_projection(access_place.as_ref()).is_some() {
                         ", as it is a captured variable in a `Fn` closure".to_string()
@@ -389,7 +387,7 @@ impl<'a, 'tcx> MirBorrowckCtxt<'a, 'tcx> {
                 local,
                 projection: [proj_base @ .., ProjectionElem::Field(upvar_index, _)],
             } => {
-                debug_assert!(is_closure_or_coroutine(
+                debug_assert!(is_closure_like(
                     Place::ty_from(local, proj_base, self.body, self.infcx.tcx).ty
                 ));
 
@@ -1474,7 +1472,8 @@ fn suggest_ampmut<'tcx>(
     }
 }
 
-fn is_closure_or_coroutine(ty: Ty<'_>) -> bool {
+/// If the type is a `Coroutine`, `Closure`, or `CoroutineClosure`
+fn is_closure_like(ty: Ty<'_>) -> bool {
     ty.is_closure() || ty.is_coroutine() || ty.is_coroutine_closure()
 }
 

compiler/rustc_borrowck/src/type_check/input_output.rs

@@ -29,7 +29,7 @@ impl<'a, 'tcx> TypeChecker<'a, 'tcx> {
     pub(super) fn check_signature_annotation(&mut self, body: &Body<'tcx>) {
         let mir_def_id = body.source.def_id().expect_local();
 
-        if !self.tcx().is_closure_or_coroutine(mir_def_id.to_def_id()) {
+        if !self.tcx().is_closure_like(mir_def_id.to_def_id()) {
             return;
         }
 

compiler/rustc_codegen_llvm/src/attributes.rs

@@ -479,7 +479,7 @@ pub fn from_fn_attrs<'ll, 'tcx>(
         // `+multivalue` feature because the purpose of the wasm abi is to match
         // the WebAssembly specification, which has this feature. This won't be
         // needed when LLVM enables this `multivalue` feature by default.
-        if !cx.tcx.is_closure_or_coroutine(instance.def_id()) {
+        if !cx.tcx.is_closure_like(instance.def_id()) {
             let abi = cx.tcx.fn_sig(instance.def_id()).skip_binder().abi();
             if abi == Abi::Wasm {
                 function_features.push("+multivalue".to_string());

compiler/rustc_codegen_ssa/src/codegen_attrs.rs

@@ -229,7 +229,7 @@ fn codegen_fn_attrs(tcx: TyCtxt<'_>, did: LocalDefId) -> CodegenFnAttrs {
             }
             sym::thread_local => codegen_fn_attrs.flags |= CodegenFnAttrFlags::THREAD_LOCAL,
             sym::track_caller => {
-                let is_closure = tcx.is_closure_or_coroutine(did.to_def_id());
+                let is_closure = tcx.is_closure_like(did.to_def_id());
 
                 if !is_closure
                     && let Some(fn_sig) = fn_sig()
@@ -274,7 +274,7 @@ fn codegen_fn_attrs(tcx: TyCtxt<'_>, did: LocalDefId) -> CodegenFnAttrs {
                 }
             }
             sym::target_feature => {
-                if !tcx.is_closure_or_coroutine(did.to_def_id())
+                if !tcx.is_closure_like(did.to_def_id())
                     && let Some(fn_sig) = fn_sig()
                     && fn_sig.skip_binder().unsafety() == hir::Unsafety::Normal
                 {
@@ -529,7 +529,7 @@ fn codegen_fn_attrs(tcx: TyCtxt<'_>, did: LocalDefId) -> CodegenFnAttrs {
     // would result in this closure being compiled without the inherited target features, but this
     // is probably a poor usage of `#[inline(always)]` and easily avoided by not using the attribute.
     if tcx.features().target_feature_11
-        && tcx.is_closure_or_coroutine(did.to_def_id())
+        && tcx.is_closure_like(did.to_def_id())
         && codegen_fn_attrs.inline != InlineAttr::Always
     {
         let owner_id = tcx.parent(did.to_def_id());

compiler/rustc_const_eval/src/interpret/util.rs

@@ -34,6 +34,7 @@ where
             match *ty.kind() {
                 ty::Param(_) => ControlFlow::Break(FoundParam),
                 ty::Closure(def_id, args)
+                | ty::CoroutineClosure(def_id, args, ..)
                 | ty::Coroutine(def_id, args, ..)
                 | ty::FnDef(def_id, args) => {
                     let instance = ty::InstanceDef::Item(def_id);

compiler/rustc_const_eval/src/interpret/validity.rs

@@ -240,8 +240,8 @@ impl<'rt, 'mir, 'tcx: 'mir, M: Machine<'mir, 'tcx>> ValidityVisitor<'rt, 'mir, '
 
         // Now we know we are projecting to a field, so figure out which one.
         match layout.ty.kind() {
-            // coroutines and closures.
-            ty::Closure(def_id, _) | ty::Coroutine(def_id, _) => {
+            // coroutines, closures, and coroutine-closures all have upvars that may be named.
+            ty::Closure(def_id, _) | ty::Coroutine(def_id, _) | ty::CoroutineClosure(def_id, _) => {
                 let mut name = None;
                 // FIXME this should be more descriptive i.e. CapturePlace instead of CapturedVar
                 // https://github.com/rust-lang/project-rfc-2229/issues/46

compiler/rustc_const_eval/src/transform/check_consts/mod.rs

@@ -72,7 +72,7 @@ impl<'mir, 'tcx> ConstCx<'mir, 'tcx> {
 
     pub fn fn_sig(&self) -> PolyFnSig<'tcx> {
         let did = self.def_id().to_def_id();
-        if self.tcx.is_closure_or_coroutine(did) {
+        if self.tcx.is_closure_like(did) {
             let ty = self.tcx.type_of(did).instantiate_identity();
             let ty::Closure(_, args) = ty.kind() else { bug!("type_of closure not ty::Closure") };
             args.as_closure().sig()

compiler/rustc_hir/src/def.rs

@@ -116,6 +116,12 @@ pub enum DefKind {
     Impl {
         of_trait: bool,
     },
+    /// A closure, coroutine, or coroutine-closure.
+    ///
+    /// These are all represented with the same `ExprKind::Closure` in the AST and HIR,
+    /// which makes it difficult to distinguish these during def collection. Therefore,
+    /// we treat them all the same, and code which needs to distinguish them can match
+    /// or `hir::ClosureKind` or `type_of`.
     Closure,
 }
 

compiler/rustc_hir_typeck/src/check.rs

@@ -90,8 +90,7 @@ pub(super) fn check_fn<'a, 'tcx>(
                 // ty.span == binding_span iff this is a closure parameter with no type ascription,
                 // or if it's an implicit `self` parameter
                 traits::SizedArgumentType(
-                    if ty_span == Some(param.span) && tcx.is_closure_or_coroutine(fn_def_id.into())
-                    {
+                    if ty_span == Some(param.span) && tcx.is_closure_like(fn_def_id.into()) {
                         None
                     } else {
                         ty.map(|ty| ty.hir_id)

compiler/rustc_hir_typeck/src/gather_locals.rs

@@ -153,7 +153,7 @@ impl<'a, 'tcx> Visitor<'tcx> for GatherLocalsVisitor<'a, 'tcx> {
                         // ascription, or if it's an implicit `self` parameter
                         traits::SizedArgumentType(
                             if ty_span == ident.span
-                                && self.fcx.tcx.is_closure_or_coroutine(self.fcx.body_id.into())
+                                && self.fcx.tcx.is_closure_like(self.fcx.body_id.into())
                             {
                                 None
                             } else {

compiler/rustc_middle/src/mir/pretty.rs

@@ -491,7 +491,7 @@ fn write_mir_sig(tcx: TyCtxt<'_>, body: &Body<'_>, w: &mut dyn io::Write) -> io:
     let kind = tcx.def_kind(def_id);
     let is_function = match kind {
         DefKind::Fn | DefKind::AssocFn | DefKind::Ctor(..) => true,
-        _ => tcx.is_closure_or_coroutine(def_id),
+        _ => tcx.is_closure_like(def_id),
     };
     match (kind, body.source.promoted) {
         (_, Some(i)) => write!(w, "{i:?} in ")?,

compiler/rustc_middle/src/ty/closure.rs

@@ -197,7 +197,7 @@ pub struct ClosureTypeInfo<'tcx> {
 }
 
 fn closure_typeinfo<'tcx>(tcx: TyCtxt<'tcx>, def: LocalDefId) -> ClosureTypeInfo<'tcx> {
-    debug_assert!(tcx.is_closure_or_coroutine(def.to_def_id()));
+    debug_assert!(tcx.is_closure_like(def.to_def_id()));
     let typeck_results = tcx.typeck(def);
     let user_provided_sig = typeck_results.user_provided_sigs[&def];
     let captures = typeck_results.closure_min_captures_flattened(def);
@@ -217,7 +217,7 @@ impl<'tcx> TyCtxt<'tcx> {
     }
 
     pub fn closure_captures(self, def_id: LocalDefId) -> &'tcx [&'tcx ty::CapturedPlace<'tcx>] {
-        if !self.is_closure_or_coroutine(def_id.to_def_id()) {
+        if !self.is_closure_like(def_id.to_def_id()) {
             return &[];
         };
         self.closure_typeinfo(def_id).captures

compiler/rustc_middle/src/ty/instance.rs

@@ -465,10 +465,7 @@ impl<'tcx> Instance<'tcx> {
     ) -> Option<Instance<'tcx>> {
         debug!("resolve(def_id={:?}, args={:?})", def_id, args);
         // Use either `resolve_closure` or `resolve_for_vtable`
-        assert!(
-            !tcx.is_closure_or_coroutine(def_id),
-            "Called `resolve_for_fn_ptr` on closure: {def_id:?}"
-        );
+        assert!(!tcx.is_closure_like(def_id), "Called `resolve_for_fn_ptr` on closure: {def_id:?}");
         Instance::resolve(tcx, param_env, def_id, args).ok().flatten().map(|mut resolved| {
             match resolved.def {
                 InstanceDef::Item(def) if resolved.def.requires_caller_location(tcx) => {
@@ -530,7 +527,7 @@ impl<'tcx> Instance<'tcx> {
                                 })
                             )
                         {
-                            if tcx.is_closure_or_coroutine(def) {
+                            if tcx.is_closure_like(def) {
                                 debug!(" => vtable fn pointer created for closure with #[track_caller]: {:?} for method {:?} {:?}",
                                        def, def_id, args);
 

compiler/rustc_middle/src/ty/util.rs

@@ -541,13 +541,15 @@ impl<'tcx> TyCtxt<'tcx> {
         Ok(())
     }
 
-    /// Returns `true` if `def_id` refers to a closure (e.g., `|x| x * 2`). Note
-    /// that closures have a `DefId`, but the closure *expression* also
-    /// has a `HirId` that is located within the context where the
-    /// closure appears (and, sadly, a corresponding `NodeId`, since
-    /// those are not yet phased out). The parent of the closure's
-    /// `DefId` will also be the context where it appears.
-    pub fn is_closure_or_coroutine(self, def_id: DefId) -> bool {
+    /// Returns `true` if `def_id` refers to a closure, coroutine, or coroutine-closure
+    /// (i.e. an async closure). These are all represented by `hir::Closure`, and all
+    /// have the same `DefKind`.
+    ///
+    /// Note that closures have a `DefId`, but the closure *expression* also has a
+    // `HirId` that is located within the context where the closure appears (and, sadly,
+    // a corresponding `NodeId`, since those are not yet phased out). The parent of
+    // the closure's `DefId` will also be the context where it appears.
+    pub fn is_closure_like(self, def_id: DefId) -> bool {
         matches!(self.def_kind(def_id), DefKind::Closure)
     }
 

compiler/rustc_mir_transform/src/coverage/spans/from_mir.rs

@@ -138,7 +138,7 @@ fn bcb_to_initial_coverage_spans<'a, 'tcx>(
             let (span, visible_macro) =
                 unexpand_into_body_span_with_visible_macro(expn_span, body_span)?;
 
-            Some(SpanFromMir::new(span, visible_macro, bcb, is_closure_or_coroutine(statement)))
+            Some(SpanFromMir::new(span, visible_macro, bcb, is_closure_like(statement)))
         });
 
         let terminator_span = Some(data.terminator()).into_iter().filter_map(move |terminator| {
@@ -153,7 +153,7 @@ fn bcb_to_initial_coverage_spans<'a, 'tcx>(
     })
 }
 
-fn is_closure_or_coroutine(statement: &Statement<'_>) -> bool {
+fn is_closure_like(statement: &Statement<'_>) -> bool {
     match statement.kind {
         StatementKind::Assign(box (_, Rvalue::Aggregate(box ref agg_kind, _))) => match agg_kind {
             AggregateKind::Closure(_, _)

compiler/rustc_monomorphize/src/collector.rs

@@ -1162,7 +1162,7 @@ fn create_fn_mono_item<'tcx>(
     let def_id = instance.def_id();
     if tcx.sess.opts.unstable_opts.profile_closures
         && def_id.is_local()
-        && tcx.is_closure_or_coroutine(def_id)
+        && tcx.is_closure_like(def_id)
     {
         crate::util::dump_closure_profile(tcx, instance);
     }

compiler/rustc_passes/src/upvars.rs

@@ -11,7 +11,7 @@ use rustc_span::Span;
 
 pub fn provide(providers: &mut Providers) {
     providers.upvars_mentioned = |tcx, def_id| {
-        if !tcx.is_closure_or_coroutine(def_id) {
+        if !tcx.is_closure_like(def_id) {
             return None;
         }
 

src/tools/clippy/clippy_lints/src/redundant_locals.rs

@@ -101,7 +101,7 @@ impl<'tcx> LateLintPass<'tcx> for RedundantLocals {
 fn is_by_value_closure_capture(cx: &LateContext<'_>, redefinition: HirId, root_variable: HirId) -> bool {
     let closure_def_id = cx.tcx.hir().enclosing_body_owner(redefinition);
 
-    cx.tcx.is_closure_or_coroutine(closure_def_id.to_def_id())
+    cx.tcx.is_closure_like(closure_def_id.to_def_id())
         && cx.tcx.closure_captures(closure_def_id).iter().any(|c| {
             matches!(c.info.capture_kind, UpvarCapture::ByValue)
                 && matches!(c.place.base, PlaceBase::Upvar(upvar) if upvar.var_path.hir_id == root_variable)

src/tools/miri/tests/pass/async-closure.rs

@@ -0,0 +1,40 @@
+#![feature(async_closure, noop_waker, async_fn_traits)]
+
+use std::future::Future;
+use std::pin::pin;
+use std::task::*;
+
+pub fn block_on<T>(fut: impl Future<Output = T>) -> T {
+    let mut fut = pin!(fut);
+    let ctx = &mut Context::from_waker(Waker::noop());
+
+    loop {
+        match fut.as_mut().poll(ctx) {
+            Poll::Pending => {}
+            Poll::Ready(t) => break t,
+        }
+    }
+}
+
+async fn call_once(f: impl async FnOnce(DropMe)) {
+    f(DropMe("world")).await;
+}
+
+#[derive(Debug)]
+struct DropMe(&'static str);
+
+impl Drop for DropMe {
+    fn drop(&mut self) {
+        println!("{}", self.0);
+    }
+}
+
+pub fn main() {
+    block_on(async {
+        let b = DropMe("hello");
+        let async_closure = async move |a: DropMe| {
+            println!("{a:?} {b:?}");
+        };
+        call_once(async_closure).await;
+    });
+}

src/tools/miri/tests/pass/async-closure.stdout

@@ -0,0 +1,3 @@
+DropMe("world") DropMe("hello")
+world
+hello