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

@@ -734,6 +734,7 @@ bidirectional_lang_item_map! {
     TransmuteTrait,
     Tuple,
     Unpin,
+    UnsafeUnpin,
     Unsize,
 // tidy-alphabetical-end
 }

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,16 @@ 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) =>
+                if matches!(
+                    self.cx().as_lang_item(goal.predicate.def_id()),
+                    Some(TraitSolverLangItem::Unpin | TraitSolverLangItem::UnsafeUnpin)
+                ) =>
             {
                 match self.cx().coroutine_movability(def_id) {
+                    // immovable coroutines are *never* [Unsafe]Unpin
                     Movability::Static => Some(Err(NoSolution)),
                     Movability::Movable => Some(
                         self.probe_builtin_trait_candidate(BuiltinImplSource::Misc).enter(|ecx| {

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

@@ -752,15 +752,16 @@ impl<'cx, 'tcx> SelectionContext<'cx, 'tcx> {
                     candidates.ambiguous = true;
                 }
                 ty::Coroutine(coroutine_def_id, _)
-                    if self.tcx().is_lang_item(def_id, LangItem::Unpin) =>
+                    if self.tcx().is_lang_item(def_id, LangItem::Unpin)
+                        || self.tcx().is_lang_item(def_id, LangItem::UnsafeUnpin) =>
                 {
                     match self.tcx().coroutine_movability(coroutine_def_id) {
                         hir::Movability::Static => {
-                            // Immovable coroutines are never `Unpin`, so
+                            // Immovable coroutines are never `[Unsafe]Unpin`, so
                             // suppress the normal auto-impl candidate for it.
                         }
                         hir::Movability::Movable => {
-                            // Movable coroutines are always `Unpin`, so add an
+                            // Movable coroutines are always `[Unsafe]Unpin`, so add an
                             // unconditional builtin candidate.
                             candidates.vec.push(BuiltinCandidate { has_nested: false });
                         }

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: {