Initial UnsafePinned impl [Part 2: Lowering]

Author: Sky9x

Diff for: compiler/rustc_ast_ir/src/lib.rs

@@ -17,16 +17,16 @@ use rustc_macros::{Decodable_NoContext, Encodable_NoContext, HashStable_NoContex
 pub mod visit;
 
 /// The movability of a coroutine / closure literal:
-/// whether a coroutine contains self-references, causing it to be `!Unpin`.
+/// whether a coroutine contains self-references, causing it to be `![Unsafe]Unpin`.
 #[derive(Clone, PartialEq, Eq, PartialOrd, Ord, Hash, Debug, Copy)]
 #[cfg_attr(
     feature = "nightly",
     derive(Encodable_NoContext, Decodable_NoContext, HashStable_NoContext)
 )]
 pub enum Movability {
-    /// May contain self-references, `!Unpin`.
+    /// May contain self-references, `!Unpin + !UnsafeUnpin`.
     Static,
-    /// Must not contain self-references, `Unpin`.
+    /// Must not contain self-references, `Unpin + UnsafeUnpin`.
     Movable,
 }
 

Diff for: compiler/rustc_borrowck/src/diagnostics/conflict_errors.rs

@@ -3385,7 +3385,7 @@ impl<'infcx, 'tcx> MirBorrowckCtxt<'_, 'infcx, 'tcx> {
                     Some(3)
                 } else if string.starts_with("static") {
                     // `static` is 6 chars long
-                    // This is used for `!Unpin` coroutines
+                    // This is used for immovable (self-referential) coroutines
                     Some(6)
                 } else {
                     None

Diff for: compiler/rustc_middle/src/mir/query.rs

@@ -29,6 +29,9 @@ pub struct CoroutineSavedTy<'tcx> {
     pub source_info: SourceInfo,
     /// Whether the local should be ignored for trait bound computations.
     pub ignore_for_traits: bool,
+    /// If this local is borrowed across a suspension point and thus is
+    /// "wrapped" in `UnsafePinned`. Always false for movable coroutines.
+    pub pinned: bool,
 }
 
 /// The layout of coroutine state.

Diff for: compiler/rustc_middle/src/ty/context.rs

@@ -335,6 +335,10 @@ impl<'tcx> Interner for TyCtxt<'tcx> {
         self.coroutine_hidden_types(def_id)
     }
 
+    fn coroutine_has_pinned_fields(self, def_id: DefId) -> Option<bool> {
+        self.coroutine_has_pinned_fields(def_id)
+    }
+
     fn fn_sig(self, def_id: DefId) -> ty::EarlyBinder<'tcx, ty::PolyFnSig<'tcx>> {
         self.fn_sig(def_id)
     }
@@ -734,6 +738,7 @@ bidirectional_lang_item_map! {
     TransmuteTrait,
     Tuple,
     Unpin,
+    UnsafeUnpin,
     Unsize,
 // tidy-alphabetical-end
 }

Diff for: compiler/rustc_middle/src/ty/util.rs

@@ -790,6 +790,13 @@ impl<'tcx> TyCtxt<'tcx> {
         ))
     }
 
