Skip to main content

apr_format/
types.rs

1//! v1 (`APRN`) container type definitions — sovereign leaf (issue #2231).
2//!
3//! Moved out of `aprender-core/src/format/` (`types.rs` + `spec.rs`). Uses
4//! [`crate::error::AprFormatError`] instead of `aprender_core::AprenderError`;
5//! `aprender-core` wraps these errors via `impl From<AprFormatError> for
6//! AprenderError` and re-exports the moved API so existing `aprender_core::
7//! format::*` paths keep working unchanged.
8//!
9//! Byte-identity (issue #2231): the field set and serde derives mirror the
10//! pre-extraction `aprender-core` types EXACTLY so the on-disk MessagePack
11//! metadata + bincode payload encodings are unchanged (the golden fixtures pin
12//! this for v1 F32).
13
14use crate::error::{AprFormatError, Result};
15use crate::model_card::ModelCard;
16use serde::{Deserialize, Serialize};
17use std::collections::HashMap;
18
19/// Magic number: "APRN" in ASCII (0x4150524E).
20pub const MAGIC: [u8; 4] = [0x41, 0x50, 0x52, 0x4E];
21
22/// Current v1 format version (1.0).
23pub const FORMAT_VERSION: (u8, u8) = (1, 0);
24
25/// v1 header size in bytes.
26pub const HEADER_SIZE: usize = 32;
27
28/// Maximum uncompressed size (1GB safety limit — compression-bomb protection).
29pub const MAX_UNCOMPRESSED_SIZE: u32 = 1024 * 1024 * 1024;
30
31/// Model type identifiers.
32#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
33#[repr(u16)]
34pub enum ModelType {
35    /// Linear regression (OLS/Ridge/Lasso)
36    LinearRegression = 0x0001,
37    /// Logistic regression (GLM Binomial)
38    LogisticRegression = 0x0002,
39    /// Decision tree (CART/ID3)
40    DecisionTree = 0x0003,
41    /// Random forest (Bagging ensemble)
42    RandomForest = 0x0004,
43    /// Gradient boosting (Boosting ensemble)
44    GradientBoosting = 0x0005,
45    /// K-means clustering (Lloyd's algorithm)
46    KMeans = 0x0006,
47    /// Principal component analysis
48    Pca = 0x0007,
49    /// Gaussian naive bayes
50    NaiveBayes = 0x0008,
51    /// K-nearest neighbors
52    Knn = 0x0009,
53    /// Support vector machine
54    Svm = 0x000A,
55    /// N-gram language model (Markov chains)
56    NgramLm = 0x0010,
57    /// TF-IDF vectorizer
58    Tfidf = 0x0011,
59    /// Count vectorizer
60    CountVectorizer = 0x0012,
61    /// Sequential neural network (Feed-forward)
62    NeuralSequential = 0x0020,
63    /// Custom neural architecture
64    NeuralCustom = 0x0021,
65    /// Content-based recommender
66    ContentRecommender = 0x0030,
67    /// Mixture of Experts (sparse/dense `MoE`)
68    MixtureOfExperts = 0x0040,
69    /// User-defined model
70    Custom = 0x00FF,
71}
72
73impl ModelType {
74    /// Convert from u16 value.
75    #[must_use]
76    pub fn from_u16(value: u16) -> Option<Self> {
77        match value {
78            0x0001 => Some(Self::LinearRegression),
79            0x0002 => Some(Self::LogisticRegression),
80            0x0003 => Some(Self::DecisionTree),
81            0x0004 => Some(Self::RandomForest),
82            0x0005 => Some(Self::GradientBoosting),
83            0x0006 => Some(Self::KMeans),
84            0x0007 => Some(Self::Pca),
85            0x0008 => Some(Self::NaiveBayes),
86            0x0009 => Some(Self::Knn),
87            0x000A => Some(Self::Svm),
88            0x0010 => Some(Self::NgramLm),
89            0x0011 => Some(Self::Tfidf),
90            0x0012 => Some(Self::CountVectorizer),
91            0x0020 => Some(Self::NeuralSequential),
92            0x0021 => Some(Self::NeuralCustom),
93            0x0030 => Some(Self::ContentRecommender),
94            0x0040 => Some(Self::MixtureOfExperts),
95            0x00FF => Some(Self::Custom),
96            _ => None,
97        }
98    }
99}
100
101/// Compression algorithm.
102#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)]
103#[repr(u8)]
104pub enum Compression {
105    /// No compression (debugging / Genchi Genbutsu)
106    None = 0x00,
107    /// Zstd level 3 (default, good balance)
108    #[default]
109    ZstdDefault = 0x01,
110    /// Zstd level 19 (maximum compression, archival)
111    ZstdMax = 0x02,
112    /// LZ4 (high-throughput streaming)
113    Lz4 = 0x03,
114}
115
116impl Compression {
117    /// Convert from u8 value.
118    #[must_use]
119    pub fn from_u8(value: u8) -> Option<Self> {
120        match value {
121            0x00 => Some(Self::None),
122            0x01 => Some(Self::ZstdDefault),
123            0x02 => Some(Self::ZstdMax),
124            0x03 => Some(Self::Lz4),
125            _ => None,
126        }
127    }
128}
129
130/// Feature flags (bitmask) — spec §3.2.
131#[derive(Debug, Clone, Copy, Default)]
132pub struct Flags(u8);
133
134impl Flags {
135    /// Payload is encrypted (AES-256-GCM)
136    pub const ENCRYPTED: u8 = 0b0000_0001;
137    /// Has digital signature (Ed25519)
138    pub const SIGNED: u8 = 0b0000_0010;
139    /// Supports chunked/streaming loading
140    pub const STREAMING: u8 = 0b0000_0100;
141    /// Has commercial license block
142    pub const LICENSED: u8 = 0b0000_1000;
143    /// 64-byte aligned tensors for zero-copy SIMD (trueno-native)
144    pub const TRUENO_NATIVE: u8 = 0b0001_0000;
145    /// Payload contains quantized tensors
146    pub const QUANTIZED: u8 = 0b0010_0000;
147    /// Has model card metadata
148    pub const HAS_MODEL_CARD: u8 = 0b0100_0000;
149
150    /// Create new empty flags.
151    #[must_use]
152    pub fn new() -> Self {
153        Self(0)
154    }
155
156    /// Set encrypted flag.
157    #[must_use]
158    pub fn with_encrypted(mut self) -> Self {
159        self.0 |= Self::ENCRYPTED;
160        self
161    }
162
163    /// Set signed flag.
164    #[must_use]
165    pub fn with_signed(mut self) -> Self {
166        self.0 |= Self::SIGNED;
167        self
168    }
169
170    /// Set streaming flag.
171    #[must_use]
172    pub fn with_streaming(mut self) -> Self {
173        self.0 |= Self::STREAMING;
174        self
175    }
176
177    /// Set licensed flag.
178    #[must_use]
179    pub fn with_licensed(mut self) -> Self {
180        self.0 |= Self::LICENSED;
181        self
182    }
183
184    /// Set trueno-native flag.
185    #[must_use]
186    pub fn with_trueno_native(mut self) -> Self {
187        self.0 |= Self::TRUENO_NATIVE;
188        self
189    }
190
191    /// Set quantized flag.
192    #[must_use]
193    pub fn with_quantized(mut self) -> Self {
194        self.0 |= Self::QUANTIZED;
195        self
196    }
197
198    /// Set model card flag.
199    #[must_use]
200    pub fn with_model_card(mut self) -> Self {
201        self.0 |= Self::HAS_MODEL_CARD;
202        self
203    }
204
205    /// Check if encrypted.
206    #[must_use]
207    pub fn is_encrypted(self) -> bool {
208        self.0 & Self::ENCRYPTED != 0
209    }
210
211    /// Check if signed.
212    #[must_use]
213    pub fn is_signed(self) -> bool {
214        self.0 & Self::SIGNED != 0
215    }
216
217    /// Check if streaming.
218    #[must_use]
219    pub fn is_streaming(self) -> bool {
220        self.0 & Self::STREAMING != 0
221    }
222
223    /// Check if licensed.
224    #[must_use]
225    pub fn is_licensed(self) -> bool {
226        self.0 & Self::LICENSED != 0
227    }
228
229    /// Check if trueno-native.
230    #[must_use]
231    pub fn is_trueno_native(self) -> bool {
232        self.0 & Self::TRUENO_NATIVE != 0
233    }
234
235    /// Check if quantized.
236    #[must_use]
237    pub fn is_quantized(self) -> bool {
238        self.0 & Self::QUANTIZED != 0
239    }
240
241    /// Check if has model card.
242    #[must_use]
243    pub fn has_model_card(self) -> bool {
244        self.0 & Self::HAS_MODEL_CARD != 0
245    }
246
247    /// Get raw value.
248    #[must_use]
249    pub fn bits(self) -> u8 {
250        self.0
251    }
252
253    /// Create from raw value (reserved high bit masked).
254    #[must_use]
255    pub fn from_bits(bits: u8) -> Self {
256        Self(bits & 0b0111_1111)
257    }
258}
259
260/// File header (32 bytes).
261#[derive(Debug, Clone)]
262pub struct Header {
263    /// Magic number (must be "APRN")
264    pub magic: [u8; 4],
265    /// Format version (major, minor)
266    pub version: (u8, u8),
267    /// Model type identifier
268    pub model_type: ModelType,
269    /// Metadata section size in bytes
270    pub metadata_size: u32,
271    /// Compressed payload size in bytes
272    pub payload_size: u32,
273    /// Uncompressed payload size (for allocation check)
274    pub uncompressed_size: u32,
275    /// Compression algorithm
276    pub compression: Compression,
277    /// Feature flags
278    pub flags: Flags,
279    /// Quality score (0-100, Poka-yoke validation) - APR-POKA-001
280    /// 0 = no validation (F), 1-59 = failing, 60-100 = passing grades
281    pub quality_score: u8,
282}
283
284impl Header {
285    /// Create a new header.
286    #[must_use]
287    pub fn new(model_type: ModelType) -> Self {
288        Self {
289            magic: MAGIC,
290            version: FORMAT_VERSION,
291            model_type,
292            metadata_size: 0,
293            payload_size: 0,
294            uncompressed_size: 0,
295            compression: Compression::default(),
296            flags: Flags::default(),
297            quality_score: 0,
298        }
299    }
300
301    /// Serialize header to bytes (32 bytes).
302    #[must_use]
303    pub fn to_bytes(&self) -> [u8; HEADER_SIZE] {
304        let mut bytes = [0u8; HEADER_SIZE];
305        bytes[0..4].copy_from_slice(&self.magic);
306        bytes[4] = self.version.0;
307        bytes[5] = self.version.1;
308        let model_type = self.model_type as u16;
309        bytes[6..8].copy_from_slice(&model_type.to_le_bytes());
310        bytes[8..12].copy_from_slice(&self.metadata_size.to_le_bytes());
311        bytes[12..16].copy_from_slice(&self.payload_size.to_le_bytes());
312        bytes[16..20].copy_from_slice(&self.uncompressed_size.to_le_bytes());
313        bytes[20] = self.compression as u8;
314        bytes[21] = self.flags.bits();
315        bytes[22] = self.quality_score;
316        // Reserved (23-31) already zero.
317        bytes
318    }
319
320    /// Parse header from bytes.
321    ///
322    /// # Errors
323    /// Returns an error on short input, bad magic, unsupported version, unknown
324    /// model type, unknown compression, or a compression-bomb-sized payload.
325    pub fn from_bytes(bytes: &[u8]) -> Result<Self> {
326        if bytes.len() < HEADER_SIZE {
327            return Err(AprFormatError::FormatError {
328                message: format!(
329                    "Header too short: {} bytes, expected {}",
330                    bytes.len(),
331                    HEADER_SIZE
332                ),
333            });
334        }
335
336        let magic: [u8; 4] = bytes[0..4]
337            .try_into()
338            .map_err(|_| AprFormatError::FormatError {
339                message: "header slice too short for magic".to_string(),
340            })?;
341        if magic != MAGIC {
342            return Err(AprFormatError::FormatError {
343                message: format!(
344                    "Invalid magic number: {:02X}{:02X}{:02X}{:02X}, expected APRN",
345                    magic[0], magic[1], magic[2], magic[3]
346                ),
347            });
348        }
349
350        let version = (bytes[4], bytes[5]);
351        if version.0 > FORMAT_VERSION.0 {
352            return Err(AprFormatError::UnsupportedVersion {
353                found: version,
354                supported: FORMAT_VERSION,
355            });
356        }
357
358        let model_type_raw = u16::from_le_bytes([bytes[6], bytes[7]]);
359        let model_type =
360            ModelType::from_u16(model_type_raw).ok_or_else(|| AprFormatError::FormatError {
361                message: format!("Unknown model type: 0x{model_type_raw:04X}"),
362            })?;
363
364        let metadata_size = u32::from_le_bytes([bytes[8], bytes[9], bytes[10], bytes[11]]);
365        let payload_size = u32::from_le_bytes([bytes[12], bytes[13], bytes[14], bytes[15]]);
366        let uncompressed_size = u32::from_le_bytes([bytes[16], bytes[17], bytes[18], bytes[19]]);
367
368        if uncompressed_size > MAX_UNCOMPRESSED_SIZE {
369            return Err(AprFormatError::FormatError {
370                message: format!(
371                    "Uncompressed size {uncompressed_size} exceeds maximum {MAX_UNCOMPRESSED_SIZE} (compression bomb protection)"
372                ),
373            });
374        }
375
376        let compression =
377            Compression::from_u8(bytes[20]).ok_or_else(|| AprFormatError::FormatError {
378                message: format!("Unknown compression algorithm: 0x{:02X}", bytes[20]),
379            })?;
380        let flags = Flags::from_bits(bytes[21]);
381        let quality_score = bytes[22];
382
383        Ok(Self {
384            magic,
385            version,
386            model_type,
387            metadata_size,
388            payload_size,
389            uncompressed_size,
390            compression,
391            flags,
392            quality_score,
393        })
394    }
395}
396
397/// Training information.
398#[derive(Debug, Clone, Serialize, Deserialize)]
399pub struct TrainingInfo {
400    /// Number of training samples
401    #[serde(default, skip_serializing_if = "Option::is_none")]
402    pub samples: Option<usize>,
403    /// Training duration in milliseconds
404    #[serde(default, skip_serializing_if = "Option::is_none")]
405    pub duration_ms: Option<u64>,
406    /// Data source description
407    #[serde(default, skip_serializing_if = "Option::is_none")]
408    pub source: Option<String>,
409}
410
411// ============================================================================
412// Knowledge Distillation Types (spec §6.3)
413// ============================================================================
414
415/// Distillation method used (spec §6.3.1)
416#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
417pub enum DistillMethod {
418    /// KL divergence on final logits (Hinton2015)
419    Standard,
420    /// Intermediate layer matching
421    Progressive,
422    /// Multiple teachers weighted average
423    Ensemble,
424}
425
426/// Teacher model provenance for audit trails (spec §6.3.2)
427#[derive(Debug, Clone, Serialize, Deserialize)]
428pub struct TeacherProvenance {
429    /// SHA256 hash of teacher .apr file
430    pub hash: String,
431    /// Ed25519 signature of teacher (if signed)
432    #[serde(default, skip_serializing_if = "Option::is_none")]
433    pub signature: Option<String>,
434    /// Teacher model type
435    pub model_type: ModelType,
436    /// Teacher parameter count
437    pub param_count: u64,
438    /// For ensemble: multiple teachers
439    #[serde(default, skip_serializing_if = "Option::is_none")]
440    pub ensemble_teachers: Option<Vec<TeacherProvenance>>,
441}
442
443/// Distillation hyperparameters (spec §6.3.2)
444#[derive(Debug, Clone, Serialize, Deserialize)]
445pub struct DistillationParams {
446    /// Temperature for softening distributions (typically 2.0-5.0)
447    pub temperature: f32,
448    /// Weight for soft vs hard loss (α in loss formula)
449    pub alpha: f32,
450    /// For progressive: weight for hidden vs logit loss (β)
451    #[serde(default, skip_serializing_if = "Option::is_none")]
452    pub beta: Option<f32>,
453    /// Training epochs for distillation
454    pub epochs: u32,
455    /// Final distillation loss achieved
456    #[serde(default, skip_serializing_if = "Option::is_none")]
457    pub final_loss: Option<f32>,
458}
459
460/// Layer mapping for progressive distillation (spec §6.3.2)
461#[derive(Debug, Clone, Serialize, Deserialize)]
462pub struct LayerMapping {
463    /// Student layer index
464    pub student_layer: usize,
465    /// Teacher layer index
466    pub teacher_layer: usize,
467    /// Weight for this layer's loss
468    pub weight: f32,
469}
470
471/// Complete distillation provenance (spec §6.3.2)
472#[derive(Debug, Clone, Serialize, Deserialize)]
473pub struct DistillationInfo {
474    /// Distillation method used
475    pub method: DistillMethod,
476    /// Teacher model provenance
477    pub teacher: TeacherProvenance,
478    /// Distillation hyperparameters
479    pub params: DistillationParams,
480    /// Optional: layer mapping for progressive distillation
481    #[serde(default, skip_serializing_if = "Option::is_none")]
482    pub layer_mapping: Option<Vec<LayerMapping>>,
483}
484
485// ============================================================================
486// Commercial License Types (spec §9)
487// ============================================================================
488
489/// License tier levels (spec §9.1).
490#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
491pub enum LicenseTier {
492    /// Personal/individual use
493    Personal,
494    /// Team/organization use (limited seats)
495    Team,
496    /// Enterprise use (unlimited seats, priority support)
497    Enterprise,
498    /// Academic/research use (non-commercial)
499    Academic,
500}
501
502/// Commercial license information (spec §9.1).
503#[derive(Debug, Clone, Serialize, Deserialize)]
504pub struct LicenseInfo {
505    /// Unique license identifier (UUID v4)
506    pub uuid: String,
507    /// Hash of the license certificate (cryptographically bound)
508    pub hash: String,
509    /// License expiration date (ISO 8601) — None for perpetual
510    #[serde(default, skip_serializing_if = "Option::is_none")]
511    pub expiry: Option<String>,
512    /// Maximum concurrent seats — None for unlimited
513    #[serde(default, skip_serializing_if = "Option::is_none")]
514    pub seats: Option<u32>,
515    /// Licensee name/organization
516    #[serde(default, skip_serializing_if = "Option::is_none")]
517    pub licensee: Option<String>,
518    /// License tier
519    pub tier: LicenseTier,
520}
521
522/// Model metadata (MessagePack-encoded).
523#[derive(Debug, Clone, Serialize, Deserialize)]
524pub struct Metadata {
525    /// Creation timestamp (ISO 8601)
526    pub created_at: String,
527    /// Aprender version that created this model
528    pub aprender_version: String,
529    /// Optional model name
530    #[serde(default, skip_serializing_if = "Option::is_none")]
531    pub model_name: Option<String>,
532    /// Optional description
533    #[serde(default, skip_serializing_if = "Option::is_none")]
534    pub description: Option<String>,
535    /// Training information
536    #[serde(default, skip_serializing_if = "Option::is_none")]
537    pub training: Option<TrainingInfo>,
538    /// Hyperparameters
539    #[serde(default, skip_serializing_if = "HashMap::is_empty")]
540    pub hyperparameters: HashMap<String, serde_json::Value>,
541    /// Model metrics
542    #[serde(default, skip_serializing_if = "HashMap::is_empty")]
543    pub metrics: HashMap<String, serde_json::Value>,
544    /// Custom user data
545    #[serde(default, skip_serializing_if = "HashMap::is_empty")]
546    pub custom: HashMap<String, serde_json::Value>,
547    /// Distillation teacher hash (spec §6.3) - simple form
548    #[serde(default, skip_serializing_if = "Option::is_none")]
549    pub distillation: Option<String>,
550    /// Full distillation provenance (spec §6.3.2) - structured form
551    #[serde(default, skip_serializing_if = "Option::is_none")]
552    pub distillation_info: Option<DistillationInfo>,
553    /// Commercial license information (spec §9.1)
554    #[serde(default, skip_serializing_if = "Option::is_none")]
555    pub license: Option<LicenseInfo>,
556    /// Model card metadata (spec §11)
557    #[serde(default, skip_serializing_if = "Option::is_none")]
558    pub model_card: Option<ModelCard>,
559}
560
561impl Default for Metadata {
562    fn default() -> Self {
563        Self {
564            created_at: chrono_lite_now(),
565            aprender_version: env!("CARGO_PKG_VERSION").to_string(),
566            model_name: None,
567            description: None,
568            training: None,
569            hyperparameters: HashMap::new(),
570            metrics: HashMap::new(),
571            custom: HashMap::new(),
572            distillation: None,
573            distillation_info: None,
574            license: None,
575            model_card: None,
576        }
577    }
578}
579
580/// Simple ISO 8601 timestamp (no chrono dependency).
581#[must_use]
582pub fn chrono_lite_now() -> String {
583    use std::time::{SystemTime, UNIX_EPOCH};
584    let duration = SystemTime::now()
585        .duration_since(UNIX_EPOCH)
586        .unwrap_or_default();
587    let secs = duration.as_secs();
588    // Convert to rough ISO 8601 (good enough for metadata)
589    format!("{secs}")
590}
591
592/// Options for saving models.
593#[derive(Debug, Clone, Default)]
594pub struct SaveOptions {
595    /// Compression algorithm
596    pub compression: Compression,
597    /// Additional metadata
598    pub metadata: Metadata,
599    /// Quality score from Poka-yoke validation (APR-POKA-001).
600    /// - `None`: no validation performed (score=0 in file)
601    /// - `Some(0)`: explicit failure — save will be REFUSED (Jidoka)
602    /// - `Some(1-59)`: validation failed but allowed to save
603    /// - `Some(60-100)`: validation passed
604    pub quality_score: Option<u8>,
605}
606
607impl SaveOptions {
608    /// Create with default compression.
609    #[must_use]
610    pub fn new() -> Self {
611        Self::default()
612    }
613
614    /// Set compression algorithm.
615    #[must_use]
616    pub fn with_compression(mut self, compression: Compression) -> Self {
617        self.compression = compression;
618        self
619    }
620
621    /// Set model name.
622    #[must_use]
623    pub fn with_name(mut self, name: impl Into<String>) -> Self {
624        self.metadata.model_name = Some(name.into());
625        self
626    }
627
628    /// Set description.
629    #[must_use]
630    pub fn with_description(mut self, desc: impl Into<String>) -> Self {
631        self.metadata.description = Some(desc.into());
632        self
633    }
634
635    /// Set distillation info (spec §6.3).
636    #[must_use]
637    pub fn with_distillation_info(mut self, info: DistillationInfo) -> Self {
638        self.metadata.distillation_info = Some(info);
639        self
640    }
641
642    /// Set license info (spec §9.1).
643    #[must_use]
644    pub fn with_license(mut self, license: LicenseInfo) -> Self {
645        self.metadata.license = Some(license);
646        self
647    }
648
649    /// Set model card (spec §11).
650    #[must_use]
651    pub fn with_model_card(mut self, card: ModelCard) -> Self {
652        self.metadata.model_card = Some(card);
653        self
654    }
655
656    /// Set quality score from Poka-yoke validation (APR-POKA-001).
657    ///
658    /// # Jidoka (Stop the Line)
659    /// - Score 0 will cause `save()` to REFUSE the write
660    /// - Score 1-59 allows save with warning
661    /// - Score 60-100 is passing
662    #[must_use]
663    pub fn with_quality_score(mut self, score: u8) -> Self {
664        self.quality_score = Some(score);
665        self
666    }
667}
668
669/// Model information (from inspection).
670#[derive(Debug, Clone)]
671#[allow(clippy::struct_excessive_bools)] // Bools represent independent flag states
672pub struct ModelInfo {
673    /// Model type
674    pub model_type: ModelType,
675    /// Format version
676    pub format_version: (u8, u8),
677    /// Metadata
678    pub metadata: Metadata,
679    /// Compressed payload size
680    pub payload_size: usize,
681    /// Uncompressed payload size
682    pub uncompressed_size: usize,
683    /// Is encrypted
684    pub encrypted: bool,
685    /// Is signed
686    pub signed: bool,
687    /// Is streaming
688    pub streaming: bool,
689    /// Has commercial license block
690    pub licensed: bool,
691    /// Uses trueno-native 64-byte aligned tensors
692    pub trueno_native: bool,
693    /// Contains quantized tensors
694    pub quantized: bool,
695    /// Has model card metadata (spec §11)
696    pub has_model_card: bool,
697}