Implement pool for any 32/64-bit architecture that supports the corresponding atomics.

Author: reitermarkus

Cargo.toml

@@ -39,6 +39,8 @@ defmt-03 = ["dep:defmt"]
 # Enable larger MPMC sizes.
 mpmc_large = []
 
+nightly = []
+
 [dependencies]
 portable-atomic = { version = "1.0", optional = true }
 hash32 = "0.3.0"
@@ -47,7 +49,7 @@ ufmt-write = { version = "0.1", optional = true }
 defmt = { version = ">=0.2.0,<0.4", optional = true }
 
 # for the pool module
-[target.'cfg(any(target_arch = "arm", target_arch = "x86"))'.dependencies]
+[target.'cfg(any(target_arch = "arm", target_pointer_width = "32", target_pointer_width = "64"))'.dependencies]
 stable_deref_trait = { version = "1", default-features = false }
 
 [dev-dependencies]

src/lib.rs

@@ -43,11 +43,11 @@
 //!
 //! List of currently implemented data structures:
 #![cfg_attr(
-    any(arm_llsc, target_arch = "x86"),
+    any(arm_llsc, target_pointer_width = "32", target_pointer_width = "64"),
     doc = "- [`Arc`](pool::arc::Arc) -- like `std::sync::Arc` but backed by a lock-free memory pool rather than `#[global_allocator]`"
 )]
 #![cfg_attr(
-    any(arm_llsc, target_arch = "x86"),
+    any(arm_llsc, target_pointer_width = "32", target_pointer_width = "64"),
     doc = "- [`Box`](pool::boxed::Box) -- like `std::boxed::Box` but backed by a lock-free memory pool rather than `#[global_allocator]`"
 )]
 //! - [`BinaryHeap`] -- priority queue
@@ -57,7 +57,7 @@
 //! - [`IndexSet`] -- hash set
 //! - [`LinearMap`]
 #![cfg_attr(
-    any(arm_llsc, target_arch = "x86"),
+    any(arm_llsc, target_pointer_width = "32", target_pointer_width = "64"),
     doc = "- [`Object`](pool::object::Object) -- objects managed by an object pool"
 )]
 //! - [`sorted_linked_list::SortedLinkedList`]
@@ -76,6 +76,14 @@
 #![cfg_attr(docsrs, feature(doc_cfg), feature(doc_auto_cfg))]
 #![cfg_attr(not(test), no_std)]
 #![deny(missing_docs)]
+#![cfg_attr(
+    all(
+        feature = "nightly",
+        target_pointer_width = "64",
+        target_has_atomic = "128"
+    ),
+    feature(integer_atomics)
+)]
 
 pub use binary_heap::BinaryHeap;
 pub use deque::Deque;
@@ -125,7 +133,20 @@ mod defmt;
     all(not(feature = "mpmc_large"), target_has_atomic = "8")
 ))]
 pub mod mpmc;
-#[cfg(any(arm_llsc, target_arch = "x86"))]
+#[cfg(any(
+    arm_llsc,
+    all(
+        target_pointer_width = "32",
+        any(target_has_atomic = "64", feature = "portable-atomic")
+    ),
+    all(
+        target_pointer_width = "64",
+        any(
+            all(target_has_atomic = "128", feature = "nightly"),
+            feature = "portable-atomic"
+        )
+    )
+))]
 pub mod pool;
 pub mod sorted_linked_list;
 #[cfg(any(

src/pool/treiber.rs

@@ -1,6 +1,6 @@
 use core::mem::ManuallyDrop;
 
-#[cfg_attr(target_arch = "x86", path = "treiber/cas.rs")]
+#[cfg_attr(not(arm_llsc), path = "treiber/cas.rs")]
 #[cfg_attr(arm_llsc, path = "treiber/llsc.rs")]
 mod impl_;
 

src/pool/treiber/cas.rs

@@ -1,27 +1,54 @@
-use core::{
-    marker::PhantomData,
-    num::{NonZeroU32, NonZeroU64},
-    ptr::NonNull,
-    sync::atomic::{AtomicU64, Ordering},
-};
+use core::{marker::PhantomData, ptr::NonNull, sync::atomic::Ordering};
+
+#[cfg(not(feature = "portable-atomic"))]
+use core::sync::atomic;
+#[cfg(feature = "portable-atomic")]
+use portable_atomic as atomic;
 
 use super::{Node, Stack};
 
+#[cfg(target_pointer_width = "32")]
+mod types {
+    use super::atomic;
+
+    pub type Inner = u64;
+    pub type InnerAtomic = atomic::AtomicU64;
+    pub type InnerNonZero = core::num::NonZeroU64;
+
+    pub type Tag = core::num::NonZeroU32;
+    pub type Address = u32;
+}
+
+#[cfg(target_pointer_width = "64")]
+mod types {
+    use super::atomic;
+
+    pub type Inner = u128;
+    pub type InnerAtomic = atomic::AtomicU128;
+    pub type InnerNonZero = core::num::NonZeroU128;
+
+    pub type Tag = core::num::NonZeroU64;
+    pub type Address = u64;
+}
+
+use types::*;
+
 pub struct AtomicPtr<N>
 where
     N: Node,
 {
-    inner: AtomicU64,
+    inner: InnerAtomic,
     _marker: PhantomData<*mut N>,
 }
 
 impl<N> AtomicPtr<N>
 where
     N: Node,
 {
+    #[inline]
     pub const fn null() -> Self {
         Self {
-            inner: AtomicU64::new(0),
+            inner: InnerAtomic::new(0),
             _marker: PhantomData,
         }
     }
@@ -35,37 +62,35 @@ where
     ) -> Result<(), Option<NonNullPtr<N>>> {
         self.inner
             .compare_exchange_weak(
-                current
-                    .map(|pointer| pointer.into_u64())
-                    .unwrap_or_default(),
-                new.map(|pointer| pointer.into_u64()).unwrap_or_default(),
+                current.map(NonNullPtr::into_inner).unwrap_or_default(),
+                new.map(NonNullPtr::into_inner).unwrap_or_default(),
                 success,
                 failure,
             )
             .map(drop)
-            .map_err(NonNullPtr::from_u64)
+            .map_err(NonNullPtr::from_inner)
     }
 
