use core::num::NonZeroUsize;
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub enum PartitionPlan {
Parts(NonZeroUsize),
AvailableParallelism,
ChunkSize(NonZeroUsize),
}
impl PartitionPlan {
#[inline]
#[must_use]
pub const fn parts(parts: usize) -> Self {
Self::Parts(nonzero_or_one(parts))
}
#[inline]
#[must_use]
pub const fn available_parallelism() -> Self {
Self::AvailableParallelism
}
#[inline]
#[must_use]
pub const fn chunk_size(chunk_size: usize) -> Self {
Self::ChunkSize(nonzero_or_one(chunk_size))
}
#[inline]
#[must_use]
pub fn chunk_len_for(self, len: usize) -> usize {
match self {
Self::Parts(parts) => chunk_for_parts(len, parts.get()),
Self::AvailableParallelism => {
let parts = std::thread::available_parallelism().map_or(1, NonZeroUsize::get);
chunk_for_parts(len, parts)
}
Self::ChunkSize(chunk_size) => chunk_size.get(),
}
}
#[inline]
pub(super) fn resolve(self, len: usize) -> usize {
self.chunk_len_for(len)
}
}
#[inline]
const fn nonzero_or_one(value: usize) -> NonZeroUsize {
if let Some(nz) = NonZeroUsize::new(value) {
nz
} else {
unsafe { NonZeroUsize::new_unchecked(1) }
}
}
#[inline]
fn chunk_for_parts(len: usize, parts: usize) -> usize {
crate::region::partition_count(len, parts).max(1)
}