fib_quant/
wire.rs

1//! Self-describing wire format for FibQuant codes.
2//!
3//! The FB1/FB2 compact formats strip all profile metadata, forcing
4//! downstream consumers to hardcode seed, dim, k, N externally. This
5//! module provides `FibCodeWireV1` — a self-describing format with a
6//! header carrying seed, dim, k, N, norm_format, and profile digest,
7//! enabling decode without any external profile knowledge.
8
9use crate::{
10    bitpack::unpack_indices,
11    codec::{FibCodeV1, CODE_SCHEMA},
12    profile::{FibQuantProfileV1, NormFormat},
13    FibQuantError, Result,
14};
15
16/// Magic bytes: "FBW1" = Fib Binary Wire v1.
17pub const WIRE_MAGIC: [u8; 4] = [b'F', b'B', b'W', b'1'];
18/// Current wire format version.
19pub const WIRE_VERSION: u8 = 1;
20/// Header size: magic(4) + version(1) + ambient_dim(4) + block_dim(4)
21/// + codebook_size(4) + rotation_seed(8) + wire_index_bits(1) + norm_format(1)
22/// + profile_digest(32) = 59 bytes.
23pub const WIRE_HEADER_SIZE: usize = 59;
24
25/// Parsed wire header (metadata without full decode).
26#[derive(Debug, Clone, PartialEq)]
27pub struct WireHeader {
28    pub version: u8,
29    pub ambient_dim: u32,
30    pub block_dim: u32,
31    pub codebook_size: u32,
32    pub rotation_seed: u64,
33    pub wire_index_bits: u8,
34    pub norm_format: NormFormat,
35    pub profile_digest: [u8; 32],
36    pub body_offset: usize,
37}
38
39/// Self-describing wire format codec.
40pub struct FibCodeWireV1;
41
42impl FibCodeWireV1 {
43    /// Encode a FibCodeV1 with profile metadata into a self-describing
44    /// wire format. The header carries all profile fields needed to
45    /// reconstruct the quantizer for decode.
46    pub fn to_wire_bytes(code: &FibCodeV1, profile: &FibQuantProfileV1) -> Result<Vec<u8>> {
47        let profile_digest_hex = profile.digest()?;
48        let digest_bytes = hex_decode_32(strip_blake3_prefix(&profile_digest_hex))?;
49        let mut out =
50            Vec::with_capacity(WIRE_HEADER_SIZE + 2 + code.norm_payload.len() + code.indices.len());
51
52        // Header
53        out.extend_from_slice(&WIRE_MAGIC);
54        out.push(WIRE_VERSION);
55        out.extend_from_slice(&profile.ambient_dim.to_le_bytes());
56        out.extend_from_slice(&profile.block_dim.to_le_bytes());
57        out.extend_from_slice(&profile.codebook_size.to_le_bytes());
58        out.extend_from_slice(&profile.rotation_seed.to_le_bytes());
59        out.push(profile.wire_index_bits);
60        out.push(norm_format_to_byte(&profile.norm_format));
61        out.extend_from_slice(&digest_bytes);
62
63        // Body: norm_len(u16) + norm_payload + indices
64        let norm_len = code.norm_payload.len() as u16;
65        out.extend_from_slice(&norm_len.to_le_bytes());
66        out.extend_from_slice(&code.norm_payload);
67        out.extend_from_slice(&code.indices);
68        Ok(out)
69    }
70
71    /// Parse only the header from wire bytes. Does not decode the body.
72    /// Useful for metadata extraction without building a quantizer.
73    pub fn parse_header(bytes: &[u8]) -> Result<WireHeader> {
74        if bytes.len() < WIRE_HEADER_SIZE {
75            return Err(FibQuantError::CorruptPayload(format!(
76                "wire format too short: {} bytes (need >= {})",
77                bytes.len(),
78                WIRE_HEADER_SIZE
79            )));
80        }
81        if bytes[0..4] != WIRE_MAGIC {
82            return Err(FibQuantError::CorruptPayload(format!(
83                "wire bad magic: {:?} (expected {:?})",
84                &bytes[0..4],
85                WIRE_MAGIC
86            )));
87        }
88        let version = bytes[4];
89        if version != WIRE_VERSION {
90            return Err(FibQuantError::CorruptPayload(format!(
91                "wire version {} not supported (need {})",
92                version, WIRE_VERSION
93            )));
94        }
95        let ambient_dim = u32::from_le_bytes([bytes[5], bytes[6], bytes[7], bytes[8]]);
96        let block_dim = u32::from_le_bytes([bytes[9], bytes[10], bytes[11], bytes[12]]);
97        let codebook_size = u32::from_le_bytes([bytes[13], bytes[14], bytes[15], bytes[16]]);
98        let rotation_seed = u64::from_le_bytes([
99            bytes[17], bytes[18], bytes[19], bytes[20], bytes[21], bytes[22], bytes[23], bytes[24],
100        ]);
101        let wire_index_bits = bytes[25];
102        let norm_format = byte_to_norm_format(bytes[26])?;
103        let mut profile_digest = [0u8; 32];
104        profile_digest.copy_from_slice(&bytes[27..59]);
105
106        Ok(WireHeader {
107            version,
108            ambient_dim,
109            block_dim,
110            codebook_size,
111            rotation_seed,
112            wire_index_bits,
113            norm_format,
114            profile_digest,
115            body_offset: WIRE_HEADER_SIZE,
116        })
117    }
118
119    /// Decode a self-describing wire payload into a FibCodeV1 and the
120    /// reconstructed FibQuantProfileV1. The caller can use the profile
121    /// to build a FibQuantizer for decode.
122    pub fn from_wire_bytes(bytes: &[u8]) -> Result<(FibCodeV1, FibQuantProfileV1)> {
123        let header = Self::parse_header(bytes)?;
124        let body = &bytes[header.body_offset..];
125        if body.len() < 2 {
126            return Err(FibQuantError::CorruptPayload(
127                "wire body too short for norm_len".into(),
128            ));
129        }
130        let norm_len = u16::from_le_bytes([body[0], body[1]]) as usize;
131        if body.len() < 2 + norm_len {
132            return Err(FibQuantError::CorruptPayload(format!(
133                "wire norm truncated: norm_len={} but only {} bytes remain",
134                norm_len,
135                body.len() - 2
136            )));
137        }
138        let norm_payload = body[2..2 + norm_len].to_vec();
139        let indices = body[2 + norm_len..].to_vec();
140
141        // Reconstruct profile from header fields
142        let profile = FibQuantProfileV1::paper_default(
143            header.ambient_dim as usize,
144            header.block_dim as usize,
145            header.codebook_size as usize,
146            header.rotation_seed,
147        )?;
148
149        // Validate wire_index_bits
150        if profile.wire_index_bits != header.wire_index_bits {
151            return Err(FibQuantError::CorruptPayload(format!(
152                "wire_index_bits {} does not match expected {}",
153                header.wire_index_bits, profile.wire_index_bits
154            )));
155        }
156
157        // Validate packed index length
158        let block_count = profile.block_count() as usize;
159        let expected_packed_len = (block_count * header.wire_index_bits as usize).div_ceil(8);
160        if indices.len() != expected_packed_len {
161            return Err(FibQuantError::CorruptPayload(format!(
162                "wire indices: got {} bytes, expected {}",
163                indices.len(),
164                expected_packed_len
165            )));
166        }
167
168        // Verify indices unpack correctly
169        let _ = unpack_indices(&indices, block_count, header.wire_index_bits)?;
170
171        let profile_digest = profile.digest()?;
172        Ok((
173            FibCodeV1 {
174                schema_version: CODE_SCHEMA.into(),
175                profile_digest,
176                codebook_digest: String::new(),
177                rotation_digest: String::new(),
178                ambient_dim: profile.ambient_dim,
179                block_dim: profile.block_dim,
180                norm_format: profile.norm_format.clone(),
181                norm_payload,
182                wire_index_bits: header.wire_index_bits,
183                block_count: profile.block_count(),
184                indices,
185            },
186            profile,
187        ))
188    }
189}
190
191fn norm_format_to_byte(format: &NormFormat) -> u8 {
192    match format {
193        NormFormat::Fp16Paper => 0,
194        NormFormat::F32Reference => 1,
195    }
196}
197
198fn byte_to_norm_format(byte: u8) -> Result<NormFormat> {
199    match byte {
200        0 => Ok(NormFormat::Fp16Paper),
201        1 => Ok(NormFormat::F32Reference),
202        _ => Err(FibQuantError::CorruptPayload(format!(
203            "unknown norm_format byte: {}",
204            byte
205        ))),
206    }
207}
208
209fn hex_decode_32(hex: &str) -> Result<[u8; 32]> {
210    if hex.len() != 64 {
211        return Err(FibQuantError::CorruptPayload(format!(
212            "digest hex length {} (expected 64)",
213            hex.len()
214        )));
215    }
216    let mut out = [0u8; 32];
217    for i in 0..32 {
218        out[i] = u8::from_str_radix(&hex[i * 2..i * 2 + 2], 16)
219            .map_err(|_| FibQuantError::CorruptPayload("invalid hex digest".into()))?;
220    }
221    Ok(out)
222}
223
224/// Strip the "blake3:" prefix from fib-quant digests.
225fn strip_blake3_prefix(digest: &str) -> &str {
226    digest.strip_prefix("blake3:").unwrap_or(digest)
227}
228
229#[allow(dead_code)]
230fn hex_encode_32(bytes: &[u8; 32]) -> String {
231    let mut out = String::with_capacity(64);
232    for b in bytes.iter() {
233        out.push_str(&format!("{:02x}", b));
234    }
235    out
236}
237
238#[cfg(test)]
239mod tests {
240    use super::*;
241    use crate::{codec::FibQuantizer, profile::FibQuantProfileV1};
242
243    fn build_test() -> Result<(FibQuantizer, FibQuantProfileV1)> {
244        let mut profile = FibQuantProfileV1::paper_default(8, 2, 8, 7)?;
245        profile.training_samples = 128;
246        profile.lloyd_restarts = 1;
247        profile.lloyd_iterations = 2;
248        let q = FibQuantizer::new(profile.clone())?;
249        Ok((q, profile))
250    }
251
252    #[test]
253    fn wire_roundtrip_preserves_data() -> Result<()> {
254        let (q, profile) = build_test()?;
255        let input = vec![0.25, -0.5, 0.75, 1.0, -1.25, 0.5, 0.125, -0.875];
256        let code = q.encode(&input)?;
257        let wire = FibCodeWireV1::to_wire_bytes(&code, &profile)?;
258        let (decoded_code, decoded_profile) = FibCodeWireV1::from_wire_bytes(&wire)?;
259        assert_eq!(decoded_code.indices, code.indices);
260        assert_eq!(decoded_code.norm_payload, code.norm_payload);
261        assert_eq!(decoded_code.wire_index_bits, code.wire_index_bits);
262        assert_eq!(decoded_code.block_count, code.block_count);
263        assert_eq!(decoded_profile.ambient_dim, profile.ambient_dim);
264        assert_eq!(decoded_profile.block_dim, profile.block_dim);
265        assert_eq!(decoded_profile.codebook_size, profile.codebook_size);
266        assert_eq!(decoded_profile.rotation_seed, profile.rotation_seed);
267        Ok(())
268    }
269
270    #[test]
271    fn wire_parse_header_extracts_metadata() -> Result<()> {
272        let (q, profile) = build_test()?;
273        let input = vec![0.25, -0.5, 0.75, 1.0, -1.25, 0.5, 0.125, -0.875];
274        let code = q.encode(&input)?;
275        let wire = FibCodeWireV1::to_wire_bytes(&code, &profile)?;
276        let header = FibCodeWireV1::parse_header(&wire)?;
277        assert_eq!(header.version, WIRE_VERSION);
278        assert_eq!(header.ambient_dim, 8);
279        assert_eq!(header.block_dim, 2);
280        assert_eq!(header.codebook_size, 8);
281        assert_eq!(header.rotation_seed, 7);
282        assert_eq!(header.wire_index_bits, profile.wire_index_bits);
283        assert_eq!(header.body_offset, WIRE_HEADER_SIZE);
284        Ok(())
285    }
286
287    #[test]
288    fn wire_rejects_bad_magic() -> Result<()> {
289        let result = FibCodeWireV1::parse_header(b"XXXX");
290        assert!(result.is_err());
291        Ok(())
292    }
293
294    #[test]
295    fn wire_rejects_truncation() -> Result<()> {
296        let result = FibCodeWireV1::parse_header(&[0; 10]);
297        assert!(result.is_err());
298        Ok(())
299    }
300
301    #[test]
302    fn wire_decoded_code_roundtrips_through_quantizer() -> Result<()> {
303        let (q, profile) = build_test()?;
304        let input = vec![0.25, -0.5, 0.75, 1.0, -1.25, 0.5, 0.125, -0.875];
305        let code = q.encode(&input)?;
306        let wire = FibCodeWireV1::to_wire_bytes(&code, &profile)?;
307        let (decoded_code, decoded_profile) = FibCodeWireV1::from_wire_bytes(&wire)?;
308        let q2 = FibQuantizer::new(decoded_profile)?;
309        let decoded = q2.decode(&decoded_code)?;
310        assert_eq!(decoded.len(), input.len());
311        for (a, b) in input.iter().zip(decoded.iter()) {
312            assert!(a.is_finite());
313            assert!(b.is_finite());
314        }
315        Ok(())
316    }
317
318    #[test]
319    fn wire_corrupt_dim_is_error() -> Result<()> {
320        let (q, profile) = build_test()?;
321        let input = vec![0.25, -0.5, 0.75, 1.0, -1.25, 0.5, 0.125, -0.875];
322        let code = q.encode(&input)?;
323        let mut wire = FibCodeWireV1::to_wire_bytes(&code, &profile)?;
324        // Corrupt ambient_dim (bytes 5..9)
325        wire[5] ^= 0xFF;
326        let result = FibCodeWireV1::from_wire_bytes(&wire);
327        // Should still work (profile is reconstructed from corrupted dim,
328        // but the decode will fail when the caller tries to use it)
329        // or fail on validation. Either way, no crash.
330        let _ = result;
331        Ok(())
332    }
333}
fib_quant/wire.rs

fib_quant/
wire.rs