+    /// True if the given coroutine has any pinned fields.
+    /// `None` if the coroutine is tainted by errors.
+    pub fn coroutine_has_pinned_fields(self, def_id: DefId) -> Option<bool> {
+        self.mir_coroutine_witnesses(def_id)
+            .map(|layout| layout.field_tys.iter().any(|ty| ty.pinned))
+    }
+
     /// Expands the given impl trait type, stopping if the type is recursive.
     #[instrument(skip(self), level = "debug", ret)]
     pub fn try_expand_impl_trait_type(

Diff for: compiler/rustc_mir_transform/src/coroutine.rs

@@ -639,6 +639,15 @@ struct LivenessInfo {
     /// Parallel vec to the above with SourceInfo for each yield terminator.
     source_info_at_suspension_points: Vec<SourceInfo>,
 
+    /// Coroutine saved locals that are borrowed across a suspension point.
+    /// This corresponds to locals that are "wrapped" with `UnsafePinned`.
+    ///
+    /// Note that movable coroutines do not allow borrowing locals across
+    /// suspension points and thus will always have this set empty.
+    ///
+    /// For more information, see [RFC 3467](https://rust-lang.github.io/rfcs/3467-unsafe-pinned.html).
+    saved_locals_borrowed_across_suspension_points: DenseBitSet<CoroutineSavedLocal>,
+
     /// For every saved local, the set of other saved locals that are
     /// storage-live at the same time as this local. We cannot overlap locals in
     /// the layout which have conflicting storage.
@@ -690,6 +699,8 @@ fn locals_live_across_suspend_points<'tcx>(
     let mut live_locals_at_suspension_points = Vec::new();
     let mut source_info_at_suspension_points = Vec::new();
     let mut live_locals_at_any_suspension_point = DenseBitSet::new_empty(body.local_decls.len());
+    let mut locals_borrowed_across_any_suspension_point =
+        DenseBitSet::new_empty(body.local_decls.len());
 
     for (block, data) in body.basic_blocks.iter_enumerated() {
         if let TerminatorKind::Yield { .. } = data.terminator().kind {
@@ -711,6 +722,7 @@ fn locals_live_across_suspend_points<'tcx>(
                 // of the local, which happens using the `intersect` operation below.
                 borrowed_locals_cursor.seek_before_primary_effect(loc);
                 live_locals.union(borrowed_locals_cursor.get());
+                locals_borrowed_across_any_suspension_point.union(borrowed_locals_cursor.get());
             }
 
             // Store the storage liveness for later use so we can restore the state
@@ -726,6 +738,7 @@ fn locals_live_across_suspend_points<'tcx>(
 
             // The coroutine argument is ignored.
             live_locals.remove(SELF_ARG);
+            locals_borrowed_across_any_suspension_point.remove(SELF_ARG);
 
             debug!("loc = {:?}, live_locals = {:?}", loc, live_locals);
 
@@ -741,13 +754,18 @@ fn locals_live_across_suspend_points<'tcx>(
     debug!("live_locals_anywhere = {:?}", live_locals_at_any_suspension_point);
     let saved_locals = CoroutineSavedLocals(live_locals_at_any_suspension_point);
 
+    debug!("borrowed_locals = {:?}", locals_borrowed_across_any_suspension_point);
+
     // Renumber our liveness_map bitsets to include only the locals we are
     // saving.
     let live_locals_at_suspension_points = live_locals_at_suspension_points
         .iter()
         .map(|live_here| saved_locals.renumber_bitset(live_here))
         .collect();
 
+    let saved_locals_borrowed_across_suspension_points =
+        saved_locals.renumber_bitset(&locals_borrowed_across_any_suspension_point);
+
     let storage_conflicts = compute_storage_conflicts(
         body,
         &saved_locals,
@@ -759,6 +777,7 @@ fn locals_live_across_suspend_points<'tcx>(
         saved_locals,
         live_locals_at_suspension_points,
         source_info_at_suspension_points,
+        saved_locals_borrowed_across_suspension_points,
         storage_conflicts,
         storage_liveness: storage_liveness_map,
     }
@@ -931,6 +950,7 @@ fn compute_layout<'tcx>(
         saved_locals,
         live_locals_at_suspension_points,
         source_info_at_suspension_points,
+        saved_locals_borrowed_across_suspension_points,
         storage_conflicts,
         storage_liveness,
     } = liveness;
@@ -960,8 +980,14 @@ fn compute_layout<'tcx>(
             ClearCrossCrate::Set(box LocalInfo::FakeBorrow) => true,
             _ => false,
         };
-        let decl =
-            CoroutineSavedTy { ty: decl.ty, source_info: decl.source_info, ignore_for_traits };
+        let pinned = saved_locals_borrowed_across_suspension_points.contains(saved_local);
+
+        let decl = CoroutineSavedTy {
+            ty: decl.ty,
+            source_info: decl.source_info,
+            ignore_for_traits,
+            pinned,
+        };
         debug!(?decl);
 
         tys.push(decl);

Diff for: compiler/rustc_next_trait_solver/src/solve/trait_goals.rs

@@ -1148,13 +1148,13 @@ where
 
             ty::Infer(_) | ty::Bound(_, _) => panic!("unexpected type `{self_ty:?}`"),
 
-            // Coroutines have one special built-in candidate, `Unpin`, which
-            // takes precedence over the structural auto trait candidate being
-            // assembled.
+            // Coroutines have two special built-in candidates, `Unpin` and `UnsafeUnpin`.
+            // These take precedence over the structural auto trait candidate being assembled.
             ty::Coroutine(def_id, _)
                 if self.cx().is_lang_item(goal.predicate.def_id(), TraitSolverLangItem::Unpin) =>
             {
                 match self.cx().coroutine_movability(def_id) {
+                    // immovable coroutines are *never* Unpin
                     Movability::Static => Some(Err(NoSolution)),
                     Movability::Movable => Some(
                         self.probe_builtin_trait_candidate(BuiltinImplSource::Misc).enter(|ecx| {
@@ -1163,6 +1163,21 @@ where
                     ),
                 }
             }
+            ty::Coroutine(def_id, _)
+                if self
+                    .cx()
+                    .is_lang_item(goal.predicate.def_id(), TraitSolverLangItem::UnsafeUnpin) =>
+            {
+                match self.cx().coroutine_has_pinned_fields(def_id) {
+                    Some(true) => Some(Err(NoSolution)),
+                    Some(false) => Some(
+                        self.probe_builtin_trait_candidate(BuiltinImplSource::Misc).enter(|ecx| {
+                            ecx.evaluate_added_goals_and_make_canonical_response(Certainty::Yes)
+                        }),
+                    ),
+                    None => None, // coro tainted by errors
+                }
+            }
 
             // If we still have an alias here, it must be rigid. For opaques, it's always
             // okay to consider auto traits because that'll reveal its hidden type. For

Diff for: compiler/rustc_trait_selection/src/traits/select/candidate_assembly.rs

@@ -751,6 +751,7 @@ impl<'cx, 'tcx> SelectionContext<'cx, 'tcx> {
                     // The auto impl might apply; we don't know.
                     candidates.ambiguous = true;
                 }
+
                 ty::Coroutine(coroutine_def_id, _)
                     if self.tcx().is_lang_item(def_id, LangItem::Unpin) =>
                 {
@@ -766,6 +767,18 @@ impl<'cx, 'tcx> SelectionContext<'cx, 'tcx> {
                         }
                     }
                 }
+                ty::Coroutine(coroutine_def_id, _)
+                    if self.tcx().is_lang_item(def_id, LangItem::UnsafeUnpin) =>
+                {
+                    match self.tcx().coroutine_has_pinned_fields(coroutine_def_id) {
+                        Some(true) => {}
+                        Some(false) => {
+                            candidates.vec.push(BuiltinCandidate { has_nested: false });
+                        }
+                        // coro tainted by errors
+                        None => candidates.vec.push(AutoImplCandidate),
+                    }
+                }
 
                 ty::Infer(ty::FreshTy(_) | ty::FreshIntTy(_) | ty::FreshFloatTy(_)) => {
                     bug!(

Diff for: compiler/rustc_type_ir/src/interner.rs

@@ -205,6 +205,8 @@ pub trait Interner:
         def_id: Self::DefId,
     ) -> ty::EarlyBinder<Self, ty::Binder<Self, Self::Tys>>;
 
+    fn coroutine_has_pinned_fields(self, def_id: Self::DefId) -> Option<bool>;
+
     fn fn_sig(
         self,
         def_id: Self::DefId,

Diff for: compiler/rustc_type_ir/src/lang_items.rs

@@ -38,6 +38,7 @@ pub enum TraitSolverLangItem {
     TransmuteTrait,
     Tuple,
     Unpin,
+    UnsafeUnpin,
     Unsize,
     // tidy-alphabetical-end
 }

Diff for: tests/mir-opt/building/async_await.b-{closure#0}.coroutine_resume.0.mir

@@ -21,6 +21,7 @@
                 scope: scope[0],
             },
             ignore_for_traits: false,
+            pinned: true,
         },
         _1: CoroutineSavedTy {
             ty: Coroutine(
@@ -42,6 +43,7 @@
                 scope: scope[0],
             },
             ignore_for_traits: false,
+            pinned: true,
         },
     },
     variant_fields: {