use super::{BStackAllocator, BStackOwnedSlice, BStackSlice};
use crate::BStack;
use std::cmp::Ordering;
use std::fmt;
use std::hash::{Hash, Hasher};
use std::io;
pub struct BStackChunk<'a> {
aligned: BStackSlice<'a>,
chunk_len: u64,
}
impl<'a> Clone for BStackChunk<'a> {
#[inline]
fn clone(&self) -> Self {
BStackChunk {
aligned: self.aligned.clone(),
chunk_len: self.chunk_len,
}
}
}
impl<'a> fmt::Debug for BStackChunk<'a> {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
f.debug_struct("BStackChunk")
.field("chunk_len", &self.chunk_len)
.field("chunk_count", &self.chunk_count())
.finish_non_exhaustive()
}
}
impl<'a> PartialEq for BStackChunk<'a> {
#[inline]
fn eq(&self, other: &Self) -> bool {
self.chunk_len == other.chunk_len && self.aligned == other.aligned
}
}
impl<'a> Eq for BStackChunk<'a> {}
impl<'a> Hash for BStackChunk<'a> {
#[inline]
fn hash<H: Hasher>(&self, state: &mut H) {
self.chunk_len.hash(state);
self.aligned.hash(state);
}
}
impl<'a> PartialOrd for BStackChunk<'a> {
#[inline]
fn partial_cmp(&self, other: &Self) -> Option<Ordering> {
Some(self.cmp(other))
}
}
impl<'a> Ord for BStackChunk<'a> {
#[inline]
fn cmp(&self, other: &Self) -> Ordering {
self.chunk_len
.cmp(&other.chunk_len)
.then_with(|| self.aligned.cmp(&other.aligned))
}
}
impl<'a> BStackChunk<'a> {
#[inline]
pub unsafe fn from_raw_parts(stack: &'a BStack, offset: u64, len: u64, chunk_len: u64) -> Self {
BStackChunk {
aligned: unsafe { BStackSlice::from_raw_parts(stack, offset, len) },
chunk_len,
}
}
#[inline]
pub unsafe fn from_raw_slice(aligned: BStackSlice<'a>, chunk_len: u64) -> Self {
BStackChunk { aligned, chunk_len }
}
#[inline]
pub fn from_slice(aligned: BStackSlice<'a>, chunk_len: u64) -> Option<Self> {
if chunk_len == 0 || !aligned.len().is_multiple_of(chunk_len) {
return None;
}
Some(BStackChunk { aligned, chunk_len })
}
#[inline]
pub fn chunk_len(&self) -> u64 {
self.chunk_len
}
#[inline]
pub fn chunk_count(&self) -> u64 {
self.aligned.len() / self.chunk_len
}
#[inline]
pub fn len(&self) -> u64 {
self.aligned.len()
}
#[inline]
pub fn is_empty(&self) -> bool {
self.aligned.is_empty()
}
#[inline]
pub fn same_stride(&self, other: &Self) -> bool {
self.chunk_len == other.chunk_len
}
#[inline]
pub fn same_phase(&self, other: &Self) -> bool {
self.same_stride(other)
&& self.aligned.start() % self.chunk_len == other.aligned.start() % self.chunk_len
}
#[inline]
pub fn adjacent_to(&self, other: &Self) -> bool {
self.aligned.adjacent_to(&other.aligned) && self.same_phase(other)
}
#[inline]
pub fn overlaps(&self, other: &Self) -> bool {
self.aligned.overlaps(&other.aligned) && self.same_phase(other)
}
pub fn merge(&self, other: &Self) -> Option<Self> {
if !self.same_stride(other) {
return None;
}
if self.is_empty() {
return Some(other.clone());
}
if other.is_empty() {
return Some(self.clone());
}
if !self.overlaps(other) {
return None;
}
let aligned = self.aligned.merge(&other.aligned)?;
Some(BStackChunk {
aligned,
chunk_len: self.chunk_len,
})
}
pub fn merge_adjacent(&self, other: &Self) -> Option<Self> {
if !self.same_stride(other) {
return None;
}
let aligned = self.aligned.merge_adjacent(&other.aligned)?;
Some(BStackChunk {
aligned,
chunk_len: self.chunk_len,
})
}
#[inline]
pub fn as_slice(&self) -> BStackSlice<'a> {
self.aligned.clone()
}
#[inline]
pub fn into_slice(self) -> BStackSlice<'a> {
self.aligned
}
#[inline]
pub fn with_stride(self, new_stride: u64) -> (BStackChunk<'a>, BStackSlice<'a>) {
self.aligned.chunks(new_stride)
}
#[inline]
pub fn stack(&self) -> &'a BStack {
self.aligned.stack()
}
#[inline]
pub fn get(&self, index: u64) -> Option<BStackSlice<'a>> {
if index >= self.chunk_count() {
return None;
}
let start = index * self.chunk_len;
Some(self.aligned.subslice(start, start + self.chunk_len))
}
#[inline]
pub fn iter(&self) -> BStackChunkIter<'a> {
BStackChunkIter {
remaining: self.aligned.clone(),
chunk_len: self.chunk_len,
}
}
#[cfg(all(feature = "set", feature = "atomic"))]
pub fn sort_by(&mut self, mut cmp: impl FnMut(&[u8], &[u8]) -> Ordering) -> io::Result<()> {
let chunk_len = self.chunk_len as usize;
let start = self.aligned.start();
let end = self.aligned.end();
self.aligned.stack().process(start, end, |buf| {
sort_chunks_by(buf, chunk_len, &mut cmp);
})
}
#[cfg(all(feature = "set", feature = "atomic"))]
pub fn sort_by_key<K: Ord>(&mut self, mut key: impl FnMut(&[u8]) -> K) -> io::Result<()> {
let chunk_len = self.chunk_len as usize;
let start = self.aligned.start();
let end = self.aligned.end();
self.aligned.stack().process(start, end, |buf| {
let keys: Vec<K> = buf.chunks_exact(chunk_len).map(&mut key).collect();
let mut order: Vec<usize> = (0..keys.len()).collect();
order.sort_by(|&i, &j| keys[i].cmp(&keys[j]));
apply_chunk_permutation(buf, chunk_len, &order);
})
}
pub fn binary_search_by(
&self,
mut cmp: impl FnMut(&[u8]) -> Ordering,
) -> io::Result<Result<u64, u64>> {
let chunk_len = self.chunk_len as usize;
let mut inline = [0u8; INLINE_SCRATCH_LEN];
let mut heap;
let buf: &mut [u8] = if chunk_len <= INLINE_SCRATCH_LEN {
&mut inline[..chunk_len]
} else {
heap = vec![0u8; chunk_len];
&mut heap[..]
};
let mut size = self.chunk_count();
let mut left = 0u64;
while size > 0 {
let half = size / 2;
let mid = left + half;
let chunk = self
.get(mid)
.expect("binary_search_by: mid computed within bounds");
chunk.read_into(buf)?;
match cmp(buf) {
Ordering::Less => {
left = mid + 1;
size -= half + 1;
}
Ordering::Greater => size = half,
Ordering::Equal => return Ok(Ok(mid)),
}
}
Ok(Err(left))
}
#[inline]
pub fn binary_search_by_key<K: Ord>(
&self,
target: &K,
mut key: impl FnMut(&[u8]) -> K,
) -> io::Result<Result<u64, u64>> {
self.binary_search_by(|bytes| key(bytes).cmp(target))
}
#[cfg(all(feature = "set", feature = "atomic"))]
pub fn select_nth_by(
&mut self,
n: u64,
mut cmp: impl FnMut(&[u8], &[u8]) -> Ordering,
) -> io::Result<()> {
assert!(
n < self.chunk_count(),
"select_nth_by: n must be < chunk_count"
);
let chunk_len = self.chunk_len as usize;
let start = self.aligned.start();
let end = self.aligned.end();
self.aligned.stack().process(start, end, |buf| {
let count = buf.len() / chunk_len;
let mut order: Vec<usize> = (0..count).collect();
order.select_nth_unstable_by(n as usize, |&i, &j| {
let a = &buf[i * chunk_len..(i + 1) * chunk_len];
let b = &buf[j * chunk_len..(j + 1) * chunk_len];
cmp(a, b)
});
apply_chunk_permutation(buf, chunk_len, &order);
})
}
#[cfg(all(feature = "set", feature = "atomic"))]
pub fn select_nth_by_key<K: Ord>(
&mut self,
n: u64,
mut key: impl FnMut(&[u8]) -> K,
) -> io::Result<()> {
assert!(
n < self.chunk_count(),
"select_nth_by_key: n must be < chunk_count"
);
let chunk_len = self.chunk_len as usize;
let start = self.aligned.start();
let end = self.aligned.end();
self.aligned.stack().process(start, end, |buf| {
let keys: Vec<K> = buf.chunks_exact(chunk_len).map(&mut key).collect();
let mut order: Vec<usize> = (0..keys.len()).collect();
order.select_nth_unstable_by(n as usize, |&i, &j| keys[i].cmp(&keys[j]));
apply_chunk_permutation(buf, chunk_len, &order);
})
}
}
impl<'a> IntoIterator for BStackChunk<'a> {
type Item = BStackSlice<'a>;
type IntoIter = BStackChunkIter<'a>;
#[inline]
fn into_iter(self) -> Self::IntoIter {
BStackChunkIter {
remaining: self.aligned,
chunk_len: self.chunk_len,
}
}
}
impl<'a> IntoIterator for &BStackChunk<'a> {
type Item = BStackSlice<'a>;
type IntoIter = BStackChunkIter<'a>;
#[inline]
fn into_iter(self) -> Self::IntoIter {
self.iter()
}
}
#[cfg(all(feature = "set", feature = "atomic"))]
fn sort_chunks_by(buf: &mut [u8], chunk_len: usize, cmp: &mut dyn FnMut(&[u8], &[u8]) -> Ordering) {
let count = buf.len() / chunk_len;
if count <= 1 {
return;
}
let mut order: Vec<usize> = (0..count).collect();
order.sort_by(|&i, &j| {
let a = &buf[i * chunk_len..(i + 1) * chunk_len];
let b = &buf[j * chunk_len..(j + 1) * chunk_len];
cmp(a, b)
});
apply_chunk_permutation(buf, chunk_len, &order);
}
const INLINE_SCRATCH_LEN: usize = 128;
#[cfg(all(feature = "set", feature = "atomic"))]
fn apply_chunk_permutation(buf: &mut [u8], chunk_len: usize, order: &[usize]) {
let mut inline = [0u8; INLINE_SCRATCH_LEN];
let mut heap;
let temp: &mut [u8] = if chunk_len <= INLINE_SCRATCH_LEN {
&mut inline[..chunk_len]
} else {
heap = vec![0u8; chunk_len];
&mut heap[..]
};
let mut visited = vec![false; order.len()];
for start in 0..order.len() {
if visited[start] || order[start] == start {
visited[start] = true;
continue;
}
temp.copy_from_slice(&buf[start * chunk_len..(start + 1) * chunk_len]);
let mut cur = start;
while order[cur] != start {
let src = order[cur];
buf.copy_within(src * chunk_len..(src + 1) * chunk_len, cur * chunk_len);
visited[cur] = true;
cur = src;
}
buf[cur * chunk_len..(cur + 1) * chunk_len].copy_from_slice(temp);
visited[cur] = true;
}
}
impl<'a> BStackSlice<'a> {
pub fn chunks(&self, chunk_len: u64) -> (BStackChunk<'a>, BStackSlice<'a>) {
assert!(chunk_len > 0, "chunks: chunk_len must be nonzero");
let len = self.len();
let aligned_len = (len / chunk_len) * chunk_len;
let chunk = BStackChunk {
aligned: self.subslice(0, aligned_len),
chunk_len,
};
(chunk, self.subslice(aligned_len, len))
}
pub fn rchunks(&self, chunk_len: u64) -> (BStackChunk<'a>, BStackSlice<'a>) {
assert!(chunk_len > 0, "rchunks: chunk_len must be nonzero");
let len = self.len();
let rem_len = len % chunk_len;
let chunk = BStackChunk {
aligned: self.subslice(rem_len, len),
chunk_len,
};
(chunk, self.subslice(0, rem_len))
}
}
impl<'a, A: BStackAllocator> BStackOwnedSlice<'a, A> {
#[inline]
pub fn chunks<'s>(&'s self, chunk_len: u64) -> (BStackChunk<'s>, BStackSlice<'s>) {
self.as_slice().chunks(chunk_len)
}
#[inline]
pub fn rchunks<'s>(&'s self, chunk_len: u64) -> (BStackChunk<'s>, BStackSlice<'s>) {
self.as_slice().rchunks(chunk_len)
}
}
pub struct BStackChunkIter<'a> {
remaining: BStackSlice<'a>,
chunk_len: u64,
}
impl<'a> Clone for BStackChunkIter<'a> {
#[inline]
fn clone(&self) -> Self {
BStackChunkIter {
remaining: self.remaining.clone(),
chunk_len: self.chunk_len,
}
}
}
impl<'a> fmt::Debug for BStackChunkIter<'a> {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
f.debug_struct("BStackChunkIter")
.field("remaining_len", &self.remaining.len())
.field("chunk_len", &self.chunk_len)
.finish_non_exhaustive()
}
}
impl<'a> Iterator for BStackChunkIter<'a> {
type Item = BStackSlice<'a>;
fn next(&mut self) -> Option<Self::Item> {
if self.remaining.len() < self.chunk_len {
return None;
}
let head = self.remaining.subslice(0, self.chunk_len);
self.remaining = self
.remaining
.subslice(self.chunk_len, self.remaining.len());
Some(head)
}
#[inline]
fn size_hint(&self) -> (usize, Option<usize>) {
let n = (self.remaining.len() / self.chunk_len).min(usize::MAX as u64) as usize;
(n, Some(n))
}
}
impl<'a> DoubleEndedIterator for BStackChunkIter<'a> {
fn next_back(&mut self) -> Option<Self::Item> {
if self.remaining.len() < self.chunk_len {
return None;
}
let len = self.remaining.len();
let tail = self.remaining.subslice(len - self.chunk_len, len);
self.remaining = self.remaining.subslice(0, len - self.chunk_len);
Some(tail)
}
}
impl<'a> ExactSizeIterator for BStackChunkIter<'a> {}
impl<'a> std::iter::FusedIterator for BStackChunkIter<'a> {}