diff --git a/sqlx-postgres-protocol/benches/decode.rs b/sqlx-postgres-protocol/benches/decode.rs index 194f3aae..eecaa4cc 100644 --- a/sqlx-postgres-protocol/benches/decode.rs +++ b/sqlx-postgres-protocol/benches/decode.rs @@ -3,13 +3,16 @@ extern crate criterion; use bytes::Bytes; use criterion::{black_box, Criterion}; -use sqlx_postgres_protocol::{BackendKeyData, Decode, ParameterStatus, ReadyForQuery, Response}; +use sqlx_postgres_protocol::{ + BackendKeyData, DataRow, Decode, ParameterStatus, ReadyForQuery, Response, +}; fn criterion_benchmark(c: &mut Criterion) { const NOTICE_RESPONSE: &[u8] = b"SNOTICE\0VNOTICE\0C42710\0Mextension \"uuid-ossp\" already exists, skipping\0Fextension.c\0L1656\0RCreateExtension\0\0"; const PARAM_STATUS: &[u8] = b"session_authorization\0postgres\0"; const BACKEND_KEY_DATA: &[u8] = b"\0\0'\xc6\x89R\xc5+"; const READY_FOR_QUERY: &[u8] = b"E"; + const DATA_ROW: &[u8] = b"\0\x03\0\0\0\x011\0\0\0\x012\0\0\0\x013"; c.bench_function("decode Response", |b| { b.iter(|| { @@ -35,6 +38,12 @@ fn criterion_benchmark(c: &mut Criterion) { let _ = ReadyForQuery::decode(black_box(Bytes::from_static(READY_FOR_QUERY))).unwrap(); }) }); + + c.bench_function("decode DataRow", |b| { + b.iter(|| { + let _ = DataRow::decode(black_box(Bytes::from_static(DATA_ROW))).unwrap(); + }) + }); } criterion_group!(benches, criterion_benchmark); diff --git a/sqlx-postgres-protocol/src/backend_key_data.rs b/sqlx-postgres-protocol/src/backend_key_data.rs index 5ee4de1f..30f16e4c 100644 --- a/sqlx-postgres-protocol/src/backend_key_data.rs +++ b/sqlx-postgres-protocol/src/backend_key_data.rs @@ -1,7 +1,6 @@ use crate::Decode; use bytes::Bytes; -use std::io; -use std::convert::TryInto; +use std::{convert::TryInto, io}; #[derive(Debug)] pub struct BackendKeyData { diff --git a/sqlx-postgres-protocol/src/data_row.rs b/sqlx-postgres-protocol/src/data_row.rs index b75ec605..0458afc6 100644 --- a/sqlx-postgres-protocol/src/data_row.rs +++ b/sqlx-postgres-protocol/src/data_row.rs @@ -1,70 +1,82 @@ use crate::Decode; -use byteorder::{BigEndian, ByteOrder}; use bytes::Bytes; -use std::io; +use std::{ + convert::TryInto, + fmt::{self, Debug}, + io, + ops::Range, +}; -// TODO: Custom Debug for DataRow - -#[derive(Debug)] pub struct DataRow { - len: u16, - data: Bytes, -} - -impl DataRow { - pub fn values(&self) -> DataValues<'_> { - DataValues { - rem: self.len, - buf: &*self.data, - } - } + ranges: Vec>>, + buf: Bytes, } impl Decode for DataRow { fn decode(src: Bytes) -> io::Result { - let len = BigEndian::read_u16(&src[..2]); + let len = u16::from_be_bytes(src.as_ref()[..2].try_into().unwrap()); - Ok(Self { - len, - data: src.slice_from(2), - }) - } -} + let mut ranges = Vec::with_capacity(len as usize); + let mut rem = len; + let mut index = 2; -pub struct DataValues<'a> { - rem: u16, - buf: &'a [u8], -} + while rem > 0 { + // The length of the column value, in bytes (this count does not include itself). + // Can be zero. As a special case, -1 indicates a NULL column value. + // No value bytes follow in the NULL case. + let value_len = + i32::from_be_bytes(src.as_ref()[index..(index + 4)].try_into().unwrap()); + index += 4; -impl<'a> Iterator for DataValues<'a> { - type Item = Option<&'a [u8]>; + if value_len == -1 { + ranges.push(None); + } else { + let value_beg = index; + let value_end = value_beg + (value_len as usize); - fn next(&mut self) -> Option { - if self.rem == 0 { - return None; + ranges.push(Some(value_beg..(value_end as usize))); + + index += value_len as usize; + } + + rem -= 1; } - let len = BigEndian::read_i32(self.buf); - let size = (if len < 0 { 0 } else { len }) as usize; - - let value = if len == -1 { - None - } else { - Some(&self.buf[4..(4 + len) as usize]) - }; - - self.rem -= 1; - self.buf = &self.buf[(size + 4)..]; - - Some(value) - } - - fn size_hint(&self) -> (usize, Option) { - (self.rem as usize, Some(self.rem as usize)) + Ok(Self { ranges, buf: src }) } } -impl<'a> ExactSizeIterator for DataValues<'a> {} +impl DataRow { + #[inline] + pub fn is_empty(&self) -> bool { + self.ranges.is_empty() + } + + #[inline] + pub fn len(&self) -> usize { + self.ranges.len() + } + + #[inline] + pub fn get(&self, index: usize) -> Option<&[u8]> { + Some(&self.buf[self.ranges[index].clone()?]) + } +} + +impl Debug for DataRow { + fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result { + f.debug_struct("DataRow") + .field( + "values", + &self + .ranges + .iter() + .map(|range| Some(Bytes::from(&self.buf[range.clone()?]))) + .collect::>(), + ) + .finish() + } +} #[cfg(test)] mod tests { @@ -79,12 +91,14 @@ mod tests { fn it_decodes_data_row() -> io::Result<()> { let src = Bytes::from_static(DATA_ROW); let message = DataRow::decode(src)?; - assert_eq!(message.values().len(), 3); - for (index, value) in message.values().enumerate() { - // "1", "2", "3" - assert_eq!(value, Some(&[(index + 1 + 48) as u8][..])); - } + println!("{:?}", message); + + assert_eq!(message.len(), 3); + + assert_eq!(message.get(0), Some(&b"1"[..])); + assert_eq!(message.get(1), Some(&b"2"[..])); + assert_eq!(message.get(2), Some(&b"3"[..])); Ok(()) } diff --git a/sqlx-postgres-protocol/src/lib.rs b/sqlx-postgres-protocol/src/lib.rs index a6128015..20ee7238 100644 --- a/sqlx-postgres-protocol/src/lib.rs +++ b/sqlx-postgres-protocol/src/lib.rs @@ -20,7 +20,7 @@ pub use self::{ authentication::Authentication, backend_key_data::BackendKeyData, command_complete::CommandComplete, - data_row::{DataRow, DataValues}, + data_row::DataRow, decode::Decode, encode::Encode, message::Message,