turbo_quant/
wire.rs

1//! Deterministic compact wire encoding for [`TurboCode`].
2//!
3//! The wire bytes are a derived acceleration artifact. They are bound to the
4//! quantizer profile on decode and are never authoritative over raw f32 vectors.
5
6use crate::{
7    bitpack,
8    error::{Result, TurboQuantError},
9    polar::PolarCode,
10    qjl::QjlSketch,
11    rotation::RotationKind,
12    turbo::{TurboCode, TurboMode, TurboQuantizer},
13};
14
15/// Magic bytes for TurboCode wire format v1.
16pub const TURBO_CODE_WIRE_MAGIC: &[u8; 4] = b"TQW1";
17
18const VERSION: u16 = 1;
19const VARIANT_TURBO_CODE: u8 = 1;
20
21/// Encoder/decoder for TurboCode wire format v1.
22pub struct TurboCodeWireV1;
23
24impl TurboCodeWireV1 {
25    /// Encode a validated TurboCode using the supplied quantizer profile.
26    pub fn encode(code: &TurboCode, profile: &TurboQuantizer) -> Result<Vec<u8>> {
27        code.validate_for(
28            profile.dim(),
29            profile.bits(),
30            profile.projections(),
31            profile.mode(),
32        )?;
33
34        let dim = checked_u32(profile.dim(), "dimension")?;
35        let polar_bits = code.polar_code.bits;
36        let qjl_projections = checked_u32(profile.projections(), "projection count")?;
37        let polar_block_count = checked_u32(code.polar_code.radii.len(), "polar block count")?;
38        let qjl_sign_count = checked_u32(
39            match profile.mode() {
40                TurboMode::PolarOnly => 0,
41                TurboMode::PolarWithQjl => code.residual_sketch.projections,
42            },
43            "qjl sign count",
44        )?;
45        let packed_angle_indices =
46            bitpack::pack_indices(&code.polar_code.angle_indices, polar_bits)?;
47        let packed_signs = match profile.mode() {
48            TurboMode::PolarOnly => Vec::new(),
49            TurboMode::PolarWithQjl => bitpack::pack_signs(&code.residual_sketch.signs)?,
50        };
51        let payload_len = checked_u64(
52            code.polar_code.radii.len() * 4 + packed_angle_indices.len() + packed_signs.len(),
53            "payload length",
54        )?;
55
56        let mut bytes = Vec::with_capacity(42 + payload_len as usize);
57        bytes.extend_from_slice(TURBO_CODE_WIRE_MAGIC);
58        bytes.extend_from_slice(&VERSION.to_le_bytes());
59        bytes.extend_from_slice(&rotation_flag(profile.rotation_kind()).to_le_bytes());
60        bytes.push(VARIANT_TURBO_CODE);
61        bytes.push(0);
62        bytes.extend_from_slice(&dim.to_le_bytes());
63        bytes.push(polar_bits);
64        bytes.extend_from_slice(&[0, 0, 0]);
65        bytes.extend_from_slice(&qjl_projections.to_le_bytes());
66        bytes.extend_from_slice(&profile.seed().to_le_bytes());
67        bytes.extend_from_slice(&polar_block_count.to_le_bytes());
68        bytes.extend_from_slice(&qjl_sign_count.to_le_bytes());
69        bytes.extend_from_slice(&payload_len.to_le_bytes());
70
71        for radius in &code.polar_code.radii {
72            bytes.extend_from_slice(&radius.to_le_bytes());
73        }
74        bytes.extend_from_slice(&packed_angle_indices);
75        bytes.extend_from_slice(&packed_signs);
76        Ok(bytes)
77    }
78
79    /// Decode and validate TurboCode wire bytes against the supplied profile.
80    pub fn decode(bytes: &[u8], profile: &TurboQuantizer) -> Result<TurboCode> {
81        let mut cursor = WireCursor::new(bytes);
82        if cursor.read_exact(TURBO_CODE_WIRE_MAGIC.len())? != TURBO_CODE_WIRE_MAGIC {
83            return Err(TurboQuantError::MalformedCode {
84                reason: "wrong TurboQuant wire magic".into(),
85            });
86        }
87        let version = cursor.read_u16()?;
88        if version != VERSION {
89            return Err(TurboQuantError::MalformedCode {
90                reason: format!("unsupported TurboQuant wire version {version}"),
91            });
92        }
93        let wire_rotation_flag = cursor.read_u16()?;
94        let expected_rotation_flag = rotation_flag(profile.rotation_kind());
95        if wire_rotation_flag != expected_rotation_flag {
96            return Err(TurboQuantError::MalformedCode {
97                reason: format!(
98                    "wire rotation flag {wire_rotation_flag} does not match quantizer profile flag {expected_rotation_flag}"
99                ),
100            });
101        }
102        let variant = cursor.read_u8()?;
103        if variant != VARIANT_TURBO_CODE {
104            return Err(TurboQuantError::MalformedCode {
105                reason: format!("unsupported TurboQuant wire variant {variant}"),
106            });
107        }
108        let reserved = cursor.read_u8()?;
109        if reserved != 0 {
110            return Err(TurboQuantError::MalformedCode {
111                reason: "nonzero TurboQuant wire reserved byte".into(),
112            });
113        }
114
115        let dim = cursor.read_u32()? as usize;
116        let polar_bits = cursor.read_u8()?;
117        let reserved2 = cursor.read_exact(3)?;
118        if reserved2 != [0, 0, 0] {
119            return Err(TurboQuantError::MalformedCode {
120                reason: "nonzero TurboQuant wire reserved bytes".into(),
121            });
122        }
123        let qjl_projections = cursor.read_u32()? as usize;
124        let seed = cursor.read_u64()?;
125        let polar_block_count = cursor.read_u32()? as usize;
126        let qjl_sign_count = cursor.read_u32()? as usize;
127        let payload_len = cursor.read_u64()?;
128        let payload_start = cursor.offset();
129
130        let expected_polar_bits = match profile.mode() {
131            TurboMode::PolarOnly => profile.bits(),
132            TurboMode::PolarWithQjl => profile.bits() - 1,
133        };
134        if dim != profile.dim()
135            || polar_bits != expected_polar_bits
136            || qjl_projections != profile.projections()
137        {
138            return Err(TurboQuantError::MalformedCode {
139                reason: "wire header does not match quantizer profile".into(),
140            });
141        }
142        if seed != profile.seed() {
143            return Err(TurboQuantError::MalformedCode {
144                reason: format!(
145                    "wire seed {seed} does not match quantizer profile seed {}",
146                    profile.seed()
147                ),
148            });
149        }
150        if polar_block_count != profile.dim() / 2 {
151            return Err(TurboQuantError::MalformedCode {
152                reason: format!(
153                    "wire polar block count {polar_block_count} does not match dimension {}",
154                    profile.dim()
155                ),
156            });
157        }
158        let expected_qjl_sign_count = match profile.mode() {
159            TurboMode::PolarOnly => 0,
160            TurboMode::PolarWithQjl => profile.projections(),
161        };
162        if qjl_sign_count != expected_qjl_sign_count {
163            return Err(TurboQuantError::MalformedCode {
164                reason: format!(
165                    "wire sign count {qjl_sign_count} does not match expected {expected_qjl_sign_count}"
166                ),
167            });
168        }
169        let angle_bytes = bitpack::packed_len(polar_block_count, polar_bits)?;
170        let sign_bytes = match profile.mode() {
171            TurboMode::PolarOnly => 0,
172            TurboMode::PolarWithQjl => profile.projections().div_ceil(8),
173        };
174        let residual_bytes = sign_bytes;
175        let expected_payload_len = checked_u64(
176            polar_block_count * 4 + angle_bytes + residual_bytes,
177            "expected payload length",
178        )?;
179        if payload_len != expected_payload_len {
180            return Err(TurboQuantError::MalformedCode {
181                reason: format!(
182                    "TurboQuant wire payload length {payload_len} does not match expected {expected_payload_len}"
183                ),
184            });
185        }
186        if payload_len > cursor.remaining_len() as u64 {
187            return Err(TurboQuantError::MalformedCode {
188                reason: "TurboQuant wire payload length exceeds remaining bytes".into(),
189            });
190        }
191
192        let mut radii = Vec::with_capacity(polar_block_count);
193        for _ in 0..polar_block_count {
194            radii.push(cursor.read_f32()?);
195        }
196        let packed_angle_indices = cursor.read_exact(angle_bytes)?.to_vec();
197        let angle_indices =
198            bitpack::unpack_indices(&packed_angle_indices, polar_block_count, polar_bits)?;
199        let residual_sketch = match profile.mode() {
200            TurboMode::PolarOnly => QjlSketch {
201                dim: profile.dim(),
202                projections: 0,
203                signs: Vec::new(),
204            },
205            TurboMode::PolarWithQjl => {
206                let packed_signs = cursor.read_exact(sign_bytes)?.to_vec();
207                let signs = bitpack::unpack_signs(&packed_signs, profile.projections())?;
208                QjlSketch {
209                    dim: profile.dim(),
210                    projections: profile.projections(),
211                    signs,
212                }
213            }
214        };
215        if cursor.offset() - payload_start != payload_len as usize {
216            return Err(TurboQuantError::MalformedCode {
217                reason: "TurboQuant wire payload length mismatch".into(),
218            });
219        }
220        cursor.finish()?;
221
222        let code = TurboCode {
223            polar_code: PolarCode {
224                dim: profile.dim(),
225                bits: polar_bits,
226                radii,
227                angle_indices,
228            },
229            residual_sketch,
230        };
231        code.validate_for(
232            profile.dim(),
233            profile.bits(),
234            profile.projections(),
235            profile.mode(),
236        )?;
237        Ok(code)
238    }
239}
240
241fn checked_u32(value: usize, field: &str) -> Result<u32> {
242    u32::try_from(value).map_err(|_| TurboQuantError::MalformedCode {
243        reason: format!("{field} {value} does not fit u32 wire field"),
244    })
245}
246
247fn checked_u64(value: usize, field: &str) -> Result<u64> {
248    u64::try_from(value).map_err(|_| TurboQuantError::MalformedCode {
249        reason: format!("{field} {value} does not fit u64 wire field"),
250    })
251}
252
253fn rotation_flag(kind: RotationKind) -> u16 {
254    match kind {
255        RotationKind::Auto => 0,
256        RotationKind::FastHadamard => 1,
257        RotationKind::StoredQr => 2,
258    }
259}
260
261struct WireCursor<'a> {
262    bytes: &'a [u8],
263    offset: usize,
264}
265
266/// Decoded TurboQuant wire header. The wire format carries the full
267/// quantizer profile (dim, bits, projections, seed, mode, rotation kind)
268/// in the first 44 bytes, so a `TurboCode` can be reconstructed from
269/// the wire bytes alone — no external quantizer required.
270#[derive(Debug, Clone, PartialEq, Eq)]
271pub struct TurboCodeWireHeader {
272    /// Original vector dimension.
273    pub dim: usize,
274    /// Polar-code bits per angle (b in the paper; b-1 for PolarWithQjl mode).
275    pub polar_bits: u8,
276    /// QJL projection count for the residual sketch.
277    pub qjl_projections: usize,
278    /// Seed used to derive the projection state.
279    pub seed: u64,
280    /// Number of polar code blocks (≈ dim / 2).
281    pub polar_block_count: usize,
282    /// QJL sign count (0 for PolarOnly mode).
283    pub qjl_sign_count: usize,
284    /// Length of the payload section following the header.
285    pub payload_len: u64,
286    /// Rotation kind embedded in the wire.
287    pub rotation_kind: RotationKind,
288}
289
290impl TurboCodeWireV1 {
291    /// Parse just the 44-byte wire header. This is the public entry point
292    /// for callers that want to extract the quantizer profile from the
293    /// wire format without validating against a specific quantizer instance.
294    pub fn parse_header(bytes: &[u8]) -> Result<TurboCodeWireHeader> {
295        if bytes.len() < 44 {
296            return Err(TurboQuantError::MalformedCode {
297                reason: format!("TurboQuant wire header is {} bytes, need 44", bytes.len()),
298            });
299        }
300        if &bytes[0..4] != TURBO_CODE_WIRE_MAGIC {
301            return Err(TurboQuantError::MalformedCode {
302                reason: "wrong TurboQuant wire magic".into(),
303            });
304        }
305        let version = u16::from_le_bytes(bytes[4..6].try_into().unwrap());
306        if version != VERSION {
307            return Err(TurboQuantError::MalformedCode {
308                reason: format!("unsupported TurboQuant wire version {version}"),
309            });
310        }
311        let wire_rotation_flag = u16::from_le_bytes(bytes[6..8].try_into().unwrap());
312        let rotation_kind = match wire_rotation_flag {
313            0 => RotationKind::Auto,
314            1 => RotationKind::FastHadamard,
315            2 => RotationKind::StoredQr,
316            _ => {
317                return Err(TurboQuantError::MalformedCode {
318                    reason: format!("unknown TurboQuant rotation flag {wire_rotation_flag}"),
319                })
320            }
321        };
322        let variant = bytes[8];
323        if variant != VARIANT_TURBO_CODE {
324            return Err(TurboQuantError::MalformedCode {
325                reason: format!("unsupported TurboQuant wire variant {variant}"),
326            });
327        }
328        let reserved = bytes[9];
329        if reserved != 0 {
330            return Err(TurboQuantError::MalformedCode {
331                reason: "nonzero TurboQuant wire reserved byte".into(),
332            });
333        }
334        let dim = u32::from_le_bytes(bytes[10..14].try_into().unwrap()) as usize;
335        let polar_bits = bytes[14];
336        let reserved2: [u8; 3] = bytes[15..18].try_into().unwrap();
337        if reserved2 != [0, 0, 0] {
338            return Err(TurboQuantError::MalformedCode {
339                reason: "nonzero TurboQuant wire reserved bytes".into(),
340            });
341        }
342        let qjl_projections = u32::from_le_bytes(bytes[18..22].try_into().unwrap()) as usize;
343        let seed = u64::from_le_bytes(bytes[22..30].try_into().unwrap());
344        let polar_block_count = u32::from_le_bytes(bytes[30..34].try_into().unwrap()) as usize;
345        let qjl_sign_count = u32::from_le_bytes(bytes[34..38].try_into().unwrap()) as usize;
346        let payload_len = u64::from_le_bytes(bytes[38..46].try_into().unwrap());
347        Ok(TurboCodeWireHeader {
348            dim,
349            polar_bits,
350            qjl_projections,
351            seed,
352            polar_block_count,
353            qjl_sign_count,
354            payload_len,
355            rotation_kind,
356        })
357    }
358}
359
360impl<'a> WireCursor<'a> {
361    fn new(bytes: &'a [u8]) -> Self {
362        Self { bytes, offset: 0 }
363    }
364
365    fn offset(&self) -> usize {
366        self.offset
367    }
368
369    fn remaining_len(&self) -> usize {
370        self.bytes.len().saturating_sub(self.offset)
371    }
372
373    fn read_exact(&mut self, len: usize) -> Result<&'a [u8]> {
374        let end = self
375            .offset
376            .checked_add(len)
377            .ok_or_else(|| TurboQuantError::MalformedCode {
378                reason: "wire offset overflow".into(),
379            })?;
380        if end > self.bytes.len() {
381            return Err(TurboQuantError::MalformedCode {
382                reason: "truncated TurboQuant wire artifact".into(),
383            });
384        }
385        let out = &self.bytes[self.offset..end];
386        self.offset = end;
387        Ok(out)
388    }
389
390    fn read_u8(&mut self) -> Result<u8> {
391        Ok(self.read_exact(1)?[0])
392    }
393
394    fn read_u16(&mut self) -> Result<u16> {
395        let bytes = self.read_exact(2)?;
396        Ok(u16::from_le_bytes([bytes[0], bytes[1]]))
397    }
398
399    fn read_u32(&mut self) -> Result<u32> {
400        let bytes = self.read_exact(4)?;
401        Ok(u32::from_le_bytes([bytes[0], bytes[1], bytes[2], bytes[3]]))
402    }
403
404    fn read_u64(&mut self) -> Result<u64> {
405        let bytes = self.read_exact(8)?;
406        Ok(u64::from_le_bytes([
407            bytes[0], bytes[1], bytes[2], bytes[3], bytes[4], bytes[5], bytes[6], bytes[7],
408        ]))
409    }
410
411    fn read_f32(&mut self) -> Result<f32> {
412        let bytes = self.read_exact(4)?;
413        Ok(f32::from_le_bytes([bytes[0], bytes[1], bytes[2], bytes[3]]))
414    }
415
416    fn finish(&self) -> Result<()> {
417        if self.offset != self.bytes.len() {
418            return Err(TurboQuantError::MalformedCode {
419                reason: "trailing bytes in TurboQuant wire artifact".into(),
420            });
421        }
422        Ok(())
423    }
424}
425
426#[cfg(test)]
427mod tests {
428    use super::*;
429
430    fn make_quantizer(dim: usize, seed: u64) -> TurboQuantizer {
431        // Use the simplest possible profile: PolarWithQjl, 8-bit, 32 projections.
432        TurboQuantizer::new(dim, 8, 32, seed).expect("quantizer build")
433    }
434
435    #[test]
436    fn parse_header_round_trips_encoded_code() {
437        let q = make_quantizer(128, 42);
438        let vector: Vec<f32> = (0..128).map(|i| (i as f32 / 128.0) - 0.5).collect();
439        let code = q.encode(&vector).expect("encode");
440        let wire = TurboCodeWireV1::encode(&code, &q).expect("wire encode");
441
442        let header = TurboCodeWireV1::parse_header(&wire).expect("parse header");
443        assert_eq!(header.dim, 128);
444        assert_eq!(header.qjl_projections, 32);
445        assert_eq!(header.seed, 42);
446        assert!(header.polar_block_count > 0);
447        assert!(header.payload_len > 0);
448    }
449
450    #[test]
451    fn parse_header_rejects_short_buffer() {
452        let bytes = vec![0u8; 10];
453        let result = TurboCodeWireV1::parse_header(&bytes);
454        assert!(result.is_err());
455    }
456
457    #[test]
458    fn parse_header_rejects_bad_magic() {
459        let mut bytes = vec![0u8; 44];
460        bytes[0..4].copy_from_slice(b"XXXX");
461        let result = TurboCodeWireV1::parse_header(&bytes);
462        assert!(result.is_err());
463    }
464
465    #[test]
466    fn parse_header_rejects_unsupported_version() {
467        let mut bytes = vec![0u8; 44];
468        bytes[0..4].copy_from_slice(TURBO_CODE_WIRE_MAGIC);
469        // version = 99 (unsupported)
470        bytes[4..6].copy_from_slice(&99u16.to_le_bytes());
471        let result = TurboCodeWireV1::parse_header(&bytes);
472        assert!(result.is_err());
473    }
474}
turbo_quant/wire.rs

turbo_quant/
wire.rs
