Optimize performance in several encode/decode programs

sqlx-postgres-protocol/benches/decode.rs

@@ -3,12 +3,13 @@ extern crate criterion;
 
 use bytes::Bytes;
 use criterion::{black_box, Criterion};
-use sqlx_postgres_protocol::{BackendKeyData, Decode, ParameterStatus, Response};
+use sqlx_postgres_protocol::{BackendKeyData, 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";
 
     c.bench_function("decode Response", |b| {
         b.iter(|| {
@@ -28,6 +29,12 @@ fn criterion_benchmark(c: &mut Criterion) {
             let _ = ParameterStatus::decode(black_box(Bytes::from_static(PARAM_STATUS))).unwrap();
         })
     });
+
+    c.bench_function("decode ReadyForQuery", |b| {
+        b.iter(|| {
+            let _ = ReadyForQuery::decode(black_box(Bytes::from_static(READY_FOR_QUERY))).unwrap();
+        })
+    });
 }
 
 criterion_group!(benches, criterion_benchmark);

sqlx-postgres-protocol/benches/encode.rs

@@ -2,25 +2,9 @@
 extern crate criterion;
 
 use criterion::Criterion;
-use sqlx_postgres_protocol::{Encode, PasswordMessage, Response, Severity, StartupMessage};
+use sqlx_postgres_protocol::{Encode, PasswordMessage, Query, StartupMessage, Terminate};
 
 fn criterion_benchmark(c: &mut Criterion) {
-    c.bench_function("encode Response::builder()", |b| {
-        let mut dst = Vec::with_capacity(1024);
-        b.iter(|| {
-            dst.clear();
-            Response::builder()
-                .severity(Severity::Notice)
-                .code("42710")
-                .message("extension \"uuid-ossp\" already exists, skipping")
-                .file("extension.c")
-                .line(1656)
-                .routine("CreateExtension")
-                .encode(&mut dst)
-                .unwrap();
-        })
-    });
-
     c.bench_function("encode PasswordMessage::cleartext", |b| {
         let mut dst = Vec::with_capacity(1024);
         b.iter(|| {
@@ -31,6 +15,22 @@ fn criterion_benchmark(c: &mut Criterion) {
         })
     });
 
+    c.bench_function("encode Query", |b| {
+        let mut dst = Vec::with_capacity(1024);
+        b.iter(|| {
+            dst.clear();
+            Query::new("SELECT 1, 2, 3").encode(&mut dst).unwrap();
+        })
+    });
+
+    c.bench_function("encode Terminate", |b| {
+        let mut dst = Vec::with_capacity(1024);
+        b.iter(|| {
+            dst.clear();
+            Terminate.encode(&mut dst).unwrap();
+        })
+    });
+
     c.bench_function("encode StartupMessage", |b| {
         let mut dst = Vec::with_capacity(1024);
         b.iter(|| {
@@ -49,7 +49,7 @@ fn criterion_benchmark(c: &mut Criterion) {
         })
     });
 
-    c.bench_function("encode Password(MD5)", |b| {
+    c.bench_function("encode Password::md5", |b| {
         let mut dst = Vec::with_capacity(1024);
         b.iter(|| {
             dst.clear();

sqlx-postgres-protocol/src/backend_key_data.rs

@@ -1,7 +1,7 @@
 use crate::Decode;
-use byteorder::{BigEndian, ByteOrder};
 use bytes::Bytes;
 use std::io;
+use std::convert::TryInto;
 
 #[derive(Debug)]
 pub struct BackendKeyData {
@@ -24,8 +24,8 @@ impl BackendKeyData {
 
 impl Decode for BackendKeyData {
     fn decode(src: Bytes) -> io::Result<Self> {
-        let process_id = BigEndian::read_u32(&src[..4]);
-        let secret_key = BigEndian::read_u32(&src[4..]);
+        let process_id = u32::from_be_bytes(src.as_ref()[0..4].try_into().unwrap());
+        let secret_key = u32::from_be_bytes(src.as_ref()[4..8].try_into().unwrap());
 
         Ok(Self {
             process_id,

sqlx-postgres-protocol/src/data_row.rs

@@ -39,10 +39,6 @@ pub struct DataValues<'a> {
 impl<'a> Iterator for DataValues<'a> {
     type Item = Option<&'a [u8]>;
 
-    fn size_hint(&self) -> (usize, Option<usize>) {
-        (self.rem as usize, Some(self.rem as usize))
-    }
-
     fn next(&mut self) -> Option<Self::Item> {
         if self.rem == 0 {
             return None;
@@ -62,6 +58,10 @@ impl<'a> Iterator for DataValues<'a> {
 
         Some(value)
     }
+
+    fn size_hint(&self) -> (usize, Option<usize>) {
+        (self.rem as usize, Some(self.rem as usize))
+    }
 }
 
 impl<'a> ExactSizeIterator for DataValues<'a> {}

sqlx-postgres-protocol/src/lib.rs

@@ -28,7 +28,7 @@ pub use self::{
     password_message::PasswordMessage,
     query::Query,
     ready_for_query::{ReadyForQuery, TransactionStatus},
-    response::{Response, ResponseBuilder, Severity},
+    response::{Response, Severity},
     row_description::{FieldDescription, FieldDescriptions, RowDescription},
     startup_message::StartupMessage,
     terminate::Terminate,

sqlx-postgres-protocol/src/query.rs

@@ -1,11 +1,11 @@
 use crate::Encode;
-use bytes::BufMut;
 use std::io;
 
 #[derive(Debug)]
 pub struct Query<'a>(&'a str);
 
 impl<'a> Query<'a> {
+    #[inline]
     pub fn new(query: &'a str) -> Self {
         Self(query)
     }
@@ -14,11 +14,10 @@ impl<'a> Query<'a> {
 impl Encode for Query<'_> {
     fn encode(&self, buf: &mut Vec<u8>) -> io::Result<()> {
         let len = self.0.len() + 4 + 1;
-        buf.reserve(len + 1);
-        buf.put_u8(b'Q');
-        buf.put_u32_be(len as u32);
-        buf.put(self.0);
-        buf.put_u8(0);
+        buf.push(b'Q');
+        buf.extend_from_slice(&(len as u32).to_be_bytes());
+        buf.extend_from_slice(self.0.as_bytes());
+        buf.push(0);
 
         Ok(())
     }

sqlx-postgres-protocol/src/ready_for_query.rs

@@ -1,5 +1,4 @@
-use crate::{Decode, Encode};
-use byteorder::{WriteBytesExt, BE};
+use crate::Decode;
 use bytes::Bytes;
 use std::io;
 
@@ -22,21 +21,6 @@ pub struct ReadyForQuery {
     pub status: TransactionStatus,
 }
 
-impl Encode for ReadyForQuery {
-    #[inline]
-    fn size_hint(&self) -> usize {
-        6
-    }
-
-    fn encode(&self, buf: &mut Vec<u8>) -> io::Result<()> {
-        buf.write_u8(b'Z')?;
-        buf.write_u32::<BE>(5)?;
-        buf.write_u8(self.status as u8)?;
-
-        Ok(())
-    }
-}
-
 impl Decode for ReadyForQuery {
     fn decode(src: Bytes) -> io::Result<Self> {
         if src.len() != 1 {
@@ -59,26 +43,12 @@ impl Decode for ReadyForQuery {
 #[cfg(test)]
 mod tests {
     use super::{ReadyForQuery, TransactionStatus};
-    use crate::{Decode, Encode};
+    use crate::Decode;
     use bytes::Bytes;
     use std::io;
 
     const READY_FOR_QUERY: &[u8] = b"E";
 
-    #[test]
-    fn it_encodes_ready_for_query() -> io::Result<()> {
-        let message = ReadyForQuery {
-            status: TransactionStatus::Error,
-        };
-
-        let mut dst = Vec::with_capacity(message.size_hint());
-        message.encode(&mut dst)?;
-
-        assert_eq!(&dst[5..], READY_FOR_QUERY);
-
-        Ok(())
-    }
-
     #[test]
     fn it_decodes_ready_for_query() -> io::Result<()> {
         let src = Bytes::from_static(READY_FOR_QUERY);

sqlx-postgres-protocol/src/response.rs

@@ -1,9 +1,7 @@
-use crate::{decode::get_str, Decode, Encode};
-use byteorder::{BigEndian, WriteBytesExt};
+use crate::{decode::get_str, Decode};
 use bytes::Bytes;
 use std::{
     fmt, io,
-    ops::Range,
     pin::Pin,
     ptr::NonNull,
     str::{self, FromStr},
@@ -104,11 +102,6 @@ unsafe impl Send for Response {}
 unsafe impl Sync for Response {}
 
 impl Response {
-    #[inline]
-    pub fn builder() -> ResponseBuilder {
-        ResponseBuilder::new()
-    }
-
     #[inline]
     pub fn severity(&self) -> Severity {
         self.severity
@@ -232,26 +225,6 @@ impl fmt::Debug for Response {
     }
 }
 
-impl Encode for Response {
-    #[inline]
-    fn size_hint(&self) -> usize {
-        self.storage.len() + 5
-    }
-
-    fn encode(&self, buf: &mut Vec<u8>) -> io::Result<()> {
-        if self.severity.is_error() {
-            buf.push(b'E');
-        } else {
-            buf.push(b'N');
-        }
-
-        buf.write_u32::<BigEndian>((4 + self.storage.len()) as u32)?;
-        buf.extend_from_slice(&self.storage);
-
-        Ok(())
-    }
-}
-
 impl Decode for Response {
     fn decode(src: Bytes) -> io::Result<Self> {
         let storage = Pin::new(src);
@@ -423,306 +396,16 @@ impl Decode for Response {
     }
 }
 
-pub struct ResponseBuilder {
-    storage: Vec<u8>,
-    severity: Option<Severity>,
-    code: Option<Range<usize>>,
-    message: Option<Range<usize>>,
-    detail: Option<Range<usize>>,
-    hint: Option<Range<usize>>,
-    position: Option<usize>,
-    internal_position: Option<usize>,
-    internal_query: Option<Range<usize>>,
-    where_: Option<Range<usize>>,
-    schema: Option<Range<usize>>,
-    table: Option<Range<usize>>,
-    column: Option<Range<usize>>,
-    data_type: Option<Range<usize>>,
-    constraint: Option<Range<usize>>,
-    file: Option<Range<usize>>,
-    line: Option<usize>,
-    routine: Option<Range<usize>>,
-}
-
-impl Default for ResponseBuilder {
-    fn default() -> Self {
-        Self {
-            storage: Vec::with_capacity(256),
-            severity: None,
-            message: None,
-            code: None,
-            detail: None,
-            hint: None,
-            position: None,
-            internal_position: None,
-            internal_query: None,
-            where_: None,
-            schema: None,
-            table: None,
-            column: None,
-            data_type: None,
-            constraint: None,
-            file: None,
-            line: None,
-            routine: None,
-        }
-    }
-}
-
-fn put_str(buf: &mut Vec<u8>, tag: u8, value: &str) -> Range<usize> {
-    buf.push(tag);
-    let beg = buf.len();
-    buf.extend_from_slice(value.as_bytes());
-    let end = buf.len();
-    buf.push(0);
-    beg..end
-}
-
-impl ResponseBuilder {
-    #[inline]
-    pub fn new() -> ResponseBuilder {
-        Self::default()
-    }
-
-    #[inline]
-    pub fn severity(mut self, severity: Severity) -> Self {
-        let sev = severity.to_str();
-
-        let _ = put_str(&mut self.storage, b'S', sev);
-        let _ = put_str(&mut self.storage, b'V', sev);
-
-        self.severity = Some(severity);
-        self
-    }
-
-    #[inline]
-    pub fn message(mut self, message: &str) -> Self {
-        self.message = Some(put_str(&mut self.storage, b'M', message));
-        self
-    }
-
-    #[inline]
-    pub fn code(mut self, code: &str) -> Self {
-        self.code = Some(put_str(&mut self.storage, b'C', code));
-        self
-    }
-
-    #[inline]
-    pub fn detail(mut self, detail: &str) -> Self {
-        self.detail = Some(put_str(&mut self.storage, b'D', detail));
-        self
-    }
-
-    #[inline]
-    pub fn hint(mut self, hint: &str) -> Self {
-        self.hint = Some(put_str(&mut self.storage, b'H', hint));
-        self
-    }
-
-    #[inline]
-    pub fn position(mut self, position: usize) -> Self {
-        self.storage.push(b'P');
-        // PANIC: Write to Vec<u8> is infallible
-        itoa::write(&mut self.storage, position).unwrap();
-        self.storage.push(0);
-
-        self.position = Some(position);
-        self
-    }
-
-    #[inline]
-    pub fn internal_position(mut self, position: usize) -> Self {
-        self.storage.push(b'p');
-        // PANIC: Write to Vec<u8> is infallible
-        itoa::write(&mut self.storage, position).unwrap();
-        self.storage.push(0);
-
-        self.internal_position = Some(position);
-        self
-    }
-
-    #[inline]
-    pub fn internal_query(mut self, query: &str) -> Self {
-        self.internal_query = Some(put_str(&mut self.storage, b'q', query));
-        self
-    }
-
-    #[inline]
-    pub fn where_(mut self, where_: &str) -> Self {
-        self.where_ = Some(put_str(&mut self.storage, b'w', where_));
-        self
-    }
-
-    #[inline]
-    pub fn schema(mut self, schema: &str) -> Self {
-        self.schema = Some(put_str(&mut self.storage, b's', schema));
-        self
-    }
-
-    #[inline]
-    pub fn table(mut self, table: &str) -> Self {
-        self.table = Some(put_str(&mut self.storage, b't', table));
-        self
-    }
-
-    #[inline]
-    pub fn column(mut self, column: &str) -> Self {
-        self.column = Some(put_str(&mut self.storage, b'c', column));
-        self
-    }
-
-    #[inline]
-    pub fn data_type(mut self, data_type: &str) -> Self {
-        self.data_type = Some(put_str(&mut self.storage, b'd', data_type));
-        self
-    }
-
-    #[inline]
-    pub fn constraint(mut self, constraint: &str) -> Self {
-        self.constraint = Some(put_str(&mut self.storage, b'n', constraint));
-        self
-    }
-
-    #[inline]
-    pub fn file(mut self, file: &str) -> Self {
-        self.file = Some(put_str(&mut self.storage, b'F', file));
-        self
-    }
-
-    #[inline]
-    pub fn line(mut self, line: usize) -> Self {
-        self.storage.push(b'L');
-        // PANIC: Write to Vec<u8> is infallible
-        itoa::write(&mut self.storage, line).unwrap();
-        self.storage.push(0);
-
-        self.line = Some(line);
-        self
-    }
-
-    #[inline]
-    pub fn routine(mut self, routine: &str) -> Self {
-        self.routine = Some(put_str(&mut self.storage, b'R', routine));
-        self
-    }
-
-    pub fn build(mut self) -> Response {
-        // Add a \0 terminator
-        self.storage.push(0);
-
-        // Freeze the storage and Pin so we can self-reference it
-        let storage = Pin::new(Bytes::from(self.storage));
-
-        let make_str_ref = |val: Option<Range<usize>>| unsafe {
-            val.map(|r| NonNull::from(str::from_utf8_unchecked(&storage[r])))
-        };
-
-        let code = make_str_ref(self.code);
-        let message = make_str_ref(self.message);
-        let detail = make_str_ref(self.detail);
-        let hint = make_str_ref(self.hint);
-        let internal_query = make_str_ref(self.internal_query);
-        let where_ = make_str_ref(self.where_);
-        let schema = make_str_ref(self.schema);
-        let table = make_str_ref(self.table);
-        let column = make_str_ref(self.column);
-        let data_type = make_str_ref(self.data_type);
-        let constraint = make_str_ref(self.constraint);
-        let file = make_str_ref(self.file);
-        let routine = make_str_ref(self.routine);
-
-        Response {
-            storage,
-            // FIXME: Default and don't panic here
-            severity: self.severity.expect("`severity` required by protocol"),
-            code: code.expect("`code` required by protocol"),
-            message: message.expect("`message` required by protocol"),
-            detail,
-            hint,
-            internal_query,
-            where_,
-            schema,
-            table,
-            column,
-            data_type,
-            constraint,
-            file,
-            routine,
-            line: self.line,
-            position: self.position,
-            internal_position: self.internal_position,
-        }
-    }
-}
-
-impl Encode for ResponseBuilder {
-    #[inline]
-    fn size_hint(&self) -> usize {
-        self.storage.len() + 6
-    }
-
-    fn encode(&self, buf: &mut Vec<u8>) -> io::Result<()> {
-        if self.severity.as_ref().map_or(false, |s| s.is_error()) {
-            buf.push(b'E');
-        } else {
-            buf.push(b'N');
-        }
-
-        buf.write_u32::<BigEndian>((5 + self.storage.len()) as u32)?;
-        buf.extend_from_slice(&self.storage);
-        buf.push(0);
-
-        Ok(())
-    }
-}
-
 #[cfg(test)]
 mod tests {
     use super::{Response, Severity};
-    use crate::{Decode, Encode};
+    use crate::Decode;
     use bytes::Bytes;
     use std::io;
 
     const RESPONSE: &[u8] = b"SNOTICE\0VNOTICE\0C42710\0Mextension \"uuid-ossp\" already exists, \
           skipping\0Fextension.c\0L1656\0RCreateExtension\0\0";
 
-    #[test]
-    fn it_encodes_response() -> io::Result<()> {
-        let message = Response::builder()
-            .severity(Severity::Notice)
-            .code("42710")
-            .message("extension \"uuid-ossp\" already exists, skipping")
-            .file("extension.c")
-            .line(1656)
-            .routine("CreateExtension")
-            .build();
-
-        let mut dst = Vec::with_capacity(message.size_hint());
-        message.encode(&mut dst)?;
-
-        assert_eq!(&dst[5..], RESPONSE);
-
-        Ok(())
-    }
-
-    #[test]
-    fn it_encodes_response_builder() -> io::Result<()> {
-        let message = Response::builder()
-            .severity(Severity::Notice)
-            .code("42710")
-            .message("extension \"uuid-ossp\" already exists, skipping")
-            .file("extension.c")
-            .line(1656)
-            .routine("CreateExtension");
-
-        let mut dst = Vec::with_capacity(message.size_hint());
-        message.encode(&mut dst)?;
-
-        assert_eq!(&dst[5..], RESPONSE);
-
-        Ok(())
-    }
-
     #[test]
     fn it_decodes_response() -> io::Result<()> {
         let src = Bytes::from_static(RESPONSE);

sqlx-postgres-protocol/src/startup_message.rs

@@ -8,10 +8,12 @@ pub struct StartupMessage<'a> {
 }
 
 impl<'a> StartupMessage<'a> {
+    #[inline]
     pub fn new(params: &'a [(&'a str, &'a str)]) -> Self {
         Self { params }
     }
 
+    #[inline]
     pub fn params(&self) -> &'a [(&'a str, &'a str)] {
         self.params
     }

sqlx-postgres-protocol/src/terminate.rs

@@ -1,5 +1,4 @@
 use crate::Encode;
-use bytes::BufMut;
 use std::io;
 
 #[derive(Debug)]
@@ -7,9 +6,8 @@ pub struct Terminate;
 
 impl Encode for Terminate {
     fn encode(&self, buf: &mut Vec<u8>) -> io::Result<()> {
-        buf.reserve(5);
-        buf.put_u8(b'X');
-        buf.put_u32_be(4);
+        buf.push(b'X');
+        buf.extend_from_slice(&4_u32.to_be_bytes());
 
         Ok(())
     }