From 8f79fc24e36e1a7540a0af6d44a27d11366fe294 Mon Sep 17 00:00:00 2001
From: Michael Goulet <michael@errs.io>
Date: Sat, 19 Nov 2022 03:06:21 +0000
Subject: [PATCH] Properly handle `Pin<&mut dyn* Trait>` receiver in codegen

---
 compiler/rustc_codegen_ssa/src/mir/block.rs   | 30 ++++++++---
 src/test/ui/dyn-star/dispatch-on-pin-mut.rs   | 51 +++++++++++++++++++
 .../ui/dyn-star/dispatch-on-pin-mut.stderr    | 11 ++++
 3 files changed, 86 insertions(+), 6 deletions(-)
 create mode 100644 src/test/ui/dyn-star/dispatch-on-pin-mut.rs
 create mode 100644 src/test/ui/dyn-star/dispatch-on-pin-mut.stderr

diff --git a/compiler/rustc_codegen_ssa/src/mir/block.rs b/compiler/rustc_codegen_ssa/src/mir/block.rs
index 7822d924c01..03d833fbba8 100644
--- a/compiler/rustc_codegen_ssa/src/mir/block.rs
+++ b/compiler/rustc_codegen_ssa/src/mir/block.rs
@@ -938,7 +938,9 @@ impl<'a, 'tcx, Bx: BuilderMethods<'a, 'tcx>> FunctionCx<'a, 'tcx, Bx> {
                         // that is understood elsewhere in the compiler as a method on
                         // `dyn Trait`.
                         // To get a `*mut RcBox<Self>`, we just keep unwrapping newtypes until
-                        // we get a value of a built-in pointer type
+                        // we get a value of a built-in pointer type.
+                        //
+                        // This is also relevant for `Pin<&mut Self>`, where we need to peel the `Pin`.
                         'descend_newtypes: while !op.layout.ty.is_unsafe_ptr()
                             && !op.layout.ty.is_region_ptr()
                         {
@@ -980,13 +982,29 @@ impl<'a, 'tcx, Bx: BuilderMethods<'a, 'tcx>> FunctionCx<'a, 'tcx, Bx> {
                         continue;
                     }
                     Immediate(_) => {
-                        let ty::Ref(_, ty, _) = op.layout.ty.kind() else {
-                            span_bug!(span, "can't codegen a virtual call on {:#?}", op);
-                        };
-                        if !ty.is_dyn_star() {
+                        // See comment above explaining why we peel these newtypes
+                        'descend_newtypes: while !op.layout.ty.is_unsafe_ptr()
+                            && !op.layout.ty.is_region_ptr()
+                        {
+                            for i in 0..op.layout.fields.count() {
+                                let field = op.extract_field(bx, i);
+                                if !field.layout.is_zst() {
+                                    // we found the one non-zero-sized field that is allowed
+                                    // now find *its* non-zero-sized field, or stop if it's a
+                                    // pointer
+                                    op = field;
+                                    continue 'descend_newtypes;
+                                }
+                            }
+
+                            span_bug!(span, "receiver has no non-zero-sized fields {:?}", op);
+                        }
+
+                        // Make sure that we've actually unwrapped the rcvr down
+                        // to a pointer or ref to `dyn* Trait`.
+                        if !op.layout.ty.builtin_deref(true).unwrap().ty.is_dyn_star() {
                             span_bug!(span, "can't codegen a virtual call on {:#?}", op);
                         }
-                        // FIXME(dyn-star): Make sure this is done on a &dyn* receiver
                         let place = op.deref(bx.cx());
                         let data_ptr = place.project_field(bx, 0);
                         let meta_ptr = place.project_field(bx, 1);
diff --git a/src/test/ui/dyn-star/dispatch-on-pin-mut.rs b/src/test/ui/dyn-star/dispatch-on-pin-mut.rs
new file mode 100644
index 00000000000..8eedff87d0d
--- /dev/null
+++ b/src/test/ui/dyn-star/dispatch-on-pin-mut.rs
@@ -0,0 +1,51 @@
+// build-pass
+// edition:2021
+
+#![feature(dyn_star)]
+//~^ WARN the feature `dyn_star` is incomplete and may not be safe to use and/or cause compiler crashes
+
+use std::future::Future;
+
+async fn foo(f: dyn* Future<Output = i32>) {
+    println!("value: {}", f.await);
+}
+
+async fn async_main() {
+    foo(Box::pin(async { 1 })).await
+}
+
+// ------------------------------------------------------------------------- //
+// Implementation Details Below...
+
+use std::pin::Pin;
+use std::task::*;
+
+pub fn noop_waker() -> Waker {
+    let raw = RawWaker::new(std::ptr::null(), &NOOP_WAKER_VTABLE);
+
+    // SAFETY: the contracts for RawWaker and RawWakerVTable are upheld
+    unsafe { Waker::from_raw(raw) }
+}
+
+const NOOP_WAKER_VTABLE: RawWakerVTable = RawWakerVTable::new(noop_clone, noop, noop, noop);
+
+unsafe fn noop_clone(_p: *const ()) -> RawWaker {
+    RawWaker::new(std::ptr::null(), &NOOP_WAKER_VTABLE)
+}
+
+unsafe fn noop(_p: *const ()) {}
+
+fn main() {
+    let mut fut = async_main();
+
+    // Poll loop, just to test the future...
+    let waker = noop_waker();
+    let ctx = &mut Context::from_waker(&waker);
+
+    loop {
+        match unsafe { Pin::new_unchecked(&mut fut).poll(ctx) } {
+            Poll::Pending => {}
+            Poll::Ready(()) => break,
+        }
+    }
+}
diff --git a/src/test/ui/dyn-star/dispatch-on-pin-mut.stderr b/src/test/ui/dyn-star/dispatch-on-pin-mut.stderr
new file mode 100644
index 00000000000..a73e3400922
--- /dev/null
+++ b/src/test/ui/dyn-star/dispatch-on-pin-mut.stderr
@@ -0,0 +1,11 @@
+warning: the feature `dyn_star` is incomplete and may not be safe to use and/or cause compiler crashes
+  --> $DIR/dispatch-on-pin-mut.rs:4:12
+   |
+LL | #![feature(dyn_star)]
+   |            ^^^^^^^^
+   |
+   = note: see issue #91611 <https://github.com/rust-lang/rust/issues/91611> for more information
+   = note: `#[warn(incomplete_features)]` on by default
+
+warning: 1 warning emitted
+