apr_format/v2/
reader_impl.rs

1//! v2 reader impls + shard manifest (issue #2231).
2//!
3//! Formerly `include!`d into `v2/mod.rs`; now a real module. The `AprV2Reader`
4//! and `AprV2ReaderRef` struct declarations live here next to their `impl`s
5//! (private fields).
6//!
7//! # Sovereignty seam (issue #2231)
8//!
9//! The dequantizing `get_tensor_as_f32` accessor (which needed the GGUF Q4_K /
10//! Q6_K dequant kernels + the local f16-scaled `dequantize_q4`) is **severed**
11//! from the leaf: it pulls quantization/physics that belongs to the framework.
12//! The leaf exposes only the raw container bytes ([`AprV2Reader::get_tensor_data`])
13//! and the trivial F32-only typed view ([`AprV2Reader::get_f32_tensor`]).
14//! `aprender-core` re-attaches the dequantizing accessor as the `AprV2DequantExt`
15//! extension trait.
16
17use super::{
18    is_aligned_64, AprV2Header, AprV2Metadata, TensorDType, TensorIndexEntry, V2FormatError,
19    HEADER_SIZE_V2,
20};
21use serde::{Deserialize, Serialize};
22use std::collections::HashMap;
23use std::io::Read;
24
25/// APR v2 format reader (owns data - copies input)
26#[derive(Debug)]
27pub struct AprV2Reader {
28    header: AprV2Header,
29    metadata: AprV2Metadata,
30    tensor_index: Vec<TensorIndexEntry>,
31    data: Vec<u8>,
32}
33
34/// APR v2 format reader with zero-copy (borrows data - for mmap)
35///
36/// This reader borrows the data slice instead of copying it, enabling
37/// true zero-copy access when used with memory-mapped files.
38///
39/// # Example
40///
41/// ```ignore
42/// use apr_format::v2::AprV2ReaderRef;
43///
44/// let bytes: &[u8] = /* mmap'd slice */;
45/// let reader = AprV2ReaderRef::from_bytes(bytes)?;
46/// let weights = reader.get_f32_tensor("embed_tokens.weight")?;
47/// ```
48#[derive(Debug)]
49pub struct AprV2ReaderRef<'a> {
50    header: AprV2Header,
51    metadata: AprV2Metadata,
52    tensor_index: Vec<TensorIndexEntry>,
53    data: &'a [u8],
54}
55
56/// Parse and bounds-check the JSON metadata section (FALSIFY-PARSE-001).
57///
58/// All offsets/sizes here come straight from the (attacker-controllable) file
59/// header — the CRC32 checksum is computed over the header itself, so a
60/// corrupted file can carry a matching checksum. Therefore every offset+size is
61/// validated with checked arithmetic and `slice::get` (never `[start..end]`
62/// indexing, which panics on out-of-range / `start > end`).
63fn parse_metadata_section(
64    data: &[u8],
65    metadata_offset: u64,
66    metadata_size: u32,
67) -> Result<AprV2Metadata, V2FormatError> {
68    let start = usize::try_from(metadata_offset)
69        .map_err(|_| V2FormatError::InvalidHeader("metadata_offset exceeds usize".to_string()))?;
70    let end = start
71        .checked_add(metadata_size as usize)
72        .ok_or_else(|| V2FormatError::InvalidHeader("metadata offset+size overflow".to_string()))?;
73    let slice = data
74        .get(start..end)
75        .ok_or_else(|| V2FormatError::InvalidHeader("file too small for metadata".to_string()))?;
76    AprV2Metadata::from_json(slice)
77}
78
79/// Parse and bounds-check the tensor index section (FALSIFY-PARSE-001).
80///
81/// `tensor_index_offset` is attacker-controllable; previously it was used
82/// directly as `&data[pos..]`, which PANICS ("range start index out of bounds")
83/// when the offset points past EOF. Now the start offset is validated against
84/// the file length before slicing, and each entry advances `pos` with the same
85/// `slice::get` guard.
86fn parse_tensor_index_section(
87    data: &[u8],
88    tensor_index_offset: u64,
89    tensor_count: u32,
90) -> Result<Vec<TensorIndexEntry>, V2FormatError> {
91    let mut pos = usize::try_from(tensor_index_offset).map_err(|_| {
92        V2FormatError::InvalidTensorIndex("tensor_index_offset exceeds usize".to_string())
93    })?;
94
95    let mut tensor_index = Vec::with_capacity(tensor_count as usize);
96    for _ in 0..tensor_count {
97        // `data.get(pos..)` returns None only when pos > data.len(); pos == len
98        // yields an empty slice, which TensorIndexEntry::from_bytes rejects
99        // cleanly. This replaces the panicking `&data[pos..]`.
100        let remaining = data.get(pos..).ok_or_else(|| {
101            V2FormatError::InvalidTensorIndex("tensor index offset past end of file".to_string())
102        })?;
103        let (entry, consumed) = TensorIndexEntry::from_bytes(remaining)?;
104        tensor_index.push(entry);
105        pos = pos.checked_add(consumed).ok_or_else(|| {
106            V2FormatError::InvalidTensorIndex("tensor index position overflow".to_string())
107        })?;
108    }
109
110    // Verify tensor names are sorted
111    for i in 1..tensor_index.len() {
112        if tensor_index[i].name < tensor_index[i - 1].name {
113            return Err(V2FormatError::InvalidTensorIndex(
114                "tensor index not sorted".to_string(),
115            ));
116        }
117    }
118
119    Ok(tensor_index)
120}
121
122impl AprV2Reader {
123    /// Read from bytes
124    ///
125    /// # Errors
126    /// Returns error if parsing fails.
127    ///
128    /// # LAYOUT-002 Jidoka Guard
129    /// Rejects APR files with `LAYOUT_COLUMN_MAJOR` flag set, as these indicate
130    /// improperly converted GGUF files that would produce garbage output.
131    pub fn from_bytes(data: &[u8]) -> Result<Self, V2FormatError> {
132        if data.len() < HEADER_SIZE_V2 {
133            return Err(V2FormatError::InvalidHeader("file too small".to_string()));
134        }
135
136        // Parse header
137        let header = AprV2Header::from_bytes(data)?;
138
139        // Verify checksum
140        if !header.verify_checksum() {
141            return Err(V2FormatError::ChecksumMismatch);
142        }
143
144        // LAYOUT-002: Jidoka Guard - Reject "dirty" APR files with column-major layout
145        if !header.flags.is_layout_valid() {
146            return Err(V2FormatError::InvalidHeader(
147                "LAYOUT-002 violation: APR file has LAYOUT_COLUMN_MAJOR flag set. \
148                 This indicates a dirty import from GGUF without proper transpose. \
149                 Re-import the model using `apr import` with LAYOUT-002 enforcement."
150                    .to_string(),
151            ));
152        }
153
154        // Parse metadata (FALSIFY-PARSE-001 / PMAT-822: checked arithmetic +
155        // .get() so a corrupted metadata_offset/size can never panic-slice).
156        let metadata = parse_metadata_section(data, header.metadata_offset, header.metadata_size)?;
157
158        // Parse tensor index
159        let tensor_index =
160            parse_tensor_index_section(data, header.tensor_index_offset, header.tensor_count)?;
161
162        Ok(Self {
163            header,
164            metadata,
165            tensor_index,
166            data: data.to_vec(),
167        })
168    }
169
170    /// Read from a Read impl
171    ///
172    /// # Errors
173    /// Returns error if read fails.
174    pub fn from_reader<R: Read>(reader: &mut R) -> Result<Self, V2FormatError> {
175        let mut data = Vec::new();
176        reader
177            .read_to_end(&mut data)
178            .map_err(|e| V2FormatError::IoError(e.to_string()))?;
179        Self::from_bytes(&data)
180    }
181
182    /// Get header
183    #[must_use]
184    pub fn header(&self) -> &AprV2Header {
185        &self.header
186    }
187
188    /// Get metadata
189    #[must_use]
190    pub fn metadata(&self) -> &AprV2Metadata {
191        &self.metadata
192    }
193
194    /// Get tensor names
195    #[must_use]
196    pub fn tensor_names(&self) -> Vec<&str> {
197        self.tensor_index.iter().map(|e| e.name.as_str()).collect()
198    }
199
200    /// Get tensor by name
201    #[must_use]
202    pub fn get_tensor(&self, name: &str) -> Option<&TensorIndexEntry> {
203        self.tensor_index.iter().find(|e| e.name == name)
204    }
205
206    /// Get tensor data by name
207    #[must_use]
208    pub fn get_tensor_data(&self, name: &str) -> Option<&[u8]> {
209        let entry = self.get_tensor(name)?;
210        // FALSIFY-PARSE-001 / PMAT-822: data_offset + offset (u64) and
211        // start + size (usize) can both wrap for a crafted header, letting a
212        // wrapped `end <= len` check pass over an OOB region. Use checked
213        // arithmetic + `slice::get` so any overflow / past-EOF range → None.
214        let abs_offset = self.header.data_offset.checked_add(entry.offset)?;
215        let start = usize::try_from(abs_offset).ok()?;
216        let end = start.checked_add(usize::try_from(entry.size).ok()?)?;
217        self.data.get(start..end)
218    }
219
220    /// Get tensor as f32 slice (F32 dtype only)
221    #[must_use]
222    pub fn get_f32_tensor(&self, name: &str) -> Option<Vec<f32>> {
223        let entry = self.get_tensor(name)?;
224        if entry.dtype != TensorDType::F32 {
225            return None;
226        }
227
228        let data = self.get_tensor_data(name)?;
229        let floats: Vec<f32> = data
230            .chunks_exact(4)
231            .map(|chunk| f32::from_le_bytes([chunk[0], chunk[1], chunk[2], chunk[3]]))
232            .collect();
233
234        Some(floats)
235    }
236
237    // NOTE (issue #2231): `get_tensor_as_f32` (the dequantizing accessor) is
238    // SEVERED from the sovereign leaf — it needed the GGUF Q4_K/Q6_K dequant
239    // kernels + the f16-scaled `dequantize_q4`, all framework/quant concerns.
240    // `aprender-core` re-attaches it via the `AprV2DequantExt` extension trait.
241
242    /// Check if all tensors are 64-byte aligned
243    #[must_use]
244    pub fn verify_alignment(&self) -> bool {
245        let data_offset = self.header.data_offset as usize;
246        self.tensor_index
247            .iter()
248            .all(|e| is_aligned_64(data_offset + e.offset as usize))
249    }
250
251    /// Borrow the parsed tensor index (used by the core dequant extension).
252    #[must_use]
253    pub fn tensor_index(&self) -> &[TensorIndexEntry] {
254        &self.tensor_index
255    }
256}
257
258impl<'a> AprV2ReaderRef<'a> {
259    /// Read from bytes (zero-copy - borrows data)
260    ///
261    /// Unlike `AprV2Reader::from_bytes`, this does NOT copy the input data.
262    /// The reader borrows the slice, making it ideal for use with mmap.
263    ///
264    /// # Errors
265    /// Returns error if parsing fails.
266    ///
267    /// # LAYOUT-002 Jidoka Guard
268    /// Rejects APR files with `LAYOUT_COLUMN_MAJOR` flag set, as these indicate
269    /// improperly converted GGUF files that would produce garbage output.
270    pub fn from_bytes(data: &'a [u8]) -> Result<Self, V2FormatError> {
271        if data.len() < HEADER_SIZE_V2 {
272            return Err(V2FormatError::InvalidHeader("file too small".to_string()));
273        }
274
275        // Parse header
276        let header = AprV2Header::from_bytes(data)?;
277
278        // Verify checksum
279        if !header.verify_checksum() {
280            return Err(V2FormatError::ChecksumMismatch);
281        }
282
283        // LAYOUT-002: Jidoka Guard - Reject "dirty" APR files with column-major layout
284        if !header.flags.is_layout_valid() {
285            return Err(V2FormatError::InvalidHeader(
286                "LAYOUT-002 violation: APR file has LAYOUT_COLUMN_MAJOR flag set. \
287                 This indicates a dirty import from GGUF without proper transpose. \
288                 Re-import the model using `apr import` with LAYOUT-002 enforcement."
289                    .to_string(),
290            ));
291        }
292
293        // Parse metadata (FALSIFY-PARSE-001 / PMAT-822: checked arithmetic +
294        // .get() so a corrupted metadata_offset/size can never panic-slice).
295        let metadata = parse_metadata_section(data, header.metadata_offset, header.metadata_size)?;
296
297        // Parse tensor index
298        let tensor_index =
299            parse_tensor_index_section(data, header.tensor_index_offset, header.tensor_count)?;
300
301        Ok(Self {
302            header,
303            metadata,
304            tensor_index,
305            data, // Borrow, no copy!
306        })
307    }
308
309    /// Get header
310    #[must_use]
311    pub fn header(&self) -> &AprV2Header {
312        &self.header
313    }
314
315    /// Get metadata
316    #[must_use]
317    pub fn metadata(&self) -> &AprV2Metadata {
318        &self.metadata
319    }
320
321    /// Get tensor names
322    #[must_use]
323    pub fn tensor_names(&self) -> Vec<&str> {
324        self.tensor_index.iter().map(|e| e.name.as_str()).collect()
325    }
326
327    /// Get tensor by name
328    #[must_use]
329    pub fn get_tensor(&self, name: &str) -> Option<&TensorIndexEntry> {
330        self.tensor_index.iter().find(|e| e.name == name)
331    }
332
333    /// Get tensor data by name (zero-copy slice into mmap)
334    #[must_use]
335    pub fn get_tensor_data(&self, name: &str) -> Option<&[u8]> {
336        let entry = self.get_tensor(name)?;
337        // FALSIFY-PARSE-001 / PMAT-822: data_offset + offset (u64) and
338        // start + size (usize) can both wrap for a crafted header, letting a
339        // wrapped `end <= len` check pass over an OOB region. Use checked
340        // arithmetic + `slice::get` so any overflow / past-EOF range → None.
341        let abs_offset = self.header.data_offset.checked_add(entry.offset)?;
342        let start = usize::try_from(abs_offset).ok()?;
343        let end = start.checked_add(usize::try_from(entry.size).ok()?)?;
344        self.data.get(start..end)
345    }
346
347    /// Get tensor as f32 Vec (copies data from mmap to `Vec<f32>`)
348    ///
349    /// Note: This allocates memory for the f32 values. For very large tensors,
350    /// consider using `get_tensor_data` and processing in chunks.
351    #[must_use]
352    pub fn get_f32_tensor(&self, name: &str) -> Option<Vec<f32>> {
353        let entry = self.get_tensor(name)?;
354        if entry.dtype != TensorDType::F32 {
355            return None;
356        }
357
358        let data = self.get_tensor_data(name)?;
359        let floats: Vec<f32> = data
360            .chunks_exact(4)
361            .map(|chunk| f32::from_le_bytes([chunk[0], chunk[1], chunk[2], chunk[3]]))
362            .collect();
363
364        Some(floats)
365    }
366
367    // NOTE (issue #2231): `get_tensor_as_f32` severed from the leaf — see the
368    // owning-reader note above. Re-attached in core via `AprV2DequantExt`.
369
370    /// Check if all tensors are 64-byte aligned
371    #[must_use]
372    pub fn verify_alignment(&self) -> bool {
373        let data_offset = self.header.data_offset as usize;
374        self.tensor_index
375            .iter()
376            .all(|e| is_aligned_64(data_offset + e.offset as usize))
377    }
378
379    /// Borrow the parsed tensor index (used by the core dequant extension).
380    #[must_use]
381    pub fn tensor_index(&self) -> &[TensorIndexEntry] {
382        &self.tensor_index
383    }
384}
385
386// ============================================================================
387// Shard Manifest
388// ============================================================================
389
390/// Shard manifest for multi-file models
391#[derive(Debug, Clone, Serialize, Deserialize)]
392pub struct ShardManifest {
393    /// Format version
394    pub version: String,
395    /// Total number of shards
396    pub shard_count: usize,
397    /// Total size in bytes
398    pub total_size: u64,
399    /// Total tensor count
400    pub tensor_count: usize,
401    /// Shard files
402    pub shards: Vec<ShardInfo>,
403    /// Tensor to shard mapping
404    pub weight_map: HashMap<String, usize>,
405}
406
407/// Information about a single shard
408#[derive(Debug, Clone, Serialize, Deserialize)]
409pub struct ShardInfo {
410    /// Shard filename
411    pub filename: String,
412    /// Shard index
413    pub index: usize,
414    /// Size in bytes
415    pub size: u64,
416    /// Tensor names in this shard
417    pub tensors: Vec<String>,
418}
apr_format/v2/reader_impl.rs

apr_format/v2/
reader_impl.rs
