diff --git a/library/alloc/src/task.rs b/library/alloc/src/task.rs index b214f4946b1b5..b40768a52b6b4 100644 --- a/library/alloc/src/task.rs +++ b/library/alloc/src/task.rs @@ -136,6 +136,15 @@ impl From> for RawWaker { #[inline(always)] fn raw_waker(waker: Arc) -> 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(waker: *const ()) -> RawWaker { unsafe { Arc::increment_strong_count(waker as *const W) }; RawWaker::new( @@ -304,6 +313,10 @@ impl From> for RawWaker { #[inline(always)] fn local_raw_waker(waker: Rc) -> 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(waker: *const ()) -> RawWaker { unsafe { Rc::increment_strong_count(waker as *const W) }; RawWaker::new( diff --git a/library/alloc/tests/lib.rs b/library/alloc/tests/lib.rs index c4e89a58a05ac..ed928994ad697 100644 --- a/library/alloc/tests/lib.rs +++ b/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; diff --git a/library/alloc/tests/task.rs b/library/alloc/tests/task.rs new file mode 100644 index 0000000000000..0f8d9218980e9 --- /dev/null +++ b/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) {} + } + + 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) {} + } + + let waker = LocalWaker::from(Rc::new(NoopWaker)); + let clone = waker.clone(); + + assert!(waker.will_wake(&clone)); + assert!(clone.will_wake(&waker)); +}