use std::io;
use std::mem::MaybeUninit;
use crate::errors::DriverError;
use crate::errors::Error;
use crate::errors::Result;
use crate::types::column::StringPool;
use crate::types::StatBuffer;
use crate::types::Unmarshal;
pub(crate) trait ReadEx {
fn read_bytes(&mut self, rv: &mut [u8]) -> Result<()>;
fn read_scalar<V>(&mut self) -> Result<V>
where V: Copy + Unmarshal<V> + StatBuffer;
fn read_string(&mut self) -> Result<String>;
fn skip_string(&mut self) -> Result<()>;
fn read_uvarint(&mut self) -> Result<u64>;
fn read_str_into_buffer(&mut self, pool: &mut StringPool) -> Result<()>;
}
const MAX_STACK_BUFFER_LEN: usize = 1024;
impl<T> ReadEx for T
where T: io::Read
{
fn read_bytes(&mut self, rv: &mut [u8]) -> Result<()> {
let mut i = 0;
while i < rv.len() {
let res_nread = {
let buf = &mut rv[i..];
self.read(buf)
};
match res_nread {
Ok(0) => {
let ret = io::Error::new(io::ErrorKind::WouldBlock, "would block");
return Err(ret.into());
}
Ok(nread) => i += nread,
Err(e) => return Err(From::from(e))
}
}
Ok(())
}
fn read_scalar<V>(&mut self) -> Result<V>
where V: Copy + Unmarshal<V> + StatBuffer {
let mut buffer = V::buffer();
self.read_bytes(buffer.as_mut())?;
Ok(V::unmarshal(buffer.as_ref()))
}
fn read_string(&mut self) -> Result<String> {
let str_len = self.read_uvarint()? as usize;
let mut buffer = vec![0_u8; str_len];
self.read_bytes(buffer.as_mut())?;
Ok(String::from_utf8(buffer)?)
}
fn skip_string(&mut self) -> Result<()> {
let str_len = self.read_uvarint()? as usize;
if str_len <= MAX_STACK_BUFFER_LEN {
unsafe {
let mut buffer: [MaybeUninit<u8>; MAX_STACK_BUFFER_LEN] =
MaybeUninit::uninit().assume_init();
self.read_bytes(
&mut *(&mut buffer[..str_len] as *mut [MaybeUninit<u8>] as *mut [u8])
)?;
}
} else {
let mut buffer = vec![0_u8; str_len];
self.read_bytes(buffer.as_mut())?;
}
Ok(())
}
fn read_uvarint(&mut self) -> Result<u64> {
let mut x = 0_u64;
let mut s = 0_u32;
let mut i = 0_usize;
loop {
let b: u8 = self.read_scalar()?;
if b < 0x80 {
if i == 9 && b > 1 {
return Err(Error::Driver(DriverError::Overflow));
}
return Ok(x | (u64::from(b) << s));
}
x |= u64::from(b & 0x7f) << s;
s += 7;
i += 1;
}
}
fn read_str_into_buffer(&mut self, pool: &mut StringPool) -> Result<()> {
let str_len = self.read_uvarint()? as usize;
let buffer = pool.allocate(str_len);
self.read_bytes(buffer)?;
Ok(())
}
}
#[test]
fn test_read_uvarint() {
use std::io::Cursor;
use super::ReadEx;
let bytes = [194_u8, 10];
let mut cursor = Cursor::new(bytes);
let actual = cursor.read_uvarint().unwrap();
assert_eq!(actual, 1346)
}