apr_format/v2/
v2format_error.rs

1//! `ShardManifest` impl + the `V2FormatError` type (issue #2231).
2//!
3//! Formerly `include!`d into `v2/mod.rs`; now a real module. `ShardManifest` /
4//! `ShardInfo` are declared in `reader_impl.rs` (all-pub fields), impl here.
5
6use super::{ShardInfo, ShardManifest};
7use std::collections::HashMap;
8
9impl ShardManifest {
10    /// Create new empty manifest
11    #[must_use]
12    pub fn new(shard_count: usize) -> Self {
13        Self {
14            version: "2.0".to_string(),
15            shard_count,
16            total_size: 0,
17            tensor_count: 0,
18            shards: Vec::with_capacity(shard_count),
19            weight_map: HashMap::new(),
20        }
21    }
22
23    /// Add shard info
24    pub fn add_shard(&mut self, info: ShardInfo) {
25        for tensor in &info.tensors {
26            self.weight_map.insert(tensor.clone(), info.index);
27        }
28        self.tensor_count += info.tensors.len();
29        self.total_size += info.size;
30        self.shards.push(info);
31    }
32
33    /// Get shard index for tensor
34    #[must_use]
35    pub fn shard_for_tensor(&self, name: &str) -> Option<usize> {
36        self.weight_map.get(name).copied()
37    }
38
39    /// Serialize to JSON
40    ///
41    /// # Errors
42    /// Returns error if serialization fails.
43    pub fn to_json(&self) -> Result<String, V2FormatError> {
44        serde_json::to_string_pretty(self).map_err(|e| V2FormatError::MetadataError(e.to_string()))
45    }
46
47    /// Deserialize from JSON
48    ///
49    /// # Errors
50    /// Returns error if deserialization fails.
51    pub fn from_json(json: &str) -> Result<Self, V2FormatError> {
52        serde_json::from_str(json).map_err(|e| V2FormatError::MetadataError(e.to_string()))
53    }
54}
55
56// ============================================================================
57// Error Type
58// ============================================================================
59
60/// APR v2 format error
61#[derive(Debug, Clone, PartialEq)]
62pub enum V2FormatError {
63    /// Invalid magic number
64    InvalidMagic([u8; 4]),
65    /// Invalid header
66    InvalidHeader(String),
67    /// Invalid tensor index
68    InvalidTensorIndex(String),
69    /// Metadata error
70    MetadataError(String),
71    /// Checksum mismatch
72    ChecksumMismatch,
73    /// Alignment error
74    AlignmentError(String),
75    /// I/O error
76    IoError(String),
77    /// Compression error
78    CompressionError(String),
79}
80
81impl std::fmt::Display for V2FormatError {
82    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
83        match self {
84            Self::InvalidMagic(magic) => {
85                write!(
86                    f,
87                    "Invalid magic: {:02x}{:02x}{:02x}{:02x}",
88                    magic[0], magic[1], magic[2], magic[3]
89                )
90            }
91            Self::InvalidHeader(msg) => write!(f, "Invalid header: {msg}"),
92            Self::InvalidTensorIndex(msg) => write!(f, "Invalid tensor index: {msg}"),
93            Self::MetadataError(msg) => write!(f, "Metadata error: {msg}"),
94            Self::ChecksumMismatch => write!(f, "Checksum mismatch"),
95            Self::AlignmentError(msg) => write!(f, "Alignment error: {msg}"),
96            Self::IoError(msg) => write!(f, "I/O error: {msg}"),
97            Self::CompressionError(msg) => write!(f, "Compression error: {msg}"),
98        }
99    }
100}
101
102impl std::error::Error for V2FormatError {}
103
104// The flat-scope `v2::tests` module is declared in `v2/mod.rs` (it uses
105// `super::*` over the whole `v2` namespace), not here.
apr_format/v2/v2format_error.rs

apr_format/v2/
v2format_error.rs
