#[cfg(not(feature = "std"))]
use alloc::vec::Vec;
use bytes::{
Buf as _, Bytes, BytesMut, TryGetError,
buf::{Chain, Take},
};
use core::any::Any;
pub trait Buf: bytes::Buf {
fn take_owned<T: OwnedBytes>(&mut self, len: usize) -> T {
T::from_bytes(self.copy_to_bytes(len))
}
fn try_take_owned<T: OwnedBytes>(&mut self, len: usize) -> Result<T, TryGetError> {
let available = self.remaining();
if available < len {
return Err(TryGetError {
requested: len,
available,
});
}
Ok(self.take_owned(len))
}
}
pub trait OwnedBytes: Sized + 'static {
fn from_bytes(bytes: Bytes) -> Self;
fn claim<V: Any>(view: V) -> Result<Self, V> {
let mut slot = Some(view);
let any: &mut dyn Any = &mut slot;
let claimed = any.downcast_mut::<Option<Self>>().and_then(Option::take);
claimed.ok_or_else(|| slot.expect("the slot must hold the view"))
}
}
impl OwnedBytes for Bytes {
#[inline]
fn from_bytes(bytes: Bytes) -> Self {
bytes
}
}
impl Buf for Bytes {}
impl Buf for BytesMut {}
impl<B: Buf + ?Sized> Buf for &mut B {
#[inline]
fn take_owned<T: OwnedBytes>(&mut self, len: usize) -> T {
(**self).take_owned(len)
}
}
impl<B: Buf> Buf for Take<B> {
#[inline]
fn take_owned<T: OwnedBytes>(&mut self, len: usize) -> T {
let remaining = self.remaining();
assert!(
len <= remaining,
"len must not exceed the remaining bytes: {len} > {remaining}"
);
let taken = self.get_mut().take_owned(len);
self.set_limit(self.limit() - len);
taken
}
}
impl<A: Buf, B: Buf> Buf for Chain<A, B> {
#[inline]
fn take_owned<T: OwnedBytes>(&mut self, len: usize) -> T {
let first = self.first_ref().remaining();
if first >= len {
self.first_mut().take_owned(len)
} else if first == 0 {
self.last_mut().take_owned(len)
} else {
T::from_bytes(self.copy_to_bytes(len))
}
}
}
pub trait Input: Into<Self::Buf> {
type Buf: Buf;
}
impl<B: Buf> Input for B {
type Buf = B;
}
impl Input for Vec<u8> {
type Buf = Bytes;
}
#[derive(Clone, Copy, Debug)]
pub struct Copying<'a>(
pub &'a [u8],
);
impl bytes::Buf for Copying<'_> {
#[inline]
fn remaining(&self) -> usize {
self.0.remaining()
}
#[inline]
fn chunk(&self) -> &[u8] {
self.0.chunk()
}
#[inline]
fn advance(&mut self, cnt: usize) {
self.0.advance(cnt);
}
#[inline]
fn copy_to_slice(&mut self, dst: &mut [u8]) {
self.0.copy_to_slice(dst);
}
}
impl Buf for Copying<'_> {}
#[derive(Clone, Copy, Debug)]
pub struct Viewing<'a> {
bytes: &'a Bytes,
pos: usize,
}
impl<'a> Viewing<'a> {
pub const fn new(bytes: &'a Bytes) -> Self {
Self { bytes, pos: 0 }
}
}
impl bytes::Buf for Viewing<'_> {
#[inline]
fn remaining(&self) -> usize {
self.bytes.len() - self.pos
}
#[inline]
fn chunk(&self) -> &[u8] {
&self.bytes[self.pos..]
}
#[inline]
fn advance(&mut self, cnt: usize) {
assert!(
cnt <= self.remaining(),
"advance must not exceed the remaining bytes"
);
self.pos += cnt;
}
#[inline]
fn copy_to_bytes(&mut self, len: usize) -> Bytes {
assert!(
len <= self.remaining(),
"length must not exceed the remaining bytes"
);
let bytes = self.bytes.slice(self.pos..self.pos + len);
self.pos += len;
bytes
}
}
impl Buf for Viewing<'_> {}
#[cfg(test)]
mod tests {
use super::*;
use crate::{Decode, DecodeExt, Encode, Read, ReadExt};
#[test]
fn test_take_owned_claims_matching_views() {
struct Reader(Bytes);
#[derive(Debug, PartialEq)]
struct Owned {
bytes: Bytes,
claimed: bool,
}
impl bytes::Buf for Reader {
fn remaining(&self) -> usize {
self.0.remaining()
}
fn chunk(&self) -> &[u8] {
self.0.chunk()
}
fn advance(&mut self, cnt: usize) {
self.0.advance(cnt);
}
}
impl Buf for Reader {
fn take_owned<T: OwnedBytes>(&mut self, len: usize) -> T {
let view = Owned {
bytes: self.0.split_to(len),
claimed: true,
};
T::claim(view).unwrap_or_else(|view| T::from_bytes(view.bytes))
}
}
impl OwnedBytes for Owned {
fn from_bytes(bytes: Bytes) -> Self {
Self {
bytes,
claimed: false,
}
}
}
let mut reader = Reader(Bytes::from_static(b"hello world"));
let claimed: Owned = reader.take_owned(5);
assert_eq!(claimed.bytes, &b"hello"[..]);
assert!(claimed.claimed);
let bytes: Bytes = reader.take_owned(6);
assert_eq!(bytes, &b" world"[..]);
let mut plain = Bytes::from_static(b"hello");
let built: Owned = plain.take_owned(5);
assert_eq!(built.bytes, &b"hello"[..]);
assert!(!built.claimed);
}
#[test]
fn test_take_owned_through_adapters() {
struct Reader(Bytes);
#[derive(Debug, PartialEq)]
struct Owned(Bytes, bool);
impl bytes::Buf for Reader {
fn remaining(&self) -> usize {
self.0.remaining()
}
fn chunk(&self) -> &[u8] {
self.0.chunk()
}
fn advance(&mut self, cnt: usize) {
self.0.advance(cnt);
}
}
impl Buf for Reader {
fn take_owned<T: OwnedBytes>(&mut self, len: usize) -> T {
T::claim(Owned(self.0.split_to(len), true))
.unwrap_or_else(|view| T::from_bytes(view.0))
}
}
impl OwnedBytes for Owned {
fn from_bytes(bytes: Bytes) -> Self {
Self(bytes, false)
}
}
let mut reader = Reader(Bytes::from_static(b"hello world"));
let claimed: Owned = (&mut &mut reader).take_owned(5);
assert_eq!(claimed, Owned(Bytes::from_static(b"hello"), true));
let mut limited = bytes::Buf::take(&mut reader, 4);
let claimed: Owned = limited.take_owned(3);
assert_eq!(claimed, Owned(Bytes::from_static(b" wo"), true));
assert_eq!(limited.limit(), 1);
assert_eq!(limited.remaining(), 1);
let claimed: Owned = limited.try_take_owned(1).unwrap();
assert_eq!(claimed, Owned(Bytes::from_static(b"r"), true));
assert_eq!(limited.limit(), 0);
let mut chained =
Reader(Bytes::from_static(b"ab")).chain(Reader(Bytes::from_static(b"cd")));
let claimed: Owned = chained.take_owned(1);
assert_eq!(claimed, Owned(Bytes::from_static(b"a"), true));
let spanning: Owned = chained.take_owned(2);
assert_eq!(spanning, Owned(Bytes::from_static(b"bc"), false));
let claimed: Owned = chained.take_owned(1);
assert_eq!(claimed, Owned(Bytes::from_static(b"d"), true));
assert!(!chained.has_remaining());
}
#[test]
fn test_try_take_owned_reports_shortfall() {
let mut input = Bytes::from_static(b"hello");
let payload = input.as_ptr();
let err = input.try_take_owned::<Bytes>(6).unwrap_err();
assert_eq!((err.requested, err.available), (6, 5));
assert_eq!(input.remaining(), 5);
let taken: Bytes = input.try_take_owned(5).unwrap();
assert_eq!(taken.as_ptr(), payload);
assert!(!input.has_remaining());
let mut inner = Bytes::from_static(b"abc");
let mut limited = bytes::Buf::take(&mut inner, 10);
let err = limited.try_take_owned::<Bytes>(5).unwrap_err();
assert_eq!((err.requested, err.available), (5, 3));
assert_eq!(limited.remaining(), 3);
assert_eq!(limited.limit(), 10);
let mut inner = Bytes::from_static(b"abcdefgh");
let mut limited = bytes::Buf::take(&mut inner, 4);
let err = limited.try_take_owned::<Bytes>(5).unwrap_err();
assert_eq!((err.requested, err.available), (5, 4));
assert_eq!(limited.remaining(), 4);
assert_eq!(inner.remaining(), 8);
}
#[test]
#[should_panic(expected = "len must not exceed the remaining bytes: 5 > 3")]
fn test_take_owned_past_inner_end() {
let mut inner = Bytes::from_static(b"abc");
let mut limited = bytes::Buf::take(&mut inner, 10);
let _: Bytes = limited.take_owned(5);
}
#[test]
fn test_vec_preserves_payload() {
let encoded = Bytes::from_static(b"hello").encode().to_vec();
let payload = encoded[1..].as_ptr();
let decoded = Bytes::decode_cfg(encoded, &(..).into()).unwrap();
assert_eq!(decoded, b"hello"[..]);
assert_eq!(decoded.as_ptr(), payload);
}
#[test]
fn test_owned_buffers_preserve_payload() {
let encoded = Bytes::from_static(b"hello").encode();
let payload = encoded[1..].as_ptr();
let cfg = (..).into();
assert_eq!(
Bytes::decode_cfg(encoded.clone(), &cfg).unwrap().as_ptr(),
payload
);
assert_eq!(
Bytes::decode_cfg(encoded.clone().take(encoded.len()), &cfg)
.unwrap()
.as_ptr(),
payload
);
assert_eq!(
Bytes::decode_cfg(encoded.clone().chain(Bytes::new()), &cfg)
.unwrap()
.as_ptr(),
payload
);
let mut input = encoded;
assert_eq!(Bytes::read_cfg(&mut input, &cfg).unwrap().as_ptr(), payload);
assert_eq!(input.remaining(), 0);
let encoded = BytesMut::from(&b"\x05hello"[..]);
let payload = encoded[1..].as_ptr();
assert_eq!(Bytes::decode_cfg(encoded, &cfg).unwrap().as_ptr(), payload);
}
#[test]
fn test_copying_scratch() {
let scratch = [0, 0, 0, 7];
let mut input = Copying(&scratch);
assert_eq!(input.chunk().as_ptr(), scratch.as_ptr());
assert_eq!(u32::read(&mut input).unwrap(), 7);
assert_eq!(input.remaining(), 0);
assert_eq!(u32::decode(Copying(&scratch)).unwrap(), 7);
}
#[test]
fn test_copying_payload() {
let encoded = Bytes::from_static(b"hello").encode();
let decoded = Bytes::decode_cfg(Copying(&encoded), &(..).into()).unwrap();
assert_eq!(decoded, b"hello"[..]);
assert_ne!(decoded.as_ptr(), encoded[1..].as_ptr());
}
#[test]
fn test_viewing_scalar_reads() {
let encoded = Bytes::copy_from_slice(b"\x00\x00\x00\x07ab");
let mut input = Viewing::new(&encoded);
assert!(encoded.is_unique());
assert_eq!(input.chunk().as_ptr(), encoded.as_ptr());
assert_eq!(u32::read(&mut input).unwrap(), 7);
let mut other = input;
assert_eq!(<[u8; 2]>::read(&mut input).unwrap(), *b"ab");
assert_eq!(other.remaining(), 2);
assert_eq!(<[u8; 2]>::read(&mut other).unwrap(), *b"ab");
assert!(!input.has_remaining());
assert!(input.chunk().is_empty());
assert_eq!(encoded.as_ref(), b"\x00\x00\x00\x07ab");
assert!(encoded.is_unique());
}
#[test]
fn test_viewing_retained_fields() {
let (first, second) = {
let encoded = Bytes::copy_from_slice(b"\x05hello\x05world");
let mut input = Viewing::new(&encoded);
let mut limited = (&mut input).take(6);
let first = Bytes::read_cfg(&mut limited, &(..).into()).unwrap();
assert_eq!(first.as_ptr(), encoded[1..].as_ptr());
assert!(!limited.has_remaining());
let second = Bytes::decode_cfg(input, &(..).into()).unwrap();
assert_eq!(second.as_ptr(), encoded[7..].as_ptr());
(first, second)
};
assert_eq!(first, b"hello"[..]);
assert_eq!(second, b"world"[..]);
}
#[test]
fn test_viewing_empty_ranges() {
for encoded in [Bytes::new(), Bytes::from_static(b"hello")] {
let mut input = Viewing::new(&encoded);
input.advance(0);
assert!(input.copy_to_bytes(0).is_empty());
assert_eq!(input.remaining(), encoded.len());
input.advance(encoded.len());
input.advance(0);
assert!(input.copy_to_bytes(0).is_empty());
assert!(!input.has_remaining());
assert!(input.chunk().is_empty());
}
}
#[test]
fn test_viewing_invalid_input() {
for bytes in [&b""[..], b"\x05hi", b"\x01ab", b"\x80"] {
let encoded = Bytes::copy_from_slice(bytes);
assert!(Bytes::decode_cfg(Viewing::new(&encoded), &(..).into()).is_err());
}
let encoded = Bytes::from_static(b"\x00");
assert!(
Bytes::decode_cfg(Viewing::new(&encoded), &(..).into())
.unwrap()
.is_empty()
);
}
#[test]
#[should_panic(expected = "advance must not exceed the remaining bytes")]
fn test_viewing_advance_past_end() {
let encoded = Bytes::from_static(b"a");
let mut input = Viewing::new(&encoded);
input.advance(1);
input.advance(usize::MAX);
}
#[test]
#[should_panic(expected = "length must not exceed the remaining bytes")]
fn test_viewing_copy_past_end() {
let encoded = Bytes::from_static(b"a");
let mut input = Viewing::new(&encoded);
input.advance(1);
input.copy_to_bytes(usize::MAX);
}
}