use core::any::{Any, TypeId};
use core::mem::{ManuallyDrop, MaybeUninit};
use core::panic::AssertUnwindSafe;
use super::aligned::AlignedBuf;
use super::vectored::Vectored;
pub unsafe trait IoBuf: Default + Send + Sync + 'static {
fn stable_ptr(&self) -> *const u8;
fn bytes_init(&self) -> usize;
fn bytes_total(&self) -> usize;
}
pub unsafe trait IoBufMut: IoBuf {
fn stable_mut_ptr(&mut self) -> *mut u8;
unsafe fn set_init(&mut self, pos: usize);
}
unsafe impl IoBuf for Vec<u8> {
fn stable_ptr(&self) -> *const u8 {
self.as_ptr()
}
fn bytes_init(&self) -> usize {
self.len()
}
fn bytes_total(&self) -> usize {
self.capacity()
}
}
unsafe impl IoBufMut for Vec<u8> {
fn stable_mut_ptr(&mut self) -> *mut u8 {
self.as_mut_ptr()
}
unsafe fn set_init(&mut self, pos: usize) {
unsafe { self.set_len(pos) }
}
}
fn cast<T, U>(value: T) -> Result<U, T>
where
T: 'static,
U: 'static,
{
if TypeId::of::<T>() != TypeId::of::<U>() {
return Err(value);
}
let value = ManuallyDrop::new(value);
Ok(unsafe { core::ptr::read((&raw const *value).cast::<U>()) })
}
pub(crate) trait Kept: Any + Send + Sync {
fn ptr(&self) -> *const u8;
fn mut_ptr(&mut self) -> *mut u8;
fn init(&self) -> usize;
fn total(&self) -> usize;
unsafe fn set_init(&mut self, pos: usize);
}
struct Written<B>(B);
impl<B> Kept for Written<B>
where
B: IoBuf,
{
fn ptr(&self) -> *const u8 {
self.0.stable_ptr()
}
fn mut_ptr(&mut self) -> *mut u8 {
self.0.stable_ptr().cast_mut()
}
fn init(&self) -> usize {
self.0.bytes_init()
}
fn total(&self) -> usize {
self.0.bytes_init()
}
unsafe fn set_init(&mut self, _pos: usize) {}
}
struct Filled<B>(B);
impl<B> Kept for Filled<B>
where
B: IoBufMut,
{
fn ptr(&self) -> *const u8 {
self.0.stable_ptr()
}
fn mut_ptr(&mut self) -> *mut u8 {
self.0.stable_mut_ptr()
}
fn init(&self) -> usize {
self.0.bytes_init()
}
fn total(&self) -> usize {
self.0.bytes_total()
}
unsafe fn set_init(&mut self, pos: usize) {
unsafe { self.0.set_init(pos) }
}
}
pub(crate) enum OwnedBuf {
Vec(Vec<u8>),
Boxed(BoxedBuf),
}
const _: () = assert!(size_of::<OwnedBuf>() == size_of::<Vec<u8>>());
pub(crate) struct BoxedBuf(ManuallyDrop<Box<Boxed>>);
impl BoxedBuf {
#[inline]
fn new(boxed: Boxed) -> Self {
Self(ManuallyDrop::new(Box::new(boxed)))
}
fn into_inner(self) -> Boxed {
let mut this = ManuallyDrop::new(self);
*unsafe { ManuallyDrop::take(&mut this.0) }
}
}
impl Drop for BoxedBuf {
#[inline(never)]
fn drop(&mut self) {
unsafe { ManuallyDrop::drop(&mut self.0) }
}
}
impl core::ops::Deref for BoxedBuf {
type Target = Boxed;
#[inline]
fn deref(&self) -> &Boxed {
&self.0
}
}
impl core::ops::DerefMut for BoxedBuf {
#[inline]
fn deref_mut(&mut self) -> &mut Boxed {
&mut self.0
}
}
pub(crate) enum Boxed {
Aligned(AlignedBuf),
Kept(AssertUnwindSafe<Box<dyn Kept>>),
Vectored(Vectored),
}
impl Default for OwnedBuf {
fn default() -> Self {
OwnedBuf::Vec(Vec::new())
}
}
impl OwnedBuf {
#[inline]
fn boxed(boxed: Boxed) -> Self {
OwnedBuf::Boxed(BoxedBuf::new(boxed))
}
#[inline]
pub(crate) fn written<B>(buf: B) -> Self
where
B: IoBuf,
{
let buf = match cast::<B, Vec<u8>>(buf) {
Ok(vec) => return OwnedBuf::Vec(vec),
Err(buf) => buf,
};
let buf = match cast::<B, OwnedBuf>(buf) {
Ok(owned) => return owned,
Err(buf) => buf,
};
match cast::<B, AlignedBuf>(buf) {
Ok(aligned) => Self::boxed(Boxed::Aligned(aligned)),
Err(buf) => Self::boxed(Boxed::Kept(AssertUnwindSafe(Box::new(Written(buf))))),
}
}
#[inline]
pub(crate) fn filled<B>(buf: B) -> Self
where
B: IoBufMut,
{
let buf = match cast::<B, Vec<u8>>(buf) {
Ok(vec) => return OwnedBuf::Vec(vec),
Err(buf) => buf,
};
let buf = match cast::<B, OwnedBuf>(buf) {
Ok(owned) => return owned,
Err(buf) => buf,
};
match cast::<B, AlignedBuf>(buf) {
Ok(aligned) => Self::boxed(Boxed::Aligned(aligned)),
Err(buf) => Self::boxed(Boxed::Kept(AssertUnwindSafe(Box::new(Filled(buf))))),
}
}
#[inline]
pub(crate) fn into_written<B>(self) -> B
where
B: IoBuf,
{
match self {
OwnedBuf::Vec(vec) => match cast::<Vec<u8>, B>(vec) {
Ok(same) => same,
Err(vec) => OwnedBuf::Vec(vec).into_other::<B, Written<B>>(|kept| kept.0),
},
other => other.into_other::<B, Written<B>>(|kept| kept.0),
}
}
#[inline]
pub(crate) fn into_filled<B>(self) -> B
where
B: IoBufMut,
{
match self {
OwnedBuf::Vec(vec) => match cast::<Vec<u8>, B>(vec) {
Ok(same) => same,
Err(vec) => OwnedBuf::Vec(vec).into_other::<B, Filled<B>>(|kept| kept.0),
},
other => other.into_other::<B, Filled<B>>(|kept| kept.0),
}
}
#[inline(never)]
fn into_other<B, K>(self, unwrap: fn(K) -> B) -> B
where
B: IoBuf,
K: Kept,
{
let boxed = match cast::<OwnedBuf, B>(self) {
Ok(same) => return same,
Err(OwnedBuf::Boxed(boxed)) => boxed,
Err(OwnedBuf::Vec(_)) => return B::default(),
};
match boxed.into_inner() {
Boxed::Aligned(aligned) => cast::<AlignedBuf, B>(aligned).unwrap_or_default(),
Boxed::Kept(AssertUnwindSafe(kept)) => {
let kept: Box<dyn Any> = kept;
kept.downcast::<K>()
.map_or_else(|_| B::default(), |k| unwrap(*k))
}
Boxed::Vectored(_) => B::default(),
}
}
pub(crate) fn vectored(bufs: Vec<OwnedBuf>) -> Self {
Self::boxed(Boxed::Vectored(Vectored::written(bufs)))
}
pub(crate) fn vectored_read(bufs: Vec<OwnedBuf>) -> Self {
Self::boxed(Boxed::Vectored(Vectored::filled(bufs)))
}
pub(crate) fn into_segments(self) -> Vec<OwnedBuf> {
match self.into_vectored() {
Ok(vectored) => vectored.into_bufs(),
Err(one) => vec![one],
}
}
pub(crate) fn into_vectored(self) -> Result<Vectored, OwnedBuf> {
match self {
OwnedBuf::Boxed(boxed) => match boxed.into_inner() {
Boxed::Vectored(vectored) => Ok(vectored),
one => Err(OwnedBuf::boxed(one)),
},
one => Err(one),
}
}
#[inline]
pub(crate) fn vectored_mut(&mut self) -> Option<&mut Vectored> {
match self {
OwnedBuf::Vec(_) => None,
OwnedBuf::Boxed(boxed) => boxed.vectored_mut(),
}
}
#[inline]
pub(crate) fn rest_from(&self, from: usize) -> &[u8] {
match self {
OwnedBuf::Vec(vec) => vec.get(from..).unwrap_or_default(),
OwnedBuf::Boxed(boxed) => boxed.rest_from(from),
}
}
#[inline]
pub(crate) fn len(&self) -> usize {
match self {
OwnedBuf::Vec(vec) => vec.len(),
OwnedBuf::Boxed(boxed) => boxed.len(),
}
}
#[inline]
pub(crate) fn is_empty(&self) -> bool {
self.len() == 0
}
#[inline]
pub(crate) fn capacity(&self) -> usize {
match self {
OwnedBuf::Vec(vec) => vec.capacity(),
OwnedBuf::Boxed(boxed) => boxed.capacity(),
}
}
#[inline]
fn as_mut_ptr(&mut self) -> *mut u8 {
match self {
OwnedBuf::Vec(vec) => vec.as_mut_ptr(),
OwnedBuf::Boxed(boxed) => boxed.as_mut_ptr(),
}
}
#[inline]
pub(crate) fn spare_capacity_mut(&mut self) -> &mut [MaybeUninit<u8>] {
match self {
OwnedBuf::Vec(vec) => vec.spare_capacity_mut(),
OwnedBuf::Boxed(boxed) => boxed.spare_capacity_mut(),
}
}
#[inline]
pub(crate) unsafe fn set_len(&mut self, len: usize) {
match self {
OwnedBuf::Vec(vec) => unsafe { vec.set_len(len) },
OwnedBuf::Boxed(boxed) => unsafe { boxed.set_len(len) },
}
}
#[inline]
pub(crate) unsafe fn add_len(&mut self, n: usize) {
match self {
OwnedBuf::Vec(vec) => unsafe { vec.set_len(vec.len() + n) },
OwnedBuf::Boxed(boxed) => {
let len = boxed.len() + n;
unsafe { boxed.set_len(len) }
}
}
}
#[inline]
pub(crate) fn truncate(&mut self, len: usize) {
if len < self.len() {
unsafe { self.set_len(len) };
}
}
#[inline]
pub(crate) fn append(&mut self, bytes: &[u8]) -> usize {
if let Some(vectored) = self.vectored_mut() {
return vectored.append(bytes);
}
let len = self.len();
let spare = self.spare_capacity_mut();
let n = bytes.len().min(spare.len());
for (to, from) in spare.iter_mut().zip(&bytes[..n]) {
to.write(*from);
}
unsafe { self.set_len(len + n) };
n
}
pub(crate) fn spare_zeroed(&mut self) -> &mut [u8] {
if self.vectored_mut().is_some() {
return &mut [];
}
let len = self.len();
let room = self.capacity().saturating_sub(len);
if room == 0 {
return &mut [];
}
let ptr = self.as_mut_ptr();
unsafe {
core::ptr::write_bytes(ptr.add(len), 0, room);
core::slice::from_raw_parts_mut(ptr.add(len), room)
}
}
}
impl core::ops::Deref for OwnedBuf {
type Target = [u8];
#[inline]
fn deref(&self) -> &[u8] {
match self {
OwnedBuf::Vec(vec) => vec,
OwnedBuf::Boxed(boxed) => boxed.bytes(),
}
}
}
impl Boxed {
#[inline(never)]
fn len(&self) -> usize {
match self {
Boxed::Aligned(aligned) => aligned.len(),
Boxed::Kept(kept) => kept.init(),
Boxed::Vectored(vectored) => vectored.len(),
}
}
#[inline(never)]
fn capacity(&self) -> usize {
match self {
Boxed::Aligned(aligned) => aligned.capacity(),
Boxed::Kept(kept) => kept.total(),
Boxed::Vectored(vectored) => vectored.capacity(),
}
}
#[inline(never)]
fn as_mut_ptr(&mut self) -> *mut u8 {
match self {
Boxed::Aligned(aligned) => aligned.as_mut_ptr(),
Boxed::Kept(kept) => kept.mut_ptr(),
Boxed::Vectored(_) => core::ptr::NonNull::dangling().as_ptr(),
}
}
#[inline(never)]
fn spare_capacity_mut(&mut self) -> &mut [MaybeUninit<u8>] {
match self {
Boxed::Aligned(aligned) => aligned.spare_capacity_mut(),
Boxed::Kept(kept) => {
let (len, total) = (kept.init(), kept.total());
if len >= total {
return &mut [];
}
unsafe {
core::slice::from_raw_parts_mut(
kept.mut_ptr().add(len).cast::<MaybeUninit<u8>>(),
total - len,
)
}
}
Boxed::Vectored(vectored) => vectored.spare_capacity_mut(),
}
}
#[inline(never)]
unsafe fn set_len(&mut self, len: usize) {
match self {
Boxed::Aligned(aligned) => unsafe { aligned.set_len(len) },
Boxed::Kept(kept) => unsafe { kept.set_init(len) },
Boxed::Vectored(vectored) => unsafe { vectored.set_len(len) },
}
}
#[inline(never)]
fn bytes(&self) -> &[u8] {
match self {
Boxed::Aligned(aligned) => aligned,
Boxed::Kept(kept) if kept.init() == 0 => &[],
Boxed::Kept(kept) => unsafe { core::slice::from_raw_parts(kept.ptr(), kept.init()) },
Boxed::Vectored(_) => &[],
}
}
#[inline(never)]
fn rest_from(&self, from: usize) -> &[u8] {
match self {
Boxed::Vectored(vectored) => vectored.rest_from(from),
other => other.bytes().get(from..).unwrap_or_default(),
}
}
#[inline(never)]
fn vectored_mut(&mut self) -> Option<&mut Vectored> {
match self {
Boxed::Vectored(vectored) => Some(vectored),
_ => None,
}
}
}
impl core::fmt::Debug for OwnedBuf {
fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
f.debug_struct("OwnedBuf")
.field("len", &self.len())
.field("capacity", &self.capacity())
.finish()
}
}
unsafe impl IoBuf for OwnedBuf {
fn stable_ptr(&self) -> *const u8 {
self.as_ptr()
}
fn bytes_init(&self) -> usize {
self.len()
}
fn bytes_total(&self) -> usize {
self.capacity()
}
}
unsafe impl IoBufMut for OwnedBuf {
fn stable_mut_ptr(&mut self) -> *mut u8 {
self.as_mut_ptr()
}
unsafe fn set_init(&mut self, pos: usize) {
unsafe { self.set_len(pos) }
}
}