+    #[inline]
     fn load(&self, order: Ordering) -> Option<NonNullPtr<N>> {
-        NonZeroU64::new(self.inner.load(order)).map(|inner| NonNullPtr {
+        InnerNonZero::new(self.inner.load(order)).map(|inner| NonNullPtr {
             inner,
             _marker: PhantomData,
         })
     }
 
+    #[inline]
     fn store(&self, value: Option<NonNullPtr<N>>, order: Ordering) {
-        self.inner.store(
-            value.map(|pointer| pointer.into_u64()).unwrap_or_default(),
-            order,
-        )
+        self.inner
+            .store(value.map(NonNullPtr::into_inner).unwrap_or_default(), order)
     }
 }
 
 pub struct NonNullPtr<N>
 where
     N: Node,
 {
-    inner: NonZeroU64,
+    inner: InnerNonZero,
     _marker: PhantomData<*mut N>,
 }
 
@@ -84,65 +109,66 @@ impl<N> NonNullPtr<N>
 where
     N: Node,
 {
+    #[inline]
     pub fn as_ptr(&self) -> *mut N {
         self.inner.get() as *mut N
     }
 
+    #[inline]
     pub fn from_static_mut_ref(ref_: &'static mut N) -> NonNullPtr<N> {
         let non_null = NonNull::from(ref_);
         Self::from_non_null(non_null)
     }
 
     fn from_non_null(ptr: NonNull<N>) -> Self {
-        let address = ptr.as_ptr() as u32;
+        let address = ptr.as_ptr() as Address;
         let tag = initial_tag().get();
 
-        let value = (u64::from(tag) << 32) | u64::from(address);
+        let value = (Inner::from(tag) << Address::BITS) | Inner::from(address);
 
         Self {
-            inner: unsafe { NonZeroU64::new_unchecked(value) },
+            inner: unsafe { InnerNonZero::new_unchecked(value) },
             _marker: PhantomData,
         }
     }
 
-    fn from_u64(value: u64) -> Option<Self> {
-        NonZeroU64::new(value).map(|inner| Self {
+    #[inline]
+    fn from_inner(value: Inner) -> Option<Self> {
+        InnerNonZero::new(value).map(|inner| Self {
             inner,
             _marker: PhantomData,
         })
     }
 
+    #[inline]
     fn non_null(&self) -> NonNull<N> {
-        unsafe { NonNull::new_unchecked(self.inner.get() as *mut N) }
+        unsafe { NonNull::new_unchecked(self.as_ptr()) }
     }
 
-    fn tag(&self) -> NonZeroU32 {
-        unsafe { NonZeroU32::new_unchecked((self.inner.get() >> 32) as u32) }
+    #[inline]
+    fn into_inner(self) -> Inner {
+        self.inner.get()
     }
 
-    fn into_u64(self) -> u64 {
-        self.inner.get()
+    #[inline]
+    fn tag(&self) -> Tag {
+        unsafe { Tag::new_unchecked((self.inner.get() >> Address::BITS) as Address) }
     }
 
     fn increase_tag(&mut self) {
-        let address = self.as_ptr() as u32;
+        let address = self.as_ptr() as Address;
 
-        let new_tag = self
-            .tag()
-            .get()
-            .checked_add(1)
-            .map(|val| unsafe { NonZeroU32::new_unchecked(val) })
-            .unwrap_or_else(initial_tag)
-            .get();
+        let new_tag = self.tag().checked_add(1).unwrap_or_else(initial_tag).get();
 
-        let value = (u64::from(new_tag) << 32) | u64::from(address);
+        let value = (Inner::from(new_tag) << Address::BITS) | Inner::from(address);
 
-        self.inner = unsafe { NonZeroU64::new_unchecked(value) };
+        self.inner = unsafe { InnerNonZero::new_unchecked(value) };
     }
 }
 
-fn initial_tag() -> NonZeroU32 {
-    unsafe { NonZeroU32::new_unchecked(1) }
+#[inline]
+const fn initial_tag() -> Tag {
+    Tag::MIN
 }
 
 pub unsafe fn push<N>(stack: &Stack<N>, new_top: NonNullPtr<N>)

src/pool/treiber/llsc.rs

@@ -16,6 +16,7 @@ impl<N> AtomicPtr<N>
 where
     N: Node,
 {
+    #[inline]
     pub const fn null() -> Self {
         Self {
             inner: UnsafeCell::new(None),
@@ -34,10 +35,12 @@ impl<N> NonNullPtr<N>
 where
     N: Node,
 {
+    #[inline]
     pub fn as_ptr(&self) -> *mut N {
         self.inner.as_ptr().cast()
     }
 
+    #[inline]
     pub fn from_static_mut_ref(ref_: &'static mut N) -> Self {
         Self {
             inner: NonNull::from(ref_),