use crate::postgres::AnyPostgresError;
use crate::postgres::types::int_types::{PostgresInt32, PostgresShort};
use crate::shared::Data;
use bun_core::String as BunString;
pub trait ReaderContext {
fn mark_message_start(&mut self);
fn peek(&self) -> &[u8];
fn skip(&mut self, count: usize);
fn ensure_length(&mut self, count: usize) -> bool;
fn read(&mut self, count: usize) -> Result<Data, AnyPostgresError>;
fn read_z(&mut self) -> Result<Data, AnyPostgresError>;
}
pub trait ProtocolInt: Sized + Copy + Eq {
const SIZE: usize;
fn from_be_slice(bytes: &[u8]) -> Self;
}
macro_rules! impl_protocol_int {
($($t:ty),*) => {$(
impl ProtocolInt for $t {
const SIZE: usize = core::mem::size_of::<$t>();
#[inline]
fn from_be_slice(bytes: &[u8]) -> Self {
let mut buf = [0u8; core::mem::size_of::<$t>()];
buf.copy_from_slice(&bytes[..Self::SIZE]);
<$t>::from_be_bytes(buf)
}
}
)*};
}
impl_protocol_int!(u8, i8, u16, i16, u32, i32, u64, i64);
impl<C: ReaderContext + ?Sized> ReaderContext for &mut C {
#[inline]
fn mark_message_start(&mut self) {
(**self).mark_message_start()
}
#[inline]
fn peek(&self) -> &[u8] {
(**self).peek()
}
#[inline]
fn skip(&mut self, count: usize) {
(**self).skip(count)
}
#[inline]
fn ensure_length(&mut self, count: usize) -> bool {
(**self).ensure_length(count)
}
#[inline]
fn read(&mut self, count: usize) -> Result<Data, AnyPostgresError> {
(**self).read(count)
}
#[inline]
fn read_z(&mut self) -> Result<Data, AnyPostgresError> {
(**self).read_z()
}
}
pub struct NewReaderWrap<Context: ReaderContext> {
pub wrapped: Context,
}
pub type Ctx<Context> = Context;
impl<Context: ReaderContext> NewReaderWrap<Context> {
#[inline]
pub fn reborrow(&mut self) -> NewReaderWrap<&mut Context> {
NewReaderWrap {
wrapped: &mut self.wrapped,
}
}
#[inline]
pub fn mark_message_start(&mut self) {
self.wrapped.mark_message_start();
}
#[inline]
pub fn read(&mut self, count: usize) -> Result<Data, AnyPostgresError> {
self.wrapped.read(count)
}
#[inline]
pub fn eat_message(&mut self, msg_: &'static [u8]) -> Result<(), AnyPostgresError> {
let msg = &msg_[1..];
self.ensure_capacity(msg.len())?;
let input = self.wrapped.read(msg.len())?;
if input.slice() == msg {
return Ok(());
}
Err(AnyPostgresError::InvalidMessage)
}
pub fn skip(&mut self, count: usize) -> Result<(), AnyPostgresError> {
self.wrapped.skip(count);
Ok(())
}
pub fn peek(&self) -> &[u8] {
self.wrapped.peek()
}
#[inline]
pub fn read_z(&mut self) -> Result<Data, AnyPostgresError> {
self.wrapped.read_z()
}
#[inline]
pub fn ensure_capacity(&mut self, count: usize) -> Result<(), AnyPostgresError> {
if !self.wrapped.ensure_length(count) {
return Err(AnyPostgresError::ShortRead);
}
Ok(())
}
pub fn int<Int: ProtocolInt>(&mut self) -> Result<Int, AnyPostgresError> {
let data = self.read(Int::SIZE)?;
let slice = data.slice();
if slice.len() < Int::SIZE {
return Err(AnyPostgresError::ShortRead);
}
Ok(Int::from_be_slice(&slice[0..Int::SIZE]))
}
pub fn peek_int<Int: ProtocolInt>(&self) -> Option<Int> {
let remain = self.peek();
if remain.len() < Int::SIZE {
return None;
}
Some(Int::from_be_slice(&remain[0..Int::SIZE]))
}
pub fn expect_int<Int: ProtocolInt>(&mut self, value: Int) -> Result<bool, AnyPostgresError> {
let actual = self.int::<Int>()?;
Ok(actual == value)
}
pub fn int4(&mut self) -> Result<PostgresInt32, AnyPostgresError> {
self.int::<PostgresInt32>()
}
pub fn short(&mut self) -> Result<PostgresShort, AnyPostgresError> {
self.int::<PostgresShort>()
}
pub fn length(&mut self) -> Result<PostgresInt32, AnyPostgresError> {
let expected = self.int::<PostgresInt32>()?;
self.ensure_capacity(expected.saturating_sub(4) as usize)?;
Ok(expected)
}
#[inline]
pub fn bytes(&mut self, count: usize) -> Result<Data, AnyPostgresError> {
self.read(count)
}
pub fn string(&mut self) -> Result<BunString, AnyPostgresError> {
let result = self.read_z()?;
Ok(BunString::borrow_utf8(result.slice()))
}
}
pub type NewReader<Context> = NewReaderWrap<Context>;