Skip to main content

qail_pg/protocol/
encoder.rs

1//! PostgreSQL Encoder (Visitor Pattern)
2//!
3//! Compiles Qail AST into PostgreSQL wire protocol bytes.
4//! This is pure, synchronous computation - no I/O, no async.
5//!
6//! # Architecture
7//!
8//! Layer 2 of the QAIL architecture:
9//! - Input: Qail (AST)
10//! - Output: BytesMut (ready to send over the wire)
11//!
12//! The async I/O layer (Layer 3) consumes these bytes.
13
14use super::EncodeError;
15use bytes::BytesMut;
16
17/// Takes a Qail and produces wire protocol bytes.
18/// This is the "Visitor" in the visitor pattern.
19pub struct PgEncoder;
20
21impl PgEncoder {
22    /// Wire format code for text columns.
23    pub const FORMAT_TEXT: i16 = 0;
24    /// Wire format code for binary columns.
25    pub const FORMAT_BINARY: i16 = 1;
26
27    #[inline(always)]
28    fn param_format_wire_len(param_format: i16) -> usize {
29        if param_format == Self::FORMAT_TEXT {
30            2 // parameter format count = 0 (server default text)
31        } else {
32            4 // parameter format count = 1 + one format code for all parameters
33        }
34    }
35
36    #[inline(always)]
37    fn encode_param_formats_vec(content: &mut Vec<u8>, param_format: i16) {
38        if param_format == Self::FORMAT_TEXT {
39            content.extend_from_slice(&0i16.to_be_bytes());
40        } else {
41            content.extend_from_slice(&1i16.to_be_bytes());
42            content.extend_from_slice(&param_format.to_be_bytes());
43        }
44    }
45
46    #[inline(always)]
47    fn encode_param_formats_bytesmut(buf: &mut BytesMut, param_format: i16) {
48        if param_format == Self::FORMAT_TEXT {
49            buf.extend_from_slice(&0i16.to_be_bytes());
50        } else {
51            buf.extend_from_slice(&1i16.to_be_bytes());
52            buf.extend_from_slice(&param_format.to_be_bytes());
53        }
54    }
55
56    #[inline(always)]
57    fn result_format_wire_len(result_format: i16) -> usize {
58        if result_format == Self::FORMAT_TEXT {
59            2 // result format count = 0
60        } else {
61            4 // result format count = 1 + one format code
62        }
63    }
64
65    #[inline(always)]
66    fn params_wire_len(params: &[Option<Vec<u8>>]) -> Result<usize, EncodeError> {
67        params.iter().try_fold(0usize, |acc, p| {
68            let field_size = 4usize
69                .checked_add(p.as_ref().map_or(0usize, |v| v.len()))
70                .ok_or(EncodeError::MessageTooLarge(usize::MAX))?;
71            acc.checked_add(field_size)
72                .ok_or(EncodeError::MessageTooLarge(usize::MAX))
73        })
74    }
75
76    #[inline(always)]
77    fn encode_result_formats_vec(content: &mut Vec<u8>, result_format: i16) {
78        if result_format == Self::FORMAT_TEXT {
79            content.extend_from_slice(&0i16.to_be_bytes());
80        } else {
81            content.extend_from_slice(&1i16.to_be_bytes());
82            content.extend_from_slice(&result_format.to_be_bytes());
83        }
84    }
85
86    #[inline(always)]
87    fn encode_result_formats_bytesmut(buf: &mut BytesMut, result_format: i16) {
88        if result_format == Self::FORMAT_TEXT {
89            buf.extend_from_slice(&0i16.to_be_bytes());
90        } else {
91            buf.extend_from_slice(&1i16.to_be_bytes());
92            buf.extend_from_slice(&result_format.to_be_bytes());
93        }
94    }
95
96    #[inline(always)]
97    fn content_len_to_wire_len(content_len: usize) -> Result<i32, EncodeError> {
98        let total = content_len
99            .checked_add(4)
100            .ok_or(EncodeError::MessageTooLarge(usize::MAX))?;
101        i32::try_from(total).map_err(|_| EncodeError::MessageTooLarge(total))
102    }
103
104    #[inline(always)]
105    fn usize_to_i16(n: usize) -> Result<i16, EncodeError> {
106        i16::try_from(n).map_err(|_| EncodeError::TooManyParameters(n))
107    }
108
109    #[inline(always)]
110    fn usize_to_i32(n: usize) -> Result<i32, EncodeError> {
111        i32::try_from(n).map_err(|_| EncodeError::MessageTooLarge(n))
112    }
113
114    #[inline(always)]
115    fn has_nul(s: &str) -> bool {
116        s.as_bytes().contains(&0)
117    }
118
119    /// Wire length of a single Bind message, including the type byte.
120    #[inline]
121    pub fn bind_wire_len_with_formats(
122        statement: &str,
123        params: &[Option<Vec<u8>>],
124        param_format: i16,
125        result_format: i16,
126    ) -> Result<usize, EncodeError> {
127        if Self::has_nul(statement) {
128            return Err(EncodeError::NullByte);
129        }
130        if params.len() > i16::MAX as usize {
131            return Err(EncodeError::TooManyParameters(params.len()));
132        }
133
134        let params_size = Self::params_wire_len(params)?;
135        let param_formats_size = Self::param_format_wire_len(param_format);
136        let result_formats_size = Self::result_format_wire_len(result_format);
137        let content_len = 1usize
138            .checked_add(statement.len())
139            .and_then(|v| v.checked_add(1))
140            .and_then(|v| v.checked_add(2))
141            .and_then(|v| v.checked_add(param_formats_size))
142            .and_then(|v| v.checked_add(params_size))
143            .and_then(|v| v.checked_add(result_formats_size))
144            .ok_or(EncodeError::MessageTooLarge(usize::MAX))?;
145        let wire_len = Self::content_len_to_wire_len(content_len)? as usize;
146        1usize
147            .checked_add(wire_len)
148            .ok_or(EncodeError::MessageTooLarge(usize::MAX))
149    }
150
151    /// Wire length of Bind + Execute, including message type bytes.
152    #[inline]
153    pub fn bind_execute_wire_len_with_formats(
154        statement: &str,
155        params: &[Option<Vec<u8>>],
156        param_format: i16,
157        result_format: i16,
158    ) -> Result<usize, EncodeError> {
159        Self::bind_wire_len_with_formats(statement, params, param_format, result_format)?
160            .checked_add(10)
161            .ok_or(EncodeError::MessageTooLarge(usize::MAX))
162    }
163
164    /// Wire length of Bind + Execute + Sync, including message type bytes.
165    #[inline]
166    pub fn bind_execute_sync_wire_len_with_formats(
167        statement: &str,
168        params: &[Option<Vec<u8>>],
169        param_format: i16,
170        result_format: i16,
171    ) -> Result<usize, EncodeError> {
172        Self::bind_execute_wire_len_with_formats(statement, params, param_format, result_format)?
173            .checked_add(5)
174            .ok_or(EncodeError::MessageTooLarge(usize::MAX))
175    }
176
177    /// Fallible simple-query encoder.
178    pub fn try_encode_query_string(sql: &str) -> Result<BytesMut, EncodeError> {
179        if Self::has_nul(sql) {
180            return Err(EncodeError::NullByte);
181        }
182
183        let mut buf = BytesMut::new();
184        let content_len = sql.len() + 1; // +1 for null terminator
185        let total_len = Self::content_len_to_wire_len(content_len)?;
186
187        buf.extend_from_slice(b"Q");
188        buf.extend_from_slice(&total_len.to_be_bytes());
189        buf.extend_from_slice(sql.as_bytes());
190        buf.extend_from_slice(&[0]);
191        Ok(buf)
192    }
193
194    /// Encode a Terminate message to close the connection.
195    pub fn encode_terminate() -> BytesMut {
196        let mut buf = BytesMut::new();
197        buf.extend_from_slice(&[b'X', 0, 0, 0, 4]);
198        buf
199    }
200
201    /// Encode a Sync message (end of pipeline in extended query protocol).
202    pub fn encode_sync() -> BytesMut {
203        let mut buf = BytesMut::new();
204        buf.extend_from_slice(&[b'S', 0, 0, 0, 4]);
205        buf
206    }
207
208    // ==================== Extended Query Protocol ====================
209
210    /// Fallible Parse message encoder.
211    pub fn try_encode_parse(
212        name: &str,
213        sql: &str,
214        param_types: &[u32],
215    ) -> Result<BytesMut, EncodeError> {
216        if Self::has_nul(name) || Self::has_nul(sql) {
217            return Err(EncodeError::NullByte);
218        }
219        if param_types.len() > i16::MAX as usize {
220            return Err(EncodeError::TooManyParameters(param_types.len()));
221        }
222
223        let mut buf = BytesMut::new();
224        buf.extend_from_slice(b"P");
225
226        let mut content = Vec::new();
227        content.extend_from_slice(name.as_bytes());
228        content.push(0);
229        content.extend_from_slice(sql.as_bytes());
230        content.push(0);
231        let param_count = Self::usize_to_i16(param_types.len())?;
232        content.extend_from_slice(&param_count.to_be_bytes());
233        for &oid in param_types {
234            content.extend_from_slice(&oid.to_be_bytes());
235        }
236
237        let len = Self::content_len_to_wire_len(content.len())?;
238        buf.extend_from_slice(&len.to_be_bytes());
239        buf.extend_from_slice(&content);
240        Ok(buf)
241    }
242
243    /// Encode a Parse message directly into an existing buffer.
244    pub fn try_encode_parse_to(
245        buf: &mut BytesMut,
246        name: &str,
247        sql: &str,
248        param_types: &[u32],
249    ) -> Result<(), EncodeError> {
250        if Self::has_nul(name) || Self::has_nul(sql) {
251            return Err(EncodeError::NullByte);
252        }
253        if param_types.len() > i16::MAX as usize {
254            return Err(EncodeError::TooManyParameters(param_types.len()));
255        }
256
257        let content_len = name
258            .len()
259            .checked_add(1)
260            .and_then(|v| v.checked_add(sql.len()))
261            .and_then(|v| v.checked_add(1))
262            .and_then(|v| v.checked_add(2))
263            .and_then(|v| v.checked_add(param_types.len().checked_mul(4)?))
264            .ok_or(EncodeError::MessageTooLarge(usize::MAX))?;
265        let wire_len = Self::content_len_to_wire_len(content_len)?;
266
267        buf.reserve(1 + 4 + content_len);
268        buf.extend_from_slice(b"P");
269        buf.extend_from_slice(&wire_len.to_be_bytes());
270        buf.extend_from_slice(name.as_bytes());
271        buf.extend_from_slice(&[0]);
272        buf.extend_from_slice(sql.as_bytes());
273        buf.extend_from_slice(&[0]);
274        let param_count = Self::usize_to_i16(param_types.len())?;
275        buf.extend_from_slice(&param_count.to_be_bytes());
276        for &oid in param_types {
277            buf.extend_from_slice(&oid.to_be_bytes());
278        }
279        Ok(())
280    }
281
282    /// Encode a Bind message (bind parameters to a prepared statement).
283    /// Wire format:
284    /// - 'B' (1 byte) - message type
285    /// - length (4 bytes)
286    /// - portal name (null-terminated)
287    /// - statement name (null-terminated)
288    /// - format code section (2-4 bytes) - default path uses 0 (all text)
289    /// - parameter count (2 bytes)
290    /// - for each parameter: length (4 bytes, -1 for NULL), data
291    /// - result format count + codes
292    ///
293    /// # Arguments
294    ///
295    /// * `portal` — Destination portal name (empty string for unnamed).
296    /// * `statement` — Source prepared statement name (empty string for unnamed).
297    /// * `params` — Parameter values; `None` entries encode as SQL NULL.
298    pub fn encode_bind(
299        portal: &str,
300        statement: &str,
301        params: &[Option<Vec<u8>>],
302    ) -> Result<BytesMut, EncodeError> {
303        Self::encode_bind_with_result_format(portal, statement, params, Self::FORMAT_TEXT)
304    }
305
306    /// Encode a Bind message with explicit result-column format.
307    ///
308    /// `result_format` is PostgreSQL wire format code: `0 = text`, `1 = binary`.
309    /// For `0`, this encodes "result format count = 0" (server default text).
310    /// For non-zero codes, this encodes one explicit result format code.
311    pub fn encode_bind_with_result_format(
312        portal: &str,
313        statement: &str,
314        params: &[Option<Vec<u8>>],
315        result_format: i16,
316    ) -> Result<BytesMut, EncodeError> {
317        Self::encode_bind_with_formats(portal, statement, params, Self::FORMAT_TEXT, result_format)
318    }
319
320    /// Encode a Bind message with explicit parameter and result format codes.
321    ///
322    /// `param_format` / `result_format` are PostgreSQL wire format codes:
323    /// `0 = text`, `1 = binary`.
324    ///
325    /// For `param_format = 0`, this encodes "parameter format count = 0"
326    /// (server default text). For non-zero, this encodes one explicit format
327    /// code applied to all parameters.
328    pub fn encode_bind_with_formats(
329        portal: &str,
330        statement: &str,
331        params: &[Option<Vec<u8>>],
332        param_format: i16,
333        result_format: i16,
334    ) -> Result<BytesMut, EncodeError> {
335        if Self::has_nul(portal) || Self::has_nul(statement) {
336            return Err(EncodeError::NullByte);
337        }
338        if params.len() > i16::MAX as usize {
339            return Err(EncodeError::TooManyParameters(params.len()));
340        }
341
342        let mut buf = BytesMut::new();
343
344        // Message type 'B'
345        buf.extend_from_slice(b"B");
346
347        let mut content = Vec::new();
348
349        // Portal name (null-terminated)
350        content.extend_from_slice(portal.as_bytes());
351        content.push(0);
352
353        // Statement name (null-terminated)
354        content.extend_from_slice(statement.as_bytes());
355        content.push(0);
356
357        // Parameter format codes
358        Self::encode_param_formats_vec(&mut content, param_format);
359
360        // Parameter count
361        let param_count = Self::usize_to_i16(params.len())?;
362        content.extend_from_slice(&param_count.to_be_bytes());
363
364        // Parameters
365        for param in params {
366            match param {
367                None => {
368                    // NULL: length = -1
369                    content.extend_from_slice(&(-1i32).to_be_bytes());
370                }
371                Some(data) => {
372                    let data_len = Self::usize_to_i32(data.len())?;
373                    content.extend_from_slice(&data_len.to_be_bytes());
374                    content.extend_from_slice(data);
375                }
376            }
377        }
378
379        // Result format codes: default text (count=0) or explicit code.
380        Self::encode_result_formats_vec(&mut content, result_format);
381
382        // Length
383        let len = Self::content_len_to_wire_len(content.len())?;
384        buf.extend_from_slice(&len.to_be_bytes());
385        buf.extend_from_slice(&content);
386
387        Ok(buf)
388    }
389
390    /// Fallible Execute message encoder.
391    pub fn try_encode_execute(portal: &str, max_rows: i32) -> Result<BytesMut, EncodeError> {
392        if Self::has_nul(portal) {
393            return Err(EncodeError::NullByte);
394        }
395        if max_rows < 0 {
396            return Err(EncodeError::InvalidMaxRows(max_rows));
397        }
398
399        let mut buf = BytesMut::new();
400        buf.extend_from_slice(b"E");
401
402        let mut content = Vec::new();
403        content.extend_from_slice(portal.as_bytes());
404        content.push(0);
405        content.extend_from_slice(&max_rows.to_be_bytes());
406
407        let len = Self::content_len_to_wire_len(content.len())?;
408        buf.extend_from_slice(&len.to_be_bytes());
409        buf.extend_from_slice(&content);
410        Ok(buf)
411    }
412
413    /// Fallible Describe message encoder.
414    pub fn try_encode_describe(is_portal: bool, name: &str) -> Result<BytesMut, EncodeError> {
415        if Self::has_nul(name) {
416            return Err(EncodeError::NullByte);
417        }
418
419        let mut buf = BytesMut::new();
420        buf.extend_from_slice(b"D");
421
422        let mut content = Vec::new();
423        content.push(if is_portal { b'P' } else { b'S' });
424        content.extend_from_slice(name.as_bytes());
425        content.push(0);
426
427        let len = Self::content_len_to_wire_len(content.len())?;
428        buf.extend_from_slice(&len.to_be_bytes());
429        buf.extend_from_slice(&content);
430        Ok(buf)
431    }
432
433    /// Encode a complete extended query pipeline (OPTIMIZED).
434    /// This combines Parse + Bind + Execute + Sync in a single buffer.
435    /// Zero intermediate allocations - writes directly to pre-sized BytesMut.
436    pub fn encode_extended_query(
437        sql: &str,
438        params: &[Option<Vec<u8>>],
439    ) -> Result<BytesMut, EncodeError> {
440        Self::encode_extended_query_with_result_format(sql, params, Self::FORMAT_TEXT)
441    }
442
443    /// Encode a complete extended query pipeline with explicit result format.
444    ///
445    /// `result_format` is PostgreSQL wire format code: `0 = text`, `1 = binary`.
446    pub fn encode_extended_query_with_result_format(
447        sql: &str,
448        params: &[Option<Vec<u8>>],
449        result_format: i16,
450    ) -> Result<BytesMut, EncodeError> {
451        Self::encode_extended_query_with_formats(sql, params, Self::FORMAT_TEXT, result_format)
452    }
453
454    /// Encode a complete extended query pipeline with explicit parameter and result formats.
455    ///
456    /// `param_format` / `result_format` are PostgreSQL wire format codes:
457    /// `0 = text`, `1 = binary`.
458    pub fn encode_extended_query_with_formats(
459        sql: &str,
460        params: &[Option<Vec<u8>>],
461        param_format: i16,
462        result_format: i16,
463    ) -> Result<BytesMut, EncodeError> {
464        if Self::has_nul(sql) {
465            return Err(EncodeError::NullByte);
466        }
467        if params.len() > i16::MAX as usize {
468            return Err(EncodeError::TooManyParameters(params.len()));
469        }
470
471        // Calculate total size upfront to avoid reallocations
472        // Bind: 1 + 4 + 1 + 1 + param_formats + 2 + params_data + result_formats
473        // Execute: 1 + 4 + 1 + 4 = 10
474        // Sync: 5
475        let params_size = params.iter().try_fold(0usize, |acc, p| {
476            let field_size = 4usize
477                .checked_add(p.as_ref().map_or(0usize, |v| v.len()))
478                .ok_or(EncodeError::MessageTooLarge(usize::MAX))?;
479            acc.checked_add(field_size)
480                .ok_or(EncodeError::MessageTooLarge(usize::MAX))
481        })?;
482        let param_formats_size = Self::param_format_wire_len(param_format);
483        let result_formats_size = Self::result_format_wire_len(result_format);
484        let total_size = 9usize
485            .checked_add(sql.len())
486            .and_then(|v| v.checked_add(9))
487            .and_then(|v| v.checked_add(params_size))
488            .and_then(|v| v.checked_add(param_formats_size))
489            .and_then(|v| v.checked_add(result_formats_size))
490            .and_then(|v| v.checked_add(10))
491            .and_then(|v| v.checked_add(5))
492            .ok_or(EncodeError::MessageTooLarge(usize::MAX))?;
493
494        let mut buf = BytesMut::with_capacity(total_size);
495
496        // ===== PARSE =====
497        buf.extend_from_slice(b"P");
498        let parse_content_len = 1usize
499            .checked_add(sql.len())
500            .and_then(|v| v.checked_add(1))
501            .and_then(|v| v.checked_add(2))
502            .ok_or(EncodeError::MessageTooLarge(usize::MAX))?;
503        let parse_len = Self::content_len_to_wire_len(parse_content_len)?;
504        buf.extend_from_slice(&parse_len.to_be_bytes());
505        buf.extend_from_slice(&[0]); // Unnamed statement
506        buf.extend_from_slice(sql.as_bytes());
507        buf.extend_from_slice(&[0]); // Null terminator
508        buf.extend_from_slice(&0i16.to_be_bytes()); // No param types (infer)
509
510        // ===== BIND =====
511        buf.extend_from_slice(b"B");
512        let bind_content_len = 1usize
513            .checked_add(1)
514            .and_then(|v| v.checked_add(2))
515            .and_then(|v| v.checked_add(param_formats_size))
516            .and_then(|v| v.checked_add(params_size))
517            .and_then(|v| v.checked_add(result_formats_size))
518            .ok_or(EncodeError::MessageTooLarge(usize::MAX))?;
519        let bind_len = Self::content_len_to_wire_len(bind_content_len)?;
520        buf.extend_from_slice(&bind_len.to_be_bytes());
521        buf.extend_from_slice(&[0]); // Unnamed portal
522        buf.extend_from_slice(&[0]); // Unnamed statement
523        Self::encode_param_formats_bytesmut(&mut buf, param_format);
524        let param_count = Self::usize_to_i16(params.len())?;
525        buf.extend_from_slice(&param_count.to_be_bytes());
526        for param in params {
527            match param {
528                None => buf.extend_from_slice(&(-1i32).to_be_bytes()),
529                Some(data) => {
530                    let data_len = Self::usize_to_i32(data.len())?;
531                    buf.extend_from_slice(&data_len.to_be_bytes());
532                    buf.extend_from_slice(data);
533                }
534            }
535        }
536        Self::encode_result_formats_bytesmut(&mut buf, result_format);
537
538        // ===== EXECUTE =====
539        buf.extend_from_slice(b"E");
540        buf.extend_from_slice(&9i32.to_be_bytes()); // len = 4 + 1 + 4
541        buf.extend_from_slice(&[0]); // Unnamed portal
542        buf.extend_from_slice(&0i32.to_be_bytes()); // Unlimited rows
543
544        // ===== SYNC =====
545        buf.extend_from_slice(&[b'S', 0, 0, 0, 4]);
546
547        Ok(buf)
548    }
549
550    /// Fallible CopyFail encoder.
551    pub fn try_encode_copy_fail(reason: &str) -> Result<BytesMut, EncodeError> {
552        if Self::has_nul(reason) {
553            return Err(EncodeError::NullByte);
554        }
555
556        let mut buf = BytesMut::new();
557        buf.extend_from_slice(b"f");
558        let content_len = reason.len() + 1; // +1 for null terminator
559        let len = Self::content_len_to_wire_len(content_len)?;
560        buf.extend_from_slice(&len.to_be_bytes());
561        buf.extend_from_slice(reason.as_bytes());
562        buf.extend_from_slice(&[0]);
563        Ok(buf)
564    }
565
566    /// Fallible Close encoder.
567    pub fn try_encode_close(is_portal: bool, name: &str) -> Result<BytesMut, EncodeError> {
568        if Self::has_nul(name) {
569            return Err(EncodeError::NullByte);
570        }
571
572        let mut buf = BytesMut::new();
573        buf.extend_from_slice(b"C");
574        let content_len = 1 + name.len() + 1; // type + name + null
575        let len = Self::content_len_to_wire_len(content_len)?;
576        buf.extend_from_slice(&len.to_be_bytes());
577        buf.extend_from_slice(&[if is_portal { b'P' } else { b'S' }]);
578        buf.extend_from_slice(name.as_bytes());
579        buf.extend_from_slice(&[0]);
580        Ok(buf)
581    }
582}
583
584// ==================== ULTRA-OPTIMIZED Hot Path Encoders ====================
585//
586// These encoders are designed to beat C:
587// - Direct integer writes (no temp arrays, no bounds checks)
588// - Borrowed slice params (zero-copy)
589// - Single store instructions via BufMut
590//
591
592use bytes::BufMut;
593
594/// Zero-copy parameter for ultra-fast encoding.
595/// Uses borrowed slices to avoid any allocation or copy.
596pub enum Param<'a> {
597    /// SQL NULL value.
598    Null,
599    /// Non-null parameter as a borrowed byte slice.
600    Bytes(&'a [u8]),
601}
602
603impl PgEncoder {
604    /// Direct i32 write - no temp array, no bounds check.
605    /// LLVM emits a single store instruction.
606    #[inline(always)]
607    fn put_i32_be(buf: &mut BytesMut, v: i32) {
608        buf.put_i32(v);
609    }
610
611    #[inline(always)]
612    fn put_i16_be(buf: &mut BytesMut, v: i16) {
613        buf.put_i16(v);
614    }
615
616    /// Encode Bind message - ULTRA OPTIMIZED.
617    /// - Direct integer writes (no temp arrays)
618    /// - Borrowed params (zero-copy)
619    /// - Single allocation check
620    #[inline]
621    pub fn encode_bind_ultra<'a>(
622        buf: &mut BytesMut,
623        statement: &str,
624        params: &[Param<'a>],
625    ) -> Result<(), EncodeError> {
626        Self::encode_bind_ultra_with_result_format(buf, statement, params, Self::FORMAT_TEXT)
627    }
628
629    /// Encode Bind message with explicit result-column format.
630    #[inline]
631    pub fn encode_bind_ultra_with_result_format<'a>(
632        buf: &mut BytesMut,
633        statement: &str,
634        params: &[Param<'a>],
635        result_format: i16,
636    ) -> Result<(), EncodeError> {
637        Self::encode_bind_ultra_with_formats(
638            buf,
639            statement,
640            params,
641            Self::FORMAT_TEXT,
642            result_format,
643        )
644    }
645
646    /// Encode Bind message with explicit parameter and result format codes.
647    #[inline]
648    pub fn encode_bind_ultra_with_formats<'a>(
649        buf: &mut BytesMut,
650        statement: &str,
651        params: &[Param<'a>],
652        param_format: i16,
653        result_format: i16,
654    ) -> Result<(), EncodeError> {
655        if Self::has_nul(statement) {
656            return Err(EncodeError::NullByte);
657        }
658        if params.len() > i16::MAX as usize {
659            return Err(EncodeError::TooManyParameters(params.len()));
660        }
661
662        // Calculate content length upfront
663        let params_size = params.iter().try_fold(0usize, |acc, p| {
664            let field_size = match p {
665                Param::Null => 4usize,
666                Param::Bytes(b) => 4usize
667                    .checked_add(b.len())
668                    .ok_or(EncodeError::MessageTooLarge(usize::MAX))?,
669            };
670            acc.checked_add(field_size)
671                .ok_or(EncodeError::MessageTooLarge(usize::MAX))
672        })?;
673        let param_formats_size = Self::param_format_wire_len(param_format);
674        let result_formats_size = Self::result_format_wire_len(result_format);
675        let content_len = 1usize
676            .checked_add(statement.len())
677            .and_then(|v| v.checked_add(1))
678            .and_then(|v| v.checked_add(2))
679            .and_then(|v| v.checked_add(param_formats_size))
680            .and_then(|v| v.checked_add(params_size))
681            .and_then(|v| v.checked_add(result_formats_size))
682            .ok_or(EncodeError::MessageTooLarge(usize::MAX))?;
683        let wire_len = Self::content_len_to_wire_len(content_len)?;
684
685        // Single reserve - no more allocations
686        buf.reserve(1 + 4 + content_len);
687
688        // Message type 'B'
689        buf.put_u8(b'B');
690
691        // Length (includes itself) - DIRECT WRITE
692        Self::put_i32_be(buf, wire_len);
693
694        // Portal name (empty, null-terminated)
695        buf.put_u8(0);
696
697        // Statement name (null-terminated)
698        buf.extend_from_slice(statement.as_bytes());
699        buf.put_u8(0);
700
701        // Parameter format codes
702        Self::encode_param_formats_bytesmut(buf, param_format);
703
704        // Parameter count
705        let param_count = Self::usize_to_i16(params.len())?;
706        Self::put_i16_be(buf, param_count);
707
708        // Parameters - ZERO COPY from borrowed slices
709        for param in params {
710            match param {
711                Param::Null => Self::put_i32_be(buf, -1),
712                Param::Bytes(data) => {
713                    let data_len = Self::usize_to_i32(data.len())?;
714                    Self::put_i32_be(buf, data_len);
715                    buf.extend_from_slice(data);
716                }
717            }
718        }
719
720        // Result format codes
721        Self::encode_result_formats_bytesmut(buf, result_format);
722        Ok(())
723    }
724
725    /// Encode Execute message - ULTRA OPTIMIZED.
726    #[inline(always)]
727    pub fn encode_execute_ultra(buf: &mut BytesMut) {
728        // Execute: 'E' + len(9) + portal("") + max_rows(0)
729        // = 'E' 00 00 00 09 00 00 00 00 00
730        buf.extend_from_slice(&[b'E', 0, 0, 0, 9, 0, 0, 0, 0, 0]);
731    }
732
733    /// Encode Sync message - ULTRA OPTIMIZED.
734    #[inline(always)]
735    pub fn encode_sync_ultra(buf: &mut BytesMut) {
736        buf.extend_from_slice(&[b'S', 0, 0, 0, 4]);
737    }
738
739    /// Encode Bind message directly into existing buffer (ZERO ALLOCATION).
740    /// This is the hot path optimization - no intermediate Vec allocation.
741    #[inline]
742    pub fn encode_bind_to(
743        buf: &mut BytesMut,
744        statement: &str,
745        params: &[Option<Vec<u8>>],
746    ) -> Result<(), EncodeError> {
747        Self::encode_bind_to_with_result_format(buf, statement, params, Self::FORMAT_TEXT)
748    }
749
750    /// Encode Bind into existing buffer with explicit result-column format.
751    #[inline]
752    pub fn encode_bind_to_with_result_format(
753        buf: &mut BytesMut,
754        statement: &str,
755        params: &[Option<Vec<u8>>],
756        result_format: i16,
757    ) -> Result<(), EncodeError> {
758        Self::encode_bind_to_with_formats(buf, statement, params, Self::FORMAT_TEXT, result_format)
759    }
760
761    /// Encode Bind into existing buffer with explicit parameter and result formats.
762    #[inline]
763    pub fn encode_bind_to_with_formats(
764        buf: &mut BytesMut,
765        statement: &str,
766        params: &[Option<Vec<u8>>],
767        param_format: i16,
768        result_format: i16,
769    ) -> Result<(), EncodeError> {
770        if Self::has_nul(statement) {
771            return Err(EncodeError::NullByte);
772        }
773        if params.len() > i16::MAX as usize {
774            return Err(EncodeError::TooManyParameters(params.len()));
775        }
776
777        // Calculate content length upfront
778        // portal(1) + statement(len+1) + param_formats + param_count(2)
779        // + params_data + result_formats(2 or 4)
780        let params_size = Self::params_wire_len(params)?;
781        let param_formats_size = Self::param_format_wire_len(param_format);
782        let result_formats_size = Self::result_format_wire_len(result_format);
783        let content_len = 1usize
784            .checked_add(statement.len())
785            .and_then(|v| v.checked_add(1))
786            .and_then(|v| v.checked_add(2))
787            .and_then(|v| v.checked_add(param_formats_size))
788            .and_then(|v| v.checked_add(params_size))
789            .and_then(|v| v.checked_add(result_formats_size))
790            .ok_or(EncodeError::MessageTooLarge(usize::MAX))?;
791        let wire_len = Self::content_len_to_wire_len(content_len)?;
792
793        buf.reserve(1 + 4 + content_len);
794
795        // Message type 'B'
796        buf.put_u8(b'B');
797
798        // Length (includes itself) - DIRECT WRITE
799        Self::put_i32_be(buf, wire_len);
800
801        // Portal name (empty, null-terminated)
802        buf.put_u8(0);
803
804        // Statement name (null-terminated)
805        buf.extend_from_slice(statement.as_bytes());
806        buf.put_u8(0);
807
808        // Parameter format codes
809        Self::encode_param_formats_bytesmut(buf, param_format);
810
811        // Parameter count
812        let param_count = Self::usize_to_i16(params.len())?;
813        Self::put_i16_be(buf, param_count);
814
815        // Parameters
816        for param in params {
817            match param {
818                None => Self::put_i32_be(buf, -1),
819                Some(data) => {
820                    let data_len = Self::usize_to_i32(data.len())?;
821                    Self::put_i32_be(buf, data_len);
822                    buf.extend_from_slice(data);
823                }
824            }
825        }
826
827        // Result format codes
828        Self::encode_result_formats_bytesmut(buf, result_format);
829        Ok(())
830    }
831
832    /// Encode Execute message directly into existing buffer (ZERO ALLOCATION).
833    #[inline]
834    pub fn encode_execute_to(buf: &mut BytesMut) {
835        // Content: portal(1) + max_rows(4) = 5 bytes
836        buf.extend_from_slice(&[b'E', 0, 0, 0, 9, 0, 0, 0, 0, 0]);
837    }
838
839    /// Encode Sync message directly into existing buffer (ZERO ALLOCATION).
840    #[inline]
841    pub fn encode_sync_to(buf: &mut BytesMut) {
842        buf.extend_from_slice(&[b'S', 0, 0, 0, 4]);
843    }
844}
845
846#[cfg(test)]
847mod tests {
848    use super::*;
849
850    // NOTE: test_encode_simple_query removed - use AstEncoder instead
851    #[test]
852    fn test_encode_query_string() {
853        let sql = "SELECT 1";
854        let bytes = PgEncoder::try_encode_query_string(sql).unwrap();
855
856        // Message type
857        assert_eq!(bytes[0], b'Q');
858
859        // Length: 4 (length field) + 8 (query) + 1 (null) = 13
860        let len = i32::from_be_bytes([bytes[1], bytes[2], bytes[3], bytes[4]]);
861        assert_eq!(len, 13);
862
863        // Query content
864        assert_eq!(&bytes[5..13], b"SELECT 1");
865
866        // Null terminator
867        assert_eq!(bytes[13], 0);
868    }
869
870    #[test]
871    fn test_encode_terminate() {
872        let bytes = PgEncoder::encode_terminate();
873        assert_eq!(bytes.as_ref(), &[b'X', 0, 0, 0, 4]);
874    }
875
876    #[test]
877    fn test_encode_sync() {
878        let bytes = PgEncoder::encode_sync();
879        assert_eq!(bytes.as_ref(), &[b'S', 0, 0, 0, 4]);
880    }
881
882    #[test]
883    fn test_encode_parse() {
884        let bytes = PgEncoder::try_encode_parse("", "SELECT $1", &[]).unwrap();
885
886        // Message type 'P'
887        assert_eq!(bytes[0], b'P');
888
889        // Content should include query
890        let content = String::from_utf8_lossy(&bytes[5..]);
891        assert!(content.contains("SELECT $1"));
892    }
893
894    #[test]
895    fn test_encode_parse_to_matches_allocate_variant() {
896        let expected = PgEncoder::try_encode_parse("stmt", "SELECT $1::int4", &[23]).unwrap();
897
898        let mut buf = BytesMut::from(&b"prefix"[..]);
899        PgEncoder::try_encode_parse_to(&mut buf, "stmt", "SELECT $1::int4", &[23]).unwrap();
900
901        assert_eq!(&buf[6..], expected.as_ref());
902    }
903
904    #[test]
905    fn test_encode_bind() {
906        let params = vec![
907            Some(b"42".to_vec()),
908            None, // NULL
909        ];
910        let bytes = PgEncoder::encode_bind("", "", &params).unwrap();
911
912        // Message type 'B'
913        assert_eq!(bytes[0], b'B');
914
915        // Should have proper length
916        let len = i32::from_be_bytes([bytes[1], bytes[2], bytes[3], bytes[4]]);
917        assert!(len > 4); // At least header
918    }
919
920    #[test]
921    fn test_encode_bind_binary_result_format() {
922        let bytes =
923            PgEncoder::encode_bind_with_result_format("", "", &[], PgEncoder::FORMAT_BINARY)
924                .unwrap();
925
926        // B + len + portal + statement + param formats + param count + result formats
927        // Result format section for binary should be: count=1, format=1.
928        assert_eq!(&bytes[11..15], &[0, 1, 0, 1]);
929    }
930
931    #[test]
932    fn test_encode_bind_binary_param_and_result_format() {
933        let bytes = PgEncoder::encode_bind_with_formats(
934            "",
935            "",
936            &[],
937            PgEncoder::FORMAT_BINARY,
938            PgEncoder::FORMAT_BINARY,
939        )
940        .unwrap();
941
942        // portal, statement, param formats(count+code), param count, result formats(count+code)
943        assert_eq!(&bytes[7..11], &[0, 1, 0, 1]);
944        assert_eq!(&bytes[11..13], &[0, 0]);
945        assert_eq!(&bytes[13..17], &[0, 1, 0, 1]);
946    }
947
948    #[test]
949    fn test_encode_execute() {
950        let bytes = PgEncoder::try_encode_execute("", 0).unwrap();
951
952        // Message type 'E'
953        assert_eq!(bytes[0], b'E');
954
955        // Length: 4 + 1 (null) + 4 (max_rows) = 9
956        let len = i32::from_be_bytes([bytes[1], bytes[2], bytes[3], bytes[4]]);
957        assert_eq!(len, 9);
958    }
959
960    #[test]
961    fn test_encode_execute_negative_max_rows_returns_error() {
962        let err = PgEncoder::try_encode_execute("", -1).expect_err("must reject negative max_rows");
963        assert_eq!(err, EncodeError::InvalidMaxRows(-1));
964    }
965
966    #[test]
967    fn test_encode_extended_query() {
968        let params = vec![Some(b"hello".to_vec())];
969        let bytes = PgEncoder::encode_extended_query("SELECT $1", &params).unwrap();
970
971        // Should contain all 4 message types: P, B, E, S
972        assert!(bytes.windows(1).any(|w| w == [b'P']));
973        assert!(bytes.windows(1).any(|w| w == [b'B']));
974        assert!(bytes.windows(1).any(|w| w == [b'E']));
975        assert!(bytes.windows(1).any(|w| w == [b'S']));
976    }
977
978    #[test]
979    fn test_encode_extended_query_binary_result_format() {
980        let bytes = PgEncoder::encode_extended_query_with_result_format(
981            "SELECT 1",
982            &[],
983            PgEncoder::FORMAT_BINARY,
984        )
985        .unwrap();
986
987        let parse_len = i32::from_be_bytes([bytes[1], bytes[2], bytes[3], bytes[4]]) as usize;
988        let bind_start = 1 + parse_len;
989        assert_eq!(bytes[bind_start], b'B');
990
991        let bind_len = i32::from_be_bytes([
992            bytes[bind_start + 1],
993            bytes[bind_start + 2],
994            bytes[bind_start + 3],
995            bytes[bind_start + 4],
996        ]);
997        assert_eq!(bind_len, 14);
998
999        let bind_content = &bytes[bind_start + 5..bind_start + 1 + bind_len as usize];
1000        assert_eq!(&bind_content[6..10], &[0, 1, 0, 1]);
1001    }
1002
1003    #[test]
1004    fn test_encode_extended_query_binary_param_and_result_format() {
1005        let bytes = PgEncoder::encode_extended_query_with_formats(
1006            "SELECT 1",
1007            &[],
1008            PgEncoder::FORMAT_BINARY,
1009            PgEncoder::FORMAT_BINARY,
1010        )
1011        .unwrap();
1012
1013        let parse_len = i32::from_be_bytes([bytes[1], bytes[2], bytes[3], bytes[4]]) as usize;
1014        let bind_start = 1 + parse_len;
1015        let bind_len = i32::from_be_bytes([
1016            bytes[bind_start + 1],
1017            bytes[bind_start + 2],
1018            bytes[bind_start + 3],
1019            bytes[bind_start + 4],
1020        ]);
1021        assert_eq!(bind_len, 16);
1022
1023        let bind_content = &bytes[bind_start + 5..bind_start + 1 + bind_len as usize];
1024        assert_eq!(&bind_content[2..6], &[0, 1, 0, 1]);
1025        assert_eq!(&bind_content[6..8], &[0, 0]);
1026        assert_eq!(&bind_content[8..12], &[0, 1, 0, 1]);
1027    }
1028
1029    #[test]
1030    fn test_encode_copy_fail() {
1031        let bytes = PgEncoder::try_encode_copy_fail("bad data").unwrap();
1032        assert_eq!(bytes[0], b'f');
1033        let len = i32::from_be_bytes([bytes[1], bytes[2], bytes[3], bytes[4]]);
1034        assert_eq!(len as usize, 4 + "bad data".len() + 1);
1035        assert_eq!(&bytes[5..13], b"bad data");
1036        assert_eq!(bytes[13], 0);
1037    }
1038
1039    #[test]
1040    fn test_encode_close_statement() {
1041        let bytes = PgEncoder::try_encode_close(false, "my_stmt").unwrap();
1042        assert_eq!(bytes[0], b'C');
1043        assert_eq!(bytes[5], b'S'); // Statement type
1044        assert_eq!(&bytes[6..13], b"my_stmt");
1045        assert_eq!(bytes[13], 0);
1046    }
1047
1048    #[test]
1049    fn test_encode_close_portal() {
1050        let bytes = PgEncoder::try_encode_close(true, "").unwrap();
1051        assert_eq!(bytes[0], b'C');
1052        assert_eq!(bytes[5], b'P'); // Portal type
1053        assert_eq!(bytes[6], 0); // Empty name null terminator
1054    }
1055
1056    #[test]
1057    fn test_encode_parse_too_many_param_types_returns_error() {
1058        let param_types = vec![0u32; (i16::MAX as usize) + 1];
1059        let err =
1060            PgEncoder::try_encode_parse("s", "SELECT 1", &param_types).expect_err("must reject");
1061        assert_eq!(err, EncodeError::TooManyParameters(param_types.len()));
1062    }
1063
1064    #[test]
1065    fn test_encode_parse_to_with_nul_rejected() {
1066        let mut buf = BytesMut::new();
1067        let err =
1068            PgEncoder::try_encode_parse_to(&mut buf, "s", "SELECT 1\0", &[]).expect_err("reject");
1069        assert_eq!(err, EncodeError::NullByte);
1070    }
1071
1072    #[test]
1073    fn test_encode_bind_to_binary_result_format() {
1074        let mut buf = BytesMut::new();
1075        PgEncoder::encode_bind_to_with_result_format(&mut buf, "", &[], PgEncoder::FORMAT_BINARY)
1076            .unwrap();
1077
1078        assert_eq!(&buf[11..15], &[0, 1, 0, 1]);
1079    }
1080
1081    #[test]
1082    fn test_encode_bind_to_binary_param_and_result_format() {
1083        let mut buf = BytesMut::new();
1084        PgEncoder::encode_bind_to_with_formats(
1085            &mut buf,
1086            "",
1087            &[],
1088            PgEncoder::FORMAT_BINARY,
1089            PgEncoder::FORMAT_BINARY,
1090        )
1091        .unwrap();
1092
1093        assert_eq!(&buf[7..11], &[0, 1, 0, 1]);
1094        assert_eq!(&buf[11..13], &[0, 0]);
1095        assert_eq!(&buf[13..17], &[0, 1, 0, 1]);
1096    }
1097
1098    #[test]
1099    fn test_bind_execute_sync_wire_len_matches_encoded_bytes() {
1100        let params = vec![Some(b"abc".to_vec()), None, Some(b"defghi".to_vec())];
1101        let mut buf = BytesMut::new();
1102        PgEncoder::encode_bind_to_with_result_format(
1103            &mut buf,
1104            "stmt",
1105            &params,
1106            PgEncoder::FORMAT_TEXT,
1107        )
1108        .unwrap();
1109        PgEncoder::encode_execute_to(&mut buf);
1110        PgEncoder::encode_sync_to(&mut buf);
1111
1112        let expected = PgEncoder::bind_execute_sync_wire_len_with_formats(
1113            "stmt",
1114            &params,
1115            PgEncoder::FORMAT_TEXT,
1116            PgEncoder::FORMAT_TEXT,
1117        )
1118        .unwrap();
1119        assert_eq!(buf.len(), expected);
1120    }
1121
1122    #[test]
1123    fn test_bind_execute_wire_len_matches_encoded_bytes_binary_formats() {
1124        let params = vec![Some(vec![1, 2, 3, 4]), Some(vec![5, 6])];
1125        let mut buf = BytesMut::new();
1126        PgEncoder::encode_bind_to_with_formats(
1127            &mut buf,
1128            "stmt",
1129            &params,
1130            PgEncoder::FORMAT_BINARY,
1131            PgEncoder::FORMAT_BINARY,
1132        )
1133        .unwrap();
1134        PgEncoder::encode_execute_to(&mut buf);
1135
1136        let expected = PgEncoder::bind_execute_wire_len_with_formats(
1137            "stmt",
1138            &params,
1139            PgEncoder::FORMAT_BINARY,
1140            PgEncoder::FORMAT_BINARY,
1141        )
1142        .unwrap();
1143        assert_eq!(buf.len(), expected);
1144    }
1145
1146    #[test]
1147    fn test_encode_bind_ultra_binary_result_format() {
1148        let mut buf = BytesMut::new();
1149        PgEncoder::encode_bind_ultra_with_result_format(
1150            &mut buf,
1151            "",
1152            &[],
1153            PgEncoder::FORMAT_BINARY,
1154        )
1155        .unwrap();
1156
1157        assert_eq!(&buf[11..15], &[0, 1, 0, 1]);
1158    }
1159
1160    #[test]
1161    fn test_encode_bind_ultra_binary_param_and_result_format() {
1162        let mut buf = BytesMut::new();
1163        PgEncoder::encode_bind_ultra_with_formats(
1164            &mut buf,
1165            "",
1166            &[],
1167            PgEncoder::FORMAT_BINARY,
1168            PgEncoder::FORMAT_BINARY,
1169        )
1170        .unwrap();
1171
1172        assert_eq!(&buf[7..11], &[0, 1, 0, 1]);
1173        assert_eq!(&buf[11..13], &[0, 0]);
1174        assert_eq!(&buf[13..17], &[0, 1, 0, 1]);
1175    }
1176
1177    #[test]
1178    fn test_encode_query_string_with_nul_returns_empty() {
1179        let err =
1180            PgEncoder::try_encode_query_string("select 1\0select 2").expect_err("must reject NUL");
1181        assert_eq!(err, EncodeError::NullByte);
1182    }
1183
1184    #[test]
1185    fn test_encode_parse_with_nul_returns_empty() {
1186        let err = PgEncoder::try_encode_parse("s", "SELECT 1\0", &[]).expect_err("must reject");
1187        assert_eq!(err, EncodeError::NullByte);
1188    }
1189
1190    #[test]
1191    fn test_encode_bind_with_nul_rejected() {
1192        let err = PgEncoder::encode_bind_with_result_format("\0", "", &[], PgEncoder::FORMAT_TEXT)
1193            .expect_err("bind with NUL portal must fail");
1194        assert_eq!(err, EncodeError::NullByte);
1195    }
1196
1197    #[test]
1198    fn test_encode_extended_query_with_nul_rejected() {
1199        let err = PgEncoder::encode_extended_query_with_result_format(
1200            "SELECT 1\0UNION SELECT 2",
1201            &[],
1202            PgEncoder::FORMAT_TEXT,
1203        )
1204        .expect_err("extended query with NUL SQL must fail");
1205        assert_eq!(err, EncodeError::NullByte);
1206    }
1207
1208    #[test]
1209    fn test_encode_copy_fail_with_nul_returns_empty() {
1210        let err = PgEncoder::try_encode_copy_fail("bad\0data").expect_err("must reject");
1211        assert_eq!(err, EncodeError::NullByte);
1212    }
1213}