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 result_format_wire_len(result_format: i16) -> usize {
29        if result_format == Self::FORMAT_TEXT {
30            2 // result format count = 0
31        } else {
32            4 // result format count = 1 + one format code
33        }
34    }
35
36    #[inline(always)]
37    fn encode_result_formats_vec(content: &mut Vec<u8>, result_format: i16) {
38        if result_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(&result_format.to_be_bytes());
43        }
44    }
45
46    #[inline(always)]
47    fn encode_result_formats_bytesmut(buf: &mut BytesMut, result_format: i16) {
48        if result_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(&result_format.to_be_bytes());
53        }
54    }
55
56    #[inline(always)]
57    fn content_len_to_wire_len(content_len: usize) -> Result<i32, EncodeError> {
58        let total = content_len
59            .checked_add(4)
60            .ok_or(EncodeError::MessageTooLarge(usize::MAX))?;
61        i32::try_from(total).map_err(|_| EncodeError::MessageTooLarge(total))
62    }
63
64    #[inline(always)]
65    fn usize_to_i16(n: usize) -> Result<i16, EncodeError> {
66        i16::try_from(n).map_err(|_| EncodeError::TooManyParameters(n))
67    }
68
69    #[inline(always)]
70    fn usize_to_i32(n: usize) -> Result<i32, EncodeError> {
71        i32::try_from(n).map_err(|_| EncodeError::MessageTooLarge(n))
72    }
73
74    #[inline(always)]
75    fn has_nul(s: &str) -> bool {
76        s.as_bytes().contains(&0)
77    }
78
79    /// Fallible simple-query encoder.
80    pub fn try_encode_query_string(sql: &str) -> Result<BytesMut, EncodeError> {
81        if Self::has_nul(sql) {
82            return Err(EncodeError::NullByte);
83        }
84
85        let mut buf = BytesMut::new();
86        let content_len = sql.len() + 1; // +1 for null terminator
87        let total_len = Self::content_len_to_wire_len(content_len)?;
88
89        buf.extend_from_slice(b"Q");
90        buf.extend_from_slice(&total_len.to_be_bytes());
91        buf.extend_from_slice(sql.as_bytes());
92        buf.extend_from_slice(&[0]);
93        Ok(buf)
94    }
95
96    /// Encode a Terminate message to close the connection.
97    pub fn encode_terminate() -> BytesMut {
98        let mut buf = BytesMut::new();
99        buf.extend_from_slice(&[b'X', 0, 0, 0, 4]);
100        buf
101    }
102
103    /// Encode a Sync message (end of pipeline in extended query protocol).
104    pub fn encode_sync() -> BytesMut {
105        let mut buf = BytesMut::new();
106        buf.extend_from_slice(&[b'S', 0, 0, 0, 4]);
107        buf
108    }
109
110    // ==================== Extended Query Protocol ====================
111
112    /// Fallible Parse message encoder.
113    pub fn try_encode_parse(
114        name: &str,
115        sql: &str,
116        param_types: &[u32],
117    ) -> Result<BytesMut, EncodeError> {
118        if Self::has_nul(name) || Self::has_nul(sql) {
119            return Err(EncodeError::NullByte);
120        }
121        if param_types.len() > i16::MAX as usize {
122            return Err(EncodeError::TooManyParameters(param_types.len()));
123        }
124
125        let mut buf = BytesMut::new();
126        buf.extend_from_slice(b"P");
127
128        let mut content = Vec::new();
129        content.extend_from_slice(name.as_bytes());
130        content.push(0);
131        content.extend_from_slice(sql.as_bytes());
132        content.push(0);
133        let param_count = Self::usize_to_i16(param_types.len())?;
134        content.extend_from_slice(&param_count.to_be_bytes());
135        for &oid in param_types {
136            content.extend_from_slice(&oid.to_be_bytes());
137        }
138
139        let len = Self::content_len_to_wire_len(content.len())?;
140        buf.extend_from_slice(&len.to_be_bytes());
141        buf.extend_from_slice(&content);
142        Ok(buf)
143    }
144
145    /// Encode a Bind message (bind parameters to a prepared statement).
146    /// Wire format:
147    /// - 'B' (1 byte) - message type
148    /// - length (4 bytes)
149    /// - portal name (null-terminated)
150    /// - statement name (null-terminated)
151    /// - format code count (2 bytes) - we use 0 (all text)
152    /// - parameter count (2 bytes)
153    /// - for each parameter: length (4 bytes, -1 for NULL), data
154    /// - result format count + codes
155    ///
156    /// # Arguments
157    ///
158    /// * `portal` — Destination portal name (empty string for unnamed).
159    /// * `statement` — Source prepared statement name (empty string for unnamed).
160    /// * `params` — Parameter values; `None` entries encode as SQL NULL.
161    pub fn encode_bind(
162        portal: &str,
163        statement: &str,
164        params: &[Option<Vec<u8>>],
165    ) -> Result<BytesMut, EncodeError> {
166        Self::encode_bind_with_result_format(portal, statement, params, Self::FORMAT_TEXT)
167    }
168
169    /// Encode a Bind message with explicit result-column format.
170    ///
171    /// `result_format` is PostgreSQL wire format code: `0 = text`, `1 = binary`.
172    /// For `0`, this encodes "result format count = 0" (server default text).
173    /// For non-zero codes, this encodes one explicit result format code.
174    pub fn encode_bind_with_result_format(
175        portal: &str,
176        statement: &str,
177        params: &[Option<Vec<u8>>],
178        result_format: i16,
179    ) -> Result<BytesMut, EncodeError> {
180        if Self::has_nul(portal) || Self::has_nul(statement) {
181            return Err(EncodeError::NullByte);
182        }
183        if params.len() > i16::MAX as usize {
184            return Err(EncodeError::TooManyParameters(params.len()));
185        }
186
187        let mut buf = BytesMut::new();
188
189        // Message type 'B'
190        buf.extend_from_slice(b"B");
191
192        let mut content = Vec::new();
193
194        // Portal name (null-terminated)
195        content.extend_from_slice(portal.as_bytes());
196        content.push(0);
197
198        // Statement name (null-terminated)
199        content.extend_from_slice(statement.as_bytes());
200        content.push(0);
201
202        // Format codes count (0 = use default text format)
203        content.extend_from_slice(&0i16.to_be_bytes());
204
205        // Parameter count
206        let param_count = Self::usize_to_i16(params.len())?;
207        content.extend_from_slice(&param_count.to_be_bytes());
208
209        // Parameters
210        for param in params {
211            match param {
212                None => {
213                    // NULL: length = -1
214                    content.extend_from_slice(&(-1i32).to_be_bytes());
215                }
216                Some(data) => {
217                    let data_len = Self::usize_to_i32(data.len())?;
218                    content.extend_from_slice(&data_len.to_be_bytes());
219                    content.extend_from_slice(data);
220                }
221            }
222        }
223
224        // Result format codes: default text (count=0) or explicit code.
225        Self::encode_result_formats_vec(&mut content, result_format);
226
227        // Length
228        let len = Self::content_len_to_wire_len(content.len())?;
229        buf.extend_from_slice(&len.to_be_bytes());
230        buf.extend_from_slice(&content);
231
232        Ok(buf)
233    }
234
235    /// Fallible Execute message encoder.
236    pub fn try_encode_execute(portal: &str, max_rows: i32) -> Result<BytesMut, EncodeError> {
237        if Self::has_nul(portal) {
238            return Err(EncodeError::NullByte);
239        }
240        if max_rows < 0 {
241            return Err(EncodeError::InvalidMaxRows(max_rows));
242        }
243
244        let mut buf = BytesMut::new();
245        buf.extend_from_slice(b"E");
246
247        let mut content = Vec::new();
248        content.extend_from_slice(portal.as_bytes());
249        content.push(0);
250        content.extend_from_slice(&max_rows.to_be_bytes());
251
252        let len = Self::content_len_to_wire_len(content.len())?;
253        buf.extend_from_slice(&len.to_be_bytes());
254        buf.extend_from_slice(&content);
255        Ok(buf)
256    }
257
258    /// Fallible Describe message encoder.
259    pub fn try_encode_describe(is_portal: bool, name: &str) -> Result<BytesMut, EncodeError> {
260        if Self::has_nul(name) {
261            return Err(EncodeError::NullByte);
262        }
263
264        let mut buf = BytesMut::new();
265        buf.extend_from_slice(b"D");
266
267        let mut content = Vec::new();
268        content.push(if is_portal { b'P' } else { b'S' });
269        content.extend_from_slice(name.as_bytes());
270        content.push(0);
271
272        let len = Self::content_len_to_wire_len(content.len())?;
273        buf.extend_from_slice(&len.to_be_bytes());
274        buf.extend_from_slice(&content);
275        Ok(buf)
276    }
277
278    /// Encode a complete extended query pipeline (OPTIMIZED).
279    /// This combines Parse + Bind + Execute + Sync in a single buffer.
280    /// Zero intermediate allocations - writes directly to pre-sized BytesMut.
281    pub fn encode_extended_query(
282        sql: &str,
283        params: &[Option<Vec<u8>>],
284    ) -> Result<BytesMut, EncodeError> {
285        Self::encode_extended_query_with_result_format(sql, params, Self::FORMAT_TEXT)
286    }
287
288    /// Encode a complete extended query pipeline with explicit result format.
289    ///
290    /// `result_format` is PostgreSQL wire format code: `0 = text`, `1 = binary`.
291    pub fn encode_extended_query_with_result_format(
292        sql: &str,
293        params: &[Option<Vec<u8>>],
294        result_format: i16,
295    ) -> Result<BytesMut, EncodeError> {
296        if Self::has_nul(sql) {
297            return Err(EncodeError::NullByte);
298        }
299        if params.len() > i16::MAX as usize {
300            return Err(EncodeError::TooManyParameters(params.len()));
301        }
302
303        // Calculate total size upfront to avoid reallocations
304        // Bind: 1 + 4 + 1 + 1 + 2 + 2 + params_data + result_formats
305        // Execute: 1 + 4 + 1 + 4 = 10
306        // Sync: 5
307        let params_size = params.iter().try_fold(0usize, |acc, p| {
308            let field_size = 4usize
309                .checked_add(p.as_ref().map_or(0usize, |v| v.len()))
310                .ok_or(EncodeError::MessageTooLarge(usize::MAX))?;
311            acc.checked_add(field_size)
312                .ok_or(EncodeError::MessageTooLarge(usize::MAX))
313        })?;
314        let result_formats_size = Self::result_format_wire_len(result_format);
315        let total_size = 9usize
316            .checked_add(sql.len())
317            .and_then(|v| v.checked_add(11))
318            .and_then(|v| v.checked_add(params_size))
319            .and_then(|v| v.checked_add(result_formats_size))
320            .and_then(|v| v.checked_add(10))
321            .and_then(|v| v.checked_add(5))
322            .ok_or(EncodeError::MessageTooLarge(usize::MAX))?;
323
324        let mut buf = BytesMut::with_capacity(total_size);
325
326        // ===== PARSE =====
327        buf.extend_from_slice(b"P");
328        let parse_content_len = 1usize
329            .checked_add(sql.len())
330            .and_then(|v| v.checked_add(1))
331            .and_then(|v| v.checked_add(2))
332            .ok_or(EncodeError::MessageTooLarge(usize::MAX))?;
333        let parse_len = Self::content_len_to_wire_len(parse_content_len)?;
334        buf.extend_from_slice(&parse_len.to_be_bytes());
335        buf.extend_from_slice(&[0]); // Unnamed statement
336        buf.extend_from_slice(sql.as_bytes());
337        buf.extend_from_slice(&[0]); // Null terminator
338        buf.extend_from_slice(&0i16.to_be_bytes()); // No param types (infer)
339
340        // ===== BIND =====
341        buf.extend_from_slice(b"B");
342        let bind_content_len = 1usize
343            .checked_add(1)
344            .and_then(|v| v.checked_add(2))
345            .and_then(|v| v.checked_add(2))
346            .and_then(|v| v.checked_add(params_size))
347            .and_then(|v| v.checked_add(result_formats_size))
348            .ok_or(EncodeError::MessageTooLarge(usize::MAX))?;
349        let bind_len = Self::content_len_to_wire_len(bind_content_len)?;
350        buf.extend_from_slice(&bind_len.to_be_bytes());
351        buf.extend_from_slice(&[0]); // Unnamed portal
352        buf.extend_from_slice(&[0]); // Unnamed statement
353        buf.extend_from_slice(&0i16.to_be_bytes()); // Format codes (default text)
354        let param_count = Self::usize_to_i16(params.len())?;
355        buf.extend_from_slice(&param_count.to_be_bytes());
356        for param in params {
357            match param {
358                None => buf.extend_from_slice(&(-1i32).to_be_bytes()),
359                Some(data) => {
360                    let data_len = Self::usize_to_i32(data.len())?;
361                    buf.extend_from_slice(&data_len.to_be_bytes());
362                    buf.extend_from_slice(data);
363                }
364            }
365        }
366        Self::encode_result_formats_bytesmut(&mut buf, result_format);
367
368        // ===== EXECUTE =====
369        buf.extend_from_slice(b"E");
370        buf.extend_from_slice(&9i32.to_be_bytes()); // len = 4 + 1 + 4
371        buf.extend_from_slice(&[0]); // Unnamed portal
372        buf.extend_from_slice(&0i32.to_be_bytes()); // Unlimited rows
373
374        // ===== SYNC =====
375        buf.extend_from_slice(&[b'S', 0, 0, 0, 4]);
376
377        Ok(buf)
378    }
379
380    /// Fallible CopyFail encoder.
381    pub fn try_encode_copy_fail(reason: &str) -> Result<BytesMut, EncodeError> {
382        if Self::has_nul(reason) {
383            return Err(EncodeError::NullByte);
384        }
385
386        let mut buf = BytesMut::new();
387        buf.extend_from_slice(b"f");
388        let content_len = reason.len() + 1; // +1 for null terminator
389        let len = Self::content_len_to_wire_len(content_len)?;
390        buf.extend_from_slice(&len.to_be_bytes());
391        buf.extend_from_slice(reason.as_bytes());
392        buf.extend_from_slice(&[0]);
393        Ok(buf)
394    }
395
396    /// Fallible Close encoder.
397    pub fn try_encode_close(is_portal: bool, name: &str) -> Result<BytesMut, EncodeError> {
398        if Self::has_nul(name) {
399            return Err(EncodeError::NullByte);
400        }
401
402        let mut buf = BytesMut::new();
403        buf.extend_from_slice(b"C");
404        let content_len = 1 + name.len() + 1; // type + name + null
405        let len = Self::content_len_to_wire_len(content_len)?;
406        buf.extend_from_slice(&len.to_be_bytes());
407        buf.extend_from_slice(&[if is_portal { b'P' } else { b'S' }]);
408        buf.extend_from_slice(name.as_bytes());
409        buf.extend_from_slice(&[0]);
410        Ok(buf)
411    }
412}
413
414// ==================== ULTRA-OPTIMIZED Hot Path Encoders ====================
415//
416// These encoders are designed to beat C:
417// - Direct integer writes (no temp arrays, no bounds checks)
418// - Borrowed slice params (zero-copy)
419// - Single store instructions via BufMut
420//
421
422use bytes::BufMut;
423
424/// Zero-copy parameter for ultra-fast encoding.
425/// Uses borrowed slices to avoid any allocation or copy.
426pub enum Param<'a> {
427    /// SQL NULL value.
428    Null,
429    /// Non-null parameter as a borrowed byte slice.
430    Bytes(&'a [u8]),
431}
432
433impl PgEncoder {
434    /// Direct i32 write - no temp array, no bounds check.
435    /// LLVM emits a single store instruction.
436    #[inline(always)]
437    fn put_i32_be(buf: &mut BytesMut, v: i32) {
438        buf.put_i32(v);
439    }
440
441    #[inline(always)]
442    fn put_i16_be(buf: &mut BytesMut, v: i16) {
443        buf.put_i16(v);
444    }
445
446    /// Encode Bind message - ULTRA OPTIMIZED.
447    /// - Direct integer writes (no temp arrays)
448    /// - Borrowed params (zero-copy)
449    /// - Single allocation check
450    #[inline]
451    pub fn encode_bind_ultra<'a>(
452        buf: &mut BytesMut,
453        statement: &str,
454        params: &[Param<'a>],
455    ) -> Result<(), EncodeError> {
456        Self::encode_bind_ultra_with_result_format(buf, statement, params, Self::FORMAT_TEXT)
457    }
458
459    /// Encode Bind message with explicit result-column format.
460    #[inline]
461    pub fn encode_bind_ultra_with_result_format<'a>(
462        buf: &mut BytesMut,
463        statement: &str,
464        params: &[Param<'a>],
465        result_format: i16,
466    ) -> Result<(), EncodeError> {
467        if Self::has_nul(statement) {
468            return Err(EncodeError::NullByte);
469        }
470        if params.len() > i16::MAX as usize {
471            return Err(EncodeError::TooManyParameters(params.len()));
472        }
473
474        // Calculate content length upfront
475        let params_size = params.iter().try_fold(0usize, |acc, p| {
476            let field_size = match p {
477                Param::Null => 4usize,
478                Param::Bytes(b) => 4usize
479                    .checked_add(b.len())
480                    .ok_or(EncodeError::MessageTooLarge(usize::MAX))?,
481            };
482            acc.checked_add(field_size)
483                .ok_or(EncodeError::MessageTooLarge(usize::MAX))
484        })?;
485        let result_formats_size = Self::result_format_wire_len(result_format);
486        let content_len = 1usize
487            .checked_add(statement.len())
488            .and_then(|v| v.checked_add(1))
489            .and_then(|v| v.checked_add(2))
490            .and_then(|v| v.checked_add(2))
491            .and_then(|v| v.checked_add(params_size))
492            .and_then(|v| v.checked_add(result_formats_size))
493            .ok_or(EncodeError::MessageTooLarge(usize::MAX))?;
494        let wire_len = Self::content_len_to_wire_len(content_len)?;
495
496        // Single reserve - no more allocations
497        buf.reserve(1 + 4 + content_len);
498
499        // Message type 'B'
500        buf.put_u8(b'B');
501
502        // Length (includes itself) - DIRECT WRITE
503        Self::put_i32_be(buf, wire_len);
504
505        // Portal name (empty, null-terminated)
506        buf.put_u8(0);
507
508        // Statement name (null-terminated)
509        buf.extend_from_slice(statement.as_bytes());
510        buf.put_u8(0);
511
512        // Format codes count (0 = default text)
513        Self::put_i16_be(buf, 0);
514
515        // Parameter count
516        let param_count = Self::usize_to_i16(params.len())?;
517        Self::put_i16_be(buf, param_count);
518
519        // Parameters - ZERO COPY from borrowed slices
520        for param in params {
521            match param {
522                Param::Null => Self::put_i32_be(buf, -1),
523                Param::Bytes(data) => {
524                    let data_len = Self::usize_to_i32(data.len())?;
525                    Self::put_i32_be(buf, data_len);
526                    buf.extend_from_slice(data);
527                }
528            }
529        }
530
531        // Result format codes
532        Self::encode_result_formats_bytesmut(buf, result_format);
533        Ok(())
534    }
535
536    /// Encode Execute message - ULTRA OPTIMIZED.
537    #[inline(always)]
538    pub fn encode_execute_ultra(buf: &mut BytesMut) {
539        // Execute: 'E' + len(9) + portal("") + max_rows(0)
540        // = 'E' 00 00 00 09 00 00 00 00 00
541        buf.extend_from_slice(&[b'E', 0, 0, 0, 9, 0, 0, 0, 0, 0]);
542    }
543
544    /// Encode Sync message - ULTRA OPTIMIZED.
545    #[inline(always)]
546    pub fn encode_sync_ultra(buf: &mut BytesMut) {
547        buf.extend_from_slice(&[b'S', 0, 0, 0, 4]);
548    }
549
550    // Keep the original methods for compatibility
551
552    /// Encode Bind message directly into existing buffer (ZERO ALLOCATION).
553    /// This is the hot path optimization - no intermediate Vec allocation.
554    #[inline]
555    pub fn encode_bind_to(
556        buf: &mut BytesMut,
557        statement: &str,
558        params: &[Option<Vec<u8>>],
559    ) -> Result<(), EncodeError> {
560        Self::encode_bind_to_with_result_format(buf, statement, params, Self::FORMAT_TEXT)
561    }
562
563    /// Encode Bind into existing buffer with explicit result-column format.
564    #[inline]
565    pub fn encode_bind_to_with_result_format(
566        buf: &mut BytesMut,
567        statement: &str,
568        params: &[Option<Vec<u8>>],
569        result_format: i16,
570    ) -> Result<(), EncodeError> {
571        if Self::has_nul(statement) {
572            return Err(EncodeError::NullByte);
573        }
574        if params.len() > i16::MAX as usize {
575            return Err(EncodeError::TooManyParameters(params.len()));
576        }
577
578        // Calculate content length upfront
579        // portal(1) + statement(len+1) + format_codes(2) + param_count(2)
580        // + params_data + result_formats(2 or 4)
581        let params_size = params.iter().try_fold(0usize, |acc, p| {
582            let field_size = 4usize
583                .checked_add(p.as_ref().map_or(0usize, |v| v.len()))
584                .ok_or(EncodeError::MessageTooLarge(usize::MAX))?;
585            acc.checked_add(field_size)
586                .ok_or(EncodeError::MessageTooLarge(usize::MAX))
587        })?;
588        let result_formats_size = Self::result_format_wire_len(result_format);
589        let content_len = 1usize
590            .checked_add(statement.len())
591            .and_then(|v| v.checked_add(1))
592            .and_then(|v| v.checked_add(2))
593            .and_then(|v| v.checked_add(2))
594            .and_then(|v| v.checked_add(params_size))
595            .and_then(|v| v.checked_add(result_formats_size))
596            .ok_or(EncodeError::MessageTooLarge(usize::MAX))?;
597        let wire_len = Self::content_len_to_wire_len(content_len)?;
598
599        buf.reserve(1 + 4 + content_len);
600
601        // Message type 'B'
602        buf.put_u8(b'B');
603
604        // Length (includes itself) - DIRECT WRITE
605        Self::put_i32_be(buf, wire_len);
606
607        // Portal name (empty, null-terminated)
608        buf.put_u8(0);
609
610        // Statement name (null-terminated)
611        buf.extend_from_slice(statement.as_bytes());
612        buf.put_u8(0);
613
614        // Format codes count (0 = default text)
615        Self::put_i16_be(buf, 0);
616
617        // Parameter count
618        let param_count = Self::usize_to_i16(params.len())?;
619        Self::put_i16_be(buf, param_count);
620
621        // Parameters
622        for param in params {
623            match param {
624                None => Self::put_i32_be(buf, -1),
625                Some(data) => {
626                    let data_len = Self::usize_to_i32(data.len())?;
627                    Self::put_i32_be(buf, data_len);
628                    buf.extend_from_slice(data);
629                }
630            }
631        }
632
633        // Result format codes
634        Self::encode_result_formats_bytesmut(buf, result_format);
635        Ok(())
636    }
637
638    /// Encode Execute message directly into existing buffer (ZERO ALLOCATION).
639    #[inline]
640    pub fn encode_execute_to(buf: &mut BytesMut) {
641        // Content: portal(1) + max_rows(4) = 5 bytes
642        buf.extend_from_slice(&[b'E', 0, 0, 0, 9, 0, 0, 0, 0, 0]);
643    }
644
645    /// Encode Sync message directly into existing buffer (ZERO ALLOCATION).
646    #[inline]
647    pub fn encode_sync_to(buf: &mut BytesMut) {
648        buf.extend_from_slice(&[b'S', 0, 0, 0, 4]);
649    }
650}
651
652#[cfg(test)]
653mod tests {
654    use super::*;
655
656    // NOTE: test_encode_simple_query removed - use AstEncoder instead
657    #[test]
658    fn test_encode_query_string() {
659        let sql = "SELECT 1";
660        let bytes = PgEncoder::try_encode_query_string(sql).unwrap();
661
662        // Message type
663        assert_eq!(bytes[0], b'Q');
664
665        // Length: 4 (length field) + 8 (query) + 1 (null) = 13
666        let len = i32::from_be_bytes([bytes[1], bytes[2], bytes[3], bytes[4]]);
667        assert_eq!(len, 13);
668
669        // Query content
670        assert_eq!(&bytes[5..13], b"SELECT 1");
671
672        // Null terminator
673        assert_eq!(bytes[13], 0);
674    }
675
676    #[test]
677    fn test_encode_terminate() {
678        let bytes = PgEncoder::encode_terminate();
679        assert_eq!(bytes.as_ref(), &[b'X', 0, 0, 0, 4]);
680    }
681
682    #[test]
683    fn test_encode_sync() {
684        let bytes = PgEncoder::encode_sync();
685        assert_eq!(bytes.as_ref(), &[b'S', 0, 0, 0, 4]);
686    }
687
688    #[test]
689    fn test_encode_parse() {
690        let bytes = PgEncoder::try_encode_parse("", "SELECT $1", &[]).unwrap();
691
692        // Message type 'P'
693        assert_eq!(bytes[0], b'P');
694
695        // Content should include query
696        let content = String::from_utf8_lossy(&bytes[5..]);
697        assert!(content.contains("SELECT $1"));
698    }
699
700    #[test]
701    fn test_encode_bind() {
702        let params = vec![
703            Some(b"42".to_vec()),
704            None, // NULL
705        ];
706        let bytes = PgEncoder::encode_bind("", "", &params).unwrap();
707
708        // Message type 'B'
709        assert_eq!(bytes[0], b'B');
710
711        // Should have proper length
712        let len = i32::from_be_bytes([bytes[1], bytes[2], bytes[3], bytes[4]]);
713        assert!(len > 4); // At least header
714    }
715
716    #[test]
717    fn test_encode_bind_binary_result_format() {
718        let bytes =
719            PgEncoder::encode_bind_with_result_format("", "", &[], PgEncoder::FORMAT_BINARY)
720                .unwrap();
721
722        // B + len + portal + statement + param formats + param count + result formats
723        // Result format section for binary should be: count=1, format=1.
724        assert_eq!(&bytes[11..15], &[0, 1, 0, 1]);
725    }
726
727    #[test]
728    fn test_encode_execute() {
729        let bytes = PgEncoder::try_encode_execute("", 0).unwrap();
730
731        // Message type 'E'
732        assert_eq!(bytes[0], b'E');
733
734        // Length: 4 + 1 (null) + 4 (max_rows) = 9
735        let len = i32::from_be_bytes([bytes[1], bytes[2], bytes[3], bytes[4]]);
736        assert_eq!(len, 9);
737    }
738
739    #[test]
740    fn test_encode_execute_negative_max_rows_returns_error() {
741        let err = PgEncoder::try_encode_execute("", -1).expect_err("must reject negative max_rows");
742        assert_eq!(err, EncodeError::InvalidMaxRows(-1));
743    }
744
745    #[test]
746    fn test_encode_extended_query() {
747        let params = vec![Some(b"hello".to_vec())];
748        let bytes = PgEncoder::encode_extended_query("SELECT $1", &params).unwrap();
749
750        // Should contain all 4 message types: P, B, E, S
751        assert!(bytes.windows(1).any(|w| w == [b'P']));
752        assert!(bytes.windows(1).any(|w| w == [b'B']));
753        assert!(bytes.windows(1).any(|w| w == [b'E']));
754        assert!(bytes.windows(1).any(|w| w == [b'S']));
755    }
756
757    #[test]
758    fn test_encode_extended_query_binary_result_format() {
759        let bytes = PgEncoder::encode_extended_query_with_result_format(
760            "SELECT 1",
761            &[],
762            PgEncoder::FORMAT_BINARY,
763        )
764        .unwrap();
765
766        let parse_len = i32::from_be_bytes([bytes[1], bytes[2], bytes[3], bytes[4]]) as usize;
767        let bind_start = 1 + parse_len;
768        assert_eq!(bytes[bind_start], b'B');
769
770        let bind_len = i32::from_be_bytes([
771            bytes[bind_start + 1],
772            bytes[bind_start + 2],
773            bytes[bind_start + 3],
774            bytes[bind_start + 4],
775        ]);
776        assert_eq!(bind_len, 14);
777
778        let bind_content = &bytes[bind_start + 5..bind_start + 1 + bind_len as usize];
779        assert_eq!(&bind_content[6..10], &[0, 1, 0, 1]);
780    }
781
782    #[test]
783    fn test_encode_copy_fail() {
784        let bytes = PgEncoder::try_encode_copy_fail("bad data").unwrap();
785        assert_eq!(bytes[0], b'f');
786        let len = i32::from_be_bytes([bytes[1], bytes[2], bytes[3], bytes[4]]);
787        assert_eq!(len as usize, 4 + "bad data".len() + 1);
788        assert_eq!(&bytes[5..13], b"bad data");
789        assert_eq!(bytes[13], 0);
790    }
791
792    #[test]
793    fn test_encode_close_statement() {
794        let bytes = PgEncoder::try_encode_close(false, "my_stmt").unwrap();
795        assert_eq!(bytes[0], b'C');
796        assert_eq!(bytes[5], b'S'); // Statement type
797        assert_eq!(&bytes[6..13], b"my_stmt");
798        assert_eq!(bytes[13], 0);
799    }
800
801    #[test]
802    fn test_encode_close_portal() {
803        let bytes = PgEncoder::try_encode_close(true, "").unwrap();
804        assert_eq!(bytes[0], b'C');
805        assert_eq!(bytes[5], b'P'); // Portal type
806        assert_eq!(bytes[6], 0); // Empty name null terminator
807    }
808
809    #[test]
810    fn test_encode_parse_too_many_param_types_returns_error() {
811        let param_types = vec![0u32; (i16::MAX as usize) + 1];
812        let err =
813            PgEncoder::try_encode_parse("s", "SELECT 1", &param_types).expect_err("must reject");
814        assert_eq!(err, EncodeError::TooManyParameters(param_types.len()));
815    }
816
817    #[test]
818    fn test_encode_bind_to_binary_result_format() {
819        let mut buf = BytesMut::new();
820        PgEncoder::encode_bind_to_with_result_format(&mut buf, "", &[], PgEncoder::FORMAT_BINARY)
821            .unwrap();
822
823        assert_eq!(&buf[11..15], &[0, 1, 0, 1]);
824    }
825
826    #[test]
827    fn test_encode_bind_ultra_binary_result_format() {
828        let mut buf = BytesMut::new();
829        PgEncoder::encode_bind_ultra_with_result_format(
830            &mut buf,
831            "",
832            &[],
833            PgEncoder::FORMAT_BINARY,
834        )
835        .unwrap();
836
837        assert_eq!(&buf[11..15], &[0, 1, 0, 1]);
838    }
839
840    #[test]
841    fn test_encode_query_string_with_nul_returns_empty() {
842        let err =
843            PgEncoder::try_encode_query_string("select 1\0select 2").expect_err("must reject NUL");
844        assert_eq!(err, EncodeError::NullByte);
845    }
846
847    #[test]
848    fn test_encode_parse_with_nul_returns_empty() {
849        let err = PgEncoder::try_encode_parse("s", "SELECT 1\0", &[]).expect_err("must reject");
850        assert_eq!(err, EncodeError::NullByte);
851    }
852
853    #[test]
854    fn test_encode_bind_with_nul_rejected() {
855        let err = PgEncoder::encode_bind_with_result_format("\0", "", &[], PgEncoder::FORMAT_TEXT)
856            .expect_err("bind with NUL portal must fail");
857        assert_eq!(err, EncodeError::NullByte);
858    }
859
860    #[test]
861    fn test_encode_extended_query_with_nul_rejected() {
862        let err = PgEncoder::encode_extended_query_with_result_format(
863            "SELECT 1\0UNION SELECT 2",
864            &[],
865            PgEncoder::FORMAT_TEXT,
866        )
867        .expect_err("extended query with NUL SQL must fail");
868        assert_eq!(err, EncodeError::NullByte);
869    }
870
871    #[test]
872    fn test_encode_copy_fail_with_nul_returns_empty() {
873        let err = PgEncoder::try_encode_copy_fail("bad\0data").expect_err("must reject");
874        assert_eq!(err, EncodeError::NullByte);
875    }
876}