use std::any::Any;
use std::fmt::{
self,
Debug,
};
use std::io::{
Cursor,
Result,
Write,
};
use bytes::{
Buf,
Bytes,
};
use super::{
Message,
Reflection,
};
#[allow(clippy::len_without_is_empty)]
pub trait PbBuffer: Any + Sized {
fn len(&self) -> usize;
fn copy_to_writer<W: Write + ?Sized>(&self, writer: &mut W) -> Result<()>;
fn copy_from_reader<B: Buf + ?Sized>(reader: &mut B) -> Result<Self>;
}
pub fn type_is<B1: 'static, B2: 'static>() -> Option<fn(B1) -> B2> {
let f: fn(B1) -> B1 = |x| x;
(&f as &dyn Any).downcast_ref::<fn(B1) -> B2>().copied()
}
pub trait PbBufferReader: Buf {
fn read_buffer<B: PbBuffer>(&mut self) -> Result<B> {
B::copy_from_reader(self)
}
fn split(&mut self, at: usize) -> Self;
}
pub trait PbBufferWriter: Write {
fn write_buffer<B: PbBuffer>(&mut self, buf: &B) -> Result<()>;
}
#[derive(Clone, PartialEq)]
pub struct Lazy<B> {
contents: Option<B>,
}
impl<B> Default for Lazy<B> {
fn default() -> Self {
Self { contents: None }
}
}
impl<B> Lazy<B> {
pub fn new(r: B) -> Self {
Self { contents: Some(r) }
}
pub fn into_buffer(self) -> B
where
B: Default,
{
self.contents.unwrap_or_default()
}
}
impl<B> Debug for Lazy<B> {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
f.debug_struct("Lazy")
.field("contents", &self.contents.as_ref().map(|_| "_"))
.finish()
}
}
impl<B: PbBuffer + PartialEq> Message for Lazy<B> {
fn compute_size(&self) -> usize {
self.contents.as_ref().map(PbBuffer::len).unwrap_or(0)
}
fn compute_grpc_slices_size(&self) -> usize {
self.contents.as_ref().map(PbBuffer::len).unwrap_or(0)
}
fn serialize<W: PbBufferWriter>(&self, w: &mut W) -> Result<()> {
if let Some(ref contents) = self.contents {
w.write_buffer(contents)?;
}
Ok(())
}
fn deserialize<R: PbBufferReader>(&mut self, r: &mut R) -> Result<()> {
self.contents = Some(r.read_buffer()?);
Ok(())
}
}
impl<B: PbBuffer + PartialEq> Reflection for Lazy<B> {}
impl<'a> PbBufferReader for Cursor<&'a [u8]> {
fn split(&mut self, at: usize) -> Self {
let pos = self.position() as usize;
self.advance(at);
let new_slice = &self.get_ref()[pos..pos + at];
Self::new(new_slice)
}
}
impl PbBuffer for Bytes {
#[inline]
fn len(&self) -> usize {
self.len()
}
fn copy_to_writer<W: Write + ?Sized>(&self, writer: &mut W) -> Result<()> {
writer.write_all(&self)
}
fn copy_from_reader<B: Buf + ?Sized>(reader: &mut B) -> Result<Self> {
let len = reader.remaining();
Ok(reader.copy_to_bytes(len))
}
}
impl PbBufferReader for Cursor<Bytes> {
fn read_buffer<B: PbBuffer>(&mut self) -> Result<B> {
if let Some(cast) = type_is::<Bytes, B>() {
let bytes = self.get_ref().slice((self.position() as usize)..);
Ok(cast(bytes))
} else {
B::copy_from_reader(self)
}
}
#[inline]
fn split(&mut self, at: usize) -> Self {
let pos = self.position() as usize;
self.advance(at);
let new_slice = self.get_ref().slice(pos..(pos + at));
Self::new(new_slice)
}
}
impl<'a> PbBufferWriter for Cursor<&'a mut Vec<u8>> {
#[inline]
fn write_buffer<B: PbBuffer>(&mut self, buf: &B) -> Result<()> {
buf.copy_to_writer(self)
}
}
impl<'a> PbBufferWriter for Cursor<&'a mut [u8]> {
#[inline]
fn write_buffer<B: PbBuffer>(&mut self, buf: &B) -> Result<()> {
buf.copy_to_writer(self)
}
}
pub struct CopyWriter<'a, W: Write> {
pub writer: &'a mut W,
}
impl<'a, W: Write + 'a> Write for CopyWriter<'a, W> {
#[inline]
fn write(&mut self, buf: &[u8]) -> Result<usize> {
self.writer.write(buf)
}
#[inline]
fn flush(&mut self) -> Result<()> {
self.writer.flush()
}
}
impl<'a, W: Write + 'a> PbBufferWriter for CopyWriter<'a, W> {
#[inline]
fn write_buffer<B: PbBuffer>(&mut self, buf: &B) -> Result<()> {
buf.copy_to_writer(self.writer)
}
}
#[test]
fn test_lazy_bytes_deserialize() {
let mut lazy = Lazy::<Bytes>::default();
let bytes = Bytes::from_static(b"asdfasdf");
lazy.deserialize(&mut Cursor::new(bytes.clone()))
.expect("failed to deserialize");
let deserialized = lazy.into_buffer();
assert_eq!(deserialized, bytes, "The entire buffer should be copied");
assert_eq!(
deserialized.as_ptr(),
bytes.as_ptr(),
"The Bytes instance should be cloned"
)
}