Skip to content

Commit

Permalink
Implement pool for any 32/64-bit architecture that supports the cor…
Browse files Browse the repository at this point in the history
…responding atomics.
  • Loading branch information
reitermarkus committed Feb 19, 2024
1 parent c593fa5 commit e0f93ba
Show file tree
Hide file tree
Showing 4 changed files with 87 additions and 40 deletions.
4 changes: 3 additions & 1 deletion Cargo.toml
Original file line number Diff line number Diff line change
Expand Up @@ -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"
Expand All @@ -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]
Expand Down
29 changes: 25 additions & 4 deletions src/lib.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand All @@ -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`]
Expand All @@ -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;
Expand Down Expand Up @@ -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(
Expand Down
2 changes: 1 addition & 1 deletion src/pool/treiber.rs
Original file line number Diff line number Diff line change
@@ -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_;

Expand Down
92 changes: 58 additions & 34 deletions src/pool/treiber/cas.rs
Original file line number Diff line number Diff line change
@@ -1,17 +1,43 @@
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>,
}

Expand All @@ -21,7 +47,7 @@ where
{
pub const fn null() -> Self {
Self {
inner: AtomicU64::new(0),
inner: InnerAtomic::new(0),
_marker: PhantomData,
}
}
Expand All @@ -35,37 +61,33 @@ 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)
}

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,
})
}

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>,
}

Expand Down Expand Up @@ -94,19 +116,20 @@ where
}

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,
})
Expand All @@ -116,33 +139,34 @@ where
unsafe { NonNull::new_unchecked(self.inner.get() as *mut N) }
}

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()
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) })
.map(|val| unsafe { Tag::new_unchecked(val) })
.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) }
fn initial_tag() -> Tag {
Tag::MIN
}

pub unsafe fn push<N>(stack: &Stack<N>, new_top: NonNullPtr<N>)
Expand Down

0 comments on commit e0f93ba

Please sign in to comment.