use std::collections::VecDeque;
use std::ops::Range;
use std::process::abort;
use super::PrefixLength;
use super::bytes::{Bytes, Retained};
pub(super) mod ring;
pub(super) mod rolling;
mod private {
pub trait Sealed {}
}
pub trait AdvanceSegment: AsRef<[u8]> + private::Sealed {
fn try_advance(&mut self, amount: usize) -> bool;
}
#[derive(Debug)]
pub struct SegmentQueue<T> {
segments: VecDeque<T>,
len: usize,
}
#[derive(Debug)]
pub struct RetainedSegmentQueue<T> {
queue: SegmentQueue<T>,
front_offset: usize,
}
impl<T> SegmentQueue<T> {
#[must_use]
pub const fn new() -> Self {
Self {
segments: VecDeque::new(),
len: 0,
}
}
pub const fn len(&self) -> usize {
self.len
}
pub fn is_empty(&self) -> bool {
self.len == 0
}
pub fn segment_count(&self) -> usize {
self.segments.len()
}
pub fn front(&self) -> Option<&T> {
self.segments.front()
}
pub fn back(&self) -> Option<&T> {
self.segments.back()
}
pub fn clear(&mut self) {
self.segments.clear();
self.len = 0;
}
pub fn clear_from(&mut self, front_offset: &mut usize) {
self.clear();
*front_offset = 0;
}
}
impl<T: AsRef<[u8]>> SegmentQueue<T> {
pub fn try_push_back(&mut self, segment: T) -> Result<(), T> {
let segment_len = segment.as_ref().len();
if segment_len == 0 {
return Ok(());
}
let Some(len) = self.len.checked_add(segment_len) else {
return Err(segment);
};
self.segments.push_back(segment);
self.len = len;
Ok(())
}
pub fn try_mutate_back<R>(
&mut self,
additional: usize,
mutate: impl FnOnce(&mut T) -> R,
) -> Option<R> {
let total_ceiling = self.len.checked_add(additional)?;
let segment = self.segments.back_mut()?;
let before = segment.as_ref().len();
let segment_ceiling = before.checked_add(additional)?;
let mutation = SegmentMutation {
segment,
total: &mut self.len,
total_ceiling,
segment_ceiling,
before,
};
Some(mutate(mutation.segment))
}
pub fn range_available(&self, front_offset: usize, offset: usize, len: usize) -> bool {
if !self.front_offset_valid(front_offset) {
return false;
}
offset.checked_add(len).is_some_and(|end| end <= self.len)
}
pub fn for_each_range(
&self,
front_offset: usize,
offset: usize,
len: usize,
mut visit: impl FnMut(&[u8]),
) -> bool {
if !self.range_available(front_offset, offset, len) {
return false;
}
if len == 0 {
return true;
}
let Some(mut skip) = front_offset.checked_add(offset) else {
return false;
};
let mut remaining = len;
for segment in &self.segments {
let bytes = segment.as_ref();
if skip >= bytes.len() {
skip -= bytes.len();
continue;
}
let bytes = &bytes[skip..];
let take = remaining.min(bytes.len());
visit(&bytes[..take]);
remaining -= take;
skip = 0;
if remaining == 0 {
return true;
}
}
false
}
pub fn copy_range_into(&self, front_offset: usize, offset: usize, output: &mut [u8]) -> bool {
let mut written = 0;
let copied = self.for_each_range(front_offset, offset, output.len(), |bytes| {
let end = written + bytes.len();
output[written..end].copy_from_slice(bytes);
written = end;
});
debug_assert!(!copied || written == output.len());
copied
}
pub fn extend_range(
&self,
front_offset: usize,
offset: usize,
len: usize,
output: &mut Vec<u8>,
) -> bool {
self.for_each_range(front_offset, offset, len, |bytes| {
output.extend_from_slice(bytes);
})
}
pub fn contiguous_segment(
&self,
front_offset: usize,
offset: usize,
len: usize,
) -> Option<(&T, Range<usize>)> {
if len == 0 || !self.range_available(front_offset, offset, len) {
return None;
}
let mut skip = front_offset.checked_add(offset)?;
for segment in &self.segments {
let segment_len = segment.as_ref().len();
if skip >= segment_len {
skip -= segment_len;
continue;
}
return (len <= segment_len - skip).then_some((segment, skip..skip + len));
}
None
}
pub fn try_consume_front_from(
&mut self,
front_offset: &mut usize,
amount: usize,
removed: impl FnMut(T),
) -> bool {
if !self.range_available(*front_offset, 0, amount) {
return false;
}
unsafe { self.consume_front_valid_from(front_offset, amount, removed) };
true
}
unsafe fn consume_front_valid_from(
&mut self,
front_offset: &mut usize,
mut amount: usize,
mut removed: impl FnMut(T),
) {
debug_assert!(self.range_available(*front_offset, 0, amount));
while amount != 0 {
let segment = unsafe { self.segments.front().unwrap_unchecked() };
let available = segment.as_ref().len() - *front_offset;
if amount < available {
*front_offset += amount;
self.len -= amount;
return;
}
amount -= available;
self.len -= available;
let segment = unsafe { self.segments.pop_front().unwrap_unchecked() };
*front_offset = 0;
removed(segment);
}
}
fn front_offset_valid(&self, front_offset: usize) -> bool {
match self.segments.front() {
Some(segment) => front_offset < segment.as_ref().len(),
None => front_offset == 0 && self.len == 0,
}
}
}
impl<T: AdvanceSegment> SegmentQueue<T> {
pub fn try_consume_front(&mut self, mut amount: usize) -> bool {
if amount > self.len {
return false;
}
while amount != 0 {
let Some(segment) = self.segments.front() else {
return false;
};
let available = segment.as_ref().len();
if amount < available {
let Some(segment) = self.segments.front_mut() else {
return false;
};
if !segment.try_advance(amount) {
return false;
}
self.len -= amount;
return true;
}
amount -= available;
self.len -= available;
self.segments.pop_front();
}
true
}
}
impl<T> Default for SegmentQueue<T> {
fn default() -> Self {
Self::new()
}
}
impl<T> PrefixLength for SegmentQueue<T> {
fn prefix_len(&self) -> usize {
self.len()
}
}
impl<T> RetainedSegmentQueue<T> {
#[must_use]
pub const fn new() -> Self {
Self {
queue: SegmentQueue::new(),
front_offset: 0,
}
}
pub const fn len(&self) -> usize {
self.queue.len()
}
pub fn is_empty(&self) -> bool {
self.queue.is_empty()
}
pub fn segment_count(&self) -> usize {
self.queue.segment_count()
}
pub fn front(&self) -> Option<&T> {
self.queue.front()
}
pub fn back(&self) -> Option<&T> {
self.queue.back()
}
pub fn clear(&mut self) {
self.queue.clear_from(&mut self.front_offset);
}
}
impl<T: AsRef<[u8]>> RetainedSegmentQueue<T> {
pub fn try_push_back(&mut self, segment: T) -> Result<(), T> {
self.queue.try_push_back(segment)
}
pub fn try_mutate_back<R>(
&mut self,
additional: usize,
mutate: impl FnOnce(&mut T) -> R,
) -> Option<R> {
self.queue.try_mutate_back(additional, mutate)
}
pub fn range_available(&self, offset: usize, len: usize) -> bool {
self.queue.range_available(self.front_offset, offset, len)
}
pub fn extend_range(&self, offset: usize, len: usize, output: &mut Vec<u8>) -> bool {
self.queue
.extend_range(self.front_offset, offset, len, output)
}
pub fn consume_prefix_up_to(&mut self, requested: usize, removed: impl FnMut(T)) -> usize {
let amount = requested.min(self.queue.len());
unsafe {
self.queue
.consume_front_valid_from(&mut self.front_offset, amount, removed)
};
amount
}
}
impl<T> Default for RetainedSegmentQueue<T> {
fn default() -> Self {
Self::new()
}
}
impl<T> PrefixLength for RetainedSegmentQueue<T> {
fn prefix_len(&self) -> usize {
self.len()
}
}
impl private::Sealed for Bytes<Retained> {}
impl AdvanceSegment for Bytes<Retained> {
fn try_advance(&mut self, amount: usize) -> bool {
self.try_advance(amount)
}
}
struct SegmentMutation<'a, T: AsRef<[u8]>> {
segment: &'a mut T,
total: &'a mut usize,
total_ceiling: usize,
segment_ceiling: usize,
before: usize,
}
impl<T: AsRef<[u8]>> Drop for SegmentMutation<'_, T> {
fn drop(&mut self) {
let after = self.segment.as_ref().len();
if after > self.segment_ceiling {
abort();
}
let next = if after >= self.before {
self.total.checked_add(after - self.before)
} else {
self.total.checked_sub(self.before - after)
};
let Some(next) = next else {
abort();
};
if next > self.total_ceiling {
abort();
}
*self.total = next;
}
}