From 29831ca74a2a7b05557b5b8d24dbe19c8a90d877 Mon Sep 17 00:00:00 2001
From: Taiki Endo <te316e89@gmail.com>
Date: Wed, 27 Sep 2023 23:41:17 +0900
Subject: [PATCH] Add zero_extend64_ptr helper

For 64-bit architecture's 32-bit ABI (e.g., AArch64 ILP32 ABI).
---
 build.rs             |  4 ++++
 src/gen/utils.rs     | 15 ++++++++++++++-
 src/utils.rs         | 31 +++++++++++++++++++++++++++++++
 tools/target_spec.sh | 15 ++++++++++++++-
 4 files changed, 63 insertions(+), 2 deletions(-)

diff --git a/build.rs b/build.rs
index a9ab2580..69cb36b1 100644
--- a/build.rs
+++ b/build.rs
@@ -80,6 +80,10 @@ fn main() {
     if !version.probe(64, 2022, 7, 18) {
         println!("cargo:rustc-cfg=portable_atomic_no_stronger_failure_ordering");
     }
+    // https://github.com/rust-lang/rust/pull/114790 merged in nightly-2023-08-24
+    if !version.probe(74, 2023, 8, 23) {
+        println!("cargo:rustc-cfg=portable_atomic_no_asm_maybe_uninit");
+    }
 
     // asm stabilized in Rust 1.59 (nightly-2021-12-16): https://github.com/rust-lang/rust/pull/91728
     let no_asm = !version.probe(59, 2021, 12, 15);
diff --git a/src/gen/utils.rs b/src/gen/utils.rs
index af017b90..61f6c99b 100644
--- a/src/gen/utils.rs
+++ b/src/gen/utils.rs
@@ -46,7 +46,20 @@
 macro_rules! ptr_reg {
     ($ptr:ident) => {{
         let _: *const _ = $ptr; // ensure $ptr is a pointer (*mut _ or *const _)
-        $ptr as u64
+        #[cfg(not(portable_atomic_no_asm_maybe_uninit))]
+        #[allow(clippy::ptr_as_ptr)]
+        {
+            // If we cast to u64 here, the provenance will be lost,
+            // so we convert to MaybeUninit<u64> via zero extend helper.
+            crate::utils::zero_extend64_ptr($ptr as *mut ())
+        }
+        #[cfg(portable_atomic_no_asm_maybe_uninit)]
+        {
+            // Use cast on old rustc because it does not support MaybeUninit
+            // registers. This is still permissive-provenance compatible and
+            // is sound.
+            $ptr as u64
+        }
     }};
 }
 #[cfg(not(all(
diff --git a/src/utils.rs b/src/utils.rs
index 6ece525e..4721a558 100644
--- a/src/utils.rs
+++ b/src/utils.rs
@@ -8,6 +8,10 @@ mod gen;
 
 use core::sync::atomic::Ordering;
 
+#[cfg(not(portable_atomic_no_asm_maybe_uninit))]
+#[cfg(target_pointer_width = "32")]
+use core::mem::{self, MaybeUninit};
+
 macro_rules! static_assert {
     ($cond:expr $(,)?) => {{
         let [] = [(); true as usize - $crate::utils::_assert_is_bool($cond) as usize];
@@ -751,6 +755,33 @@ pub(crate) fn upgrade_success_ordering(success: Ordering, failure: Ordering) ->
     }
 }
 
+/// Zero-extends the given 32-bit pointer to `MaybeUninit<u64>`.
+/// This is used for 64-bit architecture's 32-bit ABI (e.g., AArch64 ILP32 ABI).
+/// See ptr_reg! macro in src/gen/utils.rs for details.
+#[cfg(not(portable_atomic_no_asm_maybe_uninit))]
+#[cfg(target_pointer_width = "32")]
+#[allow(dead_code)]
+#[inline]
+pub(crate) fn zero_extend64_ptr(v: *mut ()) -> MaybeUninit<u64> {
+    // SAFETY: we can safely transmute any 64-bit value to MaybeUninit<u64>.
+    unsafe {
+        mem::transmute(ZeroExtended::<*mut (), 1> {
+            v: MaybeUninit::new(v),
+            pad: [core::ptr::null_mut(); 1],
+        })
+    }
+}
+#[cfg(not(portable_atomic_no_asm_maybe_uninit))]
+#[cfg(target_pointer_width = "32")]
+#[repr(C)]
+struct ZeroExtended<T, const N: usize> {
+    #[cfg(target_endian = "big")]
+    pad: [T; N],
+    v: MaybeUninit<T>,
+    #[cfg(target_endian = "little")]
+    pad: [T; N],
+}
+
 #[allow(dead_code)]
 #[cfg(any(
     target_arch = "aarch64",
diff --git a/tools/target_spec.sh b/tools/target_spec.sh
index 2b9a02f6..da4e2604 100755
--- a/tools/target_spec.sh
+++ b/tools/target_spec.sh
@@ -68,7 +68,20 @@ $(sed <<<"${known_64_bit_arch[*]}" -E 's/^/        target_arch = "/g; s/$/",/g')
 macro_rules! ptr_reg {
     (\$ptr:ident) => {{
         let _: *const _ = \$ptr; // ensure \$ptr is a pointer (*mut _ or *const _)
-        \$ptr as u64
+        #[cfg(not(portable_atomic_no_asm_maybe_uninit))]
+        #[allow(clippy::ptr_as_ptr)]
+        {
+            // If we cast to u64 here, the provenance will be lost,
+            // so we convert to MaybeUninit<u64> via zero extend helper.
+            crate::utils::zero_extend64_ptr(\$ptr as *mut ())
+        }
+        #[cfg(portable_atomic_no_asm_maybe_uninit)]
+        {
+            // Use cast on old rustc because it does not support MaybeUninit
+            // registers. This is still permissive-provenance compatible and
+            // is sound.
+            \$ptr as u64
+        }
     }};
 }
 #[cfg(not(all(