Rollup merge of #121622 - dtolnay:wake, r=cuviper

Preserve same vtable pointer when cloning raw waker, to fix Waker::will_wake

Fixes #121600.

As `@jkarneges` identified in https://github.com/rust-lang/rust/issues/121600#issuecomment-1963041051, the issue is two different const promotions produce two statics at different addresses, which may or may not later be deduplicated by the linker (in this case not).

Prior to #119863, the content of the statics was compared, and they were equal. After, the address of the statics are compared and they are not equal.

It is documented that `will_wake` _"works on a best-effort basis, and may return false even when the Wakers would awaken the same task"_ so this PR fixes a quality-of-implementation issue, not a correctness issue.

library/alloc/src/task.rs

@@ -136,6 +136,15 @@ impl<W: Wake + Send + Sync + 'static> From<Arc<W>> for RawWaker {
 #[inline(always)]
 fn raw_waker<W: Wake + Send + Sync + 'static>(waker: Arc<W>) -> RawWaker {
     // Increment the reference count of the arc to clone it.
+    //
+    // The #[inline(always)] is to ensure that raw_waker and clone_waker are
+    // always generated in the same code generation unit as one another, and
+    // therefore that the structurally identical const-promoted RawWakerVTable
+    // within both functions is deduplicated at LLVM IR code generation time.
+    // This allows optimizing Waker::will_wake to a single pointer comparison of
+    // the vtable pointers, rather than comparing all four function pointers
+    // within the vtables.
+    #[inline(always)]
     unsafe fn clone_waker<W: Wake + Send + Sync + 'static>(waker: *const ()) -> RawWaker {
         unsafe { Arc::increment_strong_count(waker as *const W) };
         RawWaker::new(
@@ -304,6 +313,10 @@ impl<W: LocalWake + 'static> From<Rc<W>> for RawWaker {
 #[inline(always)]
 fn local_raw_waker<W: LocalWake + 'static>(waker: Rc<W>) -> RawWaker {
     // Increment the reference count of the Rc to clone it.
+    //
+    // Refer to the comment on raw_waker's clone_waker regarding why this is
+    // always inline.
+    #[inline(always)]
     unsafe fn clone_waker<W: LocalWake + 'static>(waker: *const ()) -> RawWaker {
         unsafe { Rc::increment_strong_count(waker as *const W) };
         RawWaker::new(

library/alloc/tests/lib.rs

@@ -41,6 +41,7 @@
 #![feature(thin_box)]
 #![feature(strict_provenance)]
 #![feature(drain_keep_rest)]
+#![feature(local_waker)]
 #![allow(internal_features)]
 #![deny(fuzzy_provenance_casts)]
 #![deny(unsafe_op_in_unsafe_fn)]
@@ -62,6 +63,7 @@ mod rc;
 mod slice;
 mod str;
 mod string;
+mod task;
 mod thin_box;
 mod vec;
 mod vec_deque;

library/alloc/tests/task.rs

@@ -0,0 +1,34 @@
+use alloc::rc::Rc;
+use alloc::sync::Arc;
+use alloc::task::{LocalWake, Wake};
+use core::task::{LocalWaker, Waker};
+
+#[test]
+fn test_waker_will_wake_clone() {
+    struct NoopWaker;
+
+    impl Wake for NoopWaker {
+        fn wake(self: Arc<Self>) {}
+    }
+
+    let waker = Waker::from(Arc::new(NoopWaker));
+    let clone = waker.clone();
+
+    assert!(waker.will_wake(&clone));
+    assert!(clone.will_wake(&waker));
+}
+
+#[test]
+fn test_local_waker_will_wake_clone() {
+    struct NoopWaker;
+
+    impl LocalWake for NoopWaker {
+        fn wake(self: Rc<Self>) {}
+    }
+
+    let waker = LocalWaker::from(Rc::new(NoopWaker));
+    let clone = waker.clone();
+
+    assert!(waker.will_wake(&clone));
+    assert!(clone.will_wake(&waker));
+}