Add u64 wait & wake variants

Author: Grimeh

src/lib.rs

@@ -1,7 +1,7 @@
 #![no_std]
 #![doc = include_str!("../README.md")]
 
-use core::sync::atomic::{AtomicPtr, AtomicU32};
+use core::sync::atomic::{AtomicPtr, AtomicU32, AtomicU64};
 
 #[cfg(any(target_os = "linux", target_os = "android"))]
 #[path = "linux.rs"]
@@ -28,6 +28,15 @@ pub fn wait(atomic: &AtomicU32, value: u32) {
     platform::wait(atomic, value)
 }
 
+/// If the value is `value`, wait until woken up.
+///
+/// This function might also return spuriously,
+/// without a corresponding wake operation.
+#[inline]
+pub fn wait_u64(atomic: &AtomicU64, value: u64) {
+    platform::wait_u64(atomic, value)
+}
+
 /// If the value is `value`, wait until woken up.
 ///
 /// This function might also return spuriously,
@@ -44,6 +53,14 @@ pub fn wake_one(atomic: *const AtomicU32) {
     platform::wake_one(atomic);
 }
 
+/// Wake one thread that is waiting on this atomic.
+///
+/// It's okay if the pointer dangles or is null.
+#[inline]
+pub fn wake_one_u64(atomic: *const AtomicU64) {
+    platform::wake_one_u64(atomic);
+}
+
 /// Wake one thread that is waiting on this atomic.
 ///
 /// It's okay if the pointer dangles or is null.
@@ -60,6 +77,14 @@ pub fn wake_all(atomic: *const AtomicU32) {
     platform::wake_all(atomic);
 }
 
+/// Wake all threads that are waiting on this atomic.
+///
+/// It's okay if the pointer dangles or is null.
+#[inline]
+pub fn wake_all_u64(atomic: *const AtomicU64) {
+    platform::wake_all_u64(atomic);
+}
+
 /// Wake all threads that are waiting on this atomic.
 ///
 /// It's okay if the pointer dangles or is null.

src/windows.rs

@@ -1,4 +1,4 @@
-use core::sync::atomic::{AtomicPtr, AtomicU32};
+use core::sync::atomic::{AtomicPtr, AtomicU32, AtomicU64};
 use windows_sys::Win32::System::{
     Threading::{WaitOnAddress, WakeByAddressAll, WakeByAddressSingle},
     WindowsProgramming::INFINITE,
@@ -11,6 +11,14 @@ pub fn wait(a: &AtomicU32, expected: u32) {
     unsafe { WaitOnAddress(ptr.cast(), expected_ptr.cast(), 4, INFINITE) };
 }
 
+#[inline]
+pub fn wait_u64(a: &AtomicU64, expected: u64) {
+    let ptr: *const AtomicU64 = a;
+    let expected_ptr: *const u64 = &expected;
+    unsafe { WaitOnAddress(ptr.cast(), expected_ptr.cast(), 4, INFINITE) };
+}
+
+#[inline]
 pub fn wait_ptr<T>(a: *const AtomicPtr<T>, expected: *const T) {
     let expected_ptr: *const *const T = &expected as _;
     unsafe { WaitOnAddress(a.cast(), expected_ptr.cast(), 8, INFINITE )};
@@ -21,6 +29,11 @@ pub fn wake_one(ptr: *const AtomicU32) {
     unsafe { WakeByAddressSingle(ptr.cast()) };
 }
 
+#[inline]
+pub fn wake_one_u64(ptr: *const AtomicU64) {
+    unsafe { WakeByAddressSingle(ptr.cast()) };
+}
+
 #[inline]
 pub fn wake_one_ptr<T>(ptr: *const AtomicPtr<T>) {
     unsafe { WakeByAddressSingle(ptr.cast()) };
@@ -31,6 +44,11 @@ pub fn wake_all(ptr: *const AtomicU32) {
     unsafe { WakeByAddressAll(ptr.cast()) };
 }
 
+#[inline]
+pub fn wake_all_u64(ptr: *const AtomicU64) {
+    unsafe { WakeByAddressAll(ptr.cast()) };
+}
+
 #[inline]
 pub fn wake_all_ptr<T>(ptr: *const AtomicPtr<T>) {
     unsafe { WakeByAddressAll(ptr.cast()) };

tests/test.rs

@@ -1,23 +1,37 @@
-use atomic_wait::{wait, wake_all, wake_one};
+use atomic_wait::{wait, wait_u64, wake_all, wake_all_u64, wake_one, wake_one_u64};
 use std::{
     sync::atomic::{AtomicU32, Ordering::Relaxed},
     thread::sleep,
     time::{Duration, Instant},
 };
+use std::sync::atomic::AtomicU64;
 
 #[test]
 fn wake_null() {
     wake_one(std::ptr::null::<AtomicU32>());
     wake_all(std::ptr::null::<AtomicU32>());
 }
 
+#[test]
+fn wake_null_u64() {
+    wake_one_u64(std::ptr::null::<AtomicU64>());
+    wake_all_u64(std::ptr::null::<AtomicU64>());
+}
+
 #[test]
 fn wake_nothing() {
     let a = AtomicU32::new(0);
     wake_one(&a);
     wake_all(&a);
 }
 
+#[test]
+fn wake_nothing_u64() {
+    let a = AtomicU64::new(0);
+    wake_one_u64(&a);
+    wake_all_u64(&a);
+}
+
 #[test]
 fn wait_unexpected() {
     let t = Instant::now();
@@ -26,6 +40,14 @@ fn wait_unexpected() {
     assert!(t.elapsed().as_millis() < 100);
 }
 
+#[test]
+fn wait_unexpected_u64() {
+    let t = Instant::now();
+    let a = AtomicU64::new(0);
+    wait_u64(&a, 1);
+    assert!(t.elapsed().as_millis() < 100);
+}
+
 #[test]
 fn wait_wake() {
     let t = Instant::now();
@@ -43,3 +65,21 @@ fn wait_wake() {
         assert!((90..400).contains(&t.elapsed().as_millis()));
     });
 }
+
+#[test]
+fn wait_wake_u64() {
+    let t = Instant::now();
+    let a = AtomicU64::new(0);
+    std::thread::scope(|s| {
+        s.spawn(|| {
+            sleep(Duration::from_millis(100));
+            a.store(1, Relaxed);
+            wake_one_u64(&a);
+        });
+        while a.load(Relaxed) == 0 {
+            wait_u64(&a, 0);
+        }
+        assert_eq!(a.load(Relaxed), 1);
+        assert!((90..400).contains(&t.elapsed().as_millis()));
+    });
+}