infotheory/spec/
core.rs

1//! Canonical validated and compiled backend plans.
2//!
3//! This module is the internal/public bridge between wrapper AST specs
4//! (`RateBackend`, `CompressionBackend`) and the compiled runtime plans used by
5//! the generic Rust API.
6
7use crate::api::{
8    CalibratedSpec, CalibrationContextKind, CompressionBackend, MAX_MIXTURE_NESTING,
9    MixtureExpertSpec, MixtureKind, MixtureScheduleMode, MixtureSpec, ParticleSpec, RateBackend,
10};
11use crate::coders::CoderType;
12use crate::compression::FramingMode;
13use crate::spec::{SpecError, SpecResult};
14use std::fmt;
15use std::path::{Path, PathBuf};
16use std::sync::Arc;
17
18mod adapters;
19mod compile;
20pub(crate) use adapters::*;
21pub(crate) use compile::*;
22
23/// Compilation environment for backend/spec validation and canonicalization.
24#[derive(Clone, Debug, Default, Eq, PartialEq)]
25#[non_exhaustive]
26pub struct SpecEnvironment {
27    base_dir: PathBuf,
28}
29
30impl SpecEnvironment {
31    /// Create an environment rooted at `base_dir`.
32    pub fn new(base_dir: impl Into<PathBuf>) -> Self {
33        Self {
34            base_dir: base_dir.into(),
35        }
36    }
37
38    /// Base directory used to resolve relative asset references.
39    pub fn base_dir(&self) -> &Path {
40        &self.base_dir
41    }
42}
43
44/// Typed external asset reference captured by a validated/compiled spec.
45#[derive(Clone, Debug, Eq, PartialEq, Hash)]
46#[non_exhaustive]
47pub enum AssetRef {
48    /// Filesystem-backed asset resolved relative to a [`SpecEnvironment`].
49    Filesystem(PathBuf),
50}
51
52/// Deterministic binary canonical code for a validated spec.
53#[derive(Clone, Eq, PartialEq, Hash)]
54pub struct CanonicalBytes(Arc<[u8]>);
55
56impl CanonicalBytes {
57    /// View the canonical bytes.
58    pub fn as_slice(&self) -> &[u8] {
59        &self.0
60    }
61
62    /// Length of the canonical code in bytes.
63    pub fn len(&self) -> usize {
64        self.0.len()
65    }
66
67    /// Whether the canonical code is empty.
68    pub fn is_empty(&self) -> bool {
69        self.0.is_empty()
70    }
71}
72
73impl fmt::Debug for CanonicalBytes {
74    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
75        f.debug_tuple("CanonicalBytes")
76            .field(&self.0.len())
77            .finish()
78    }
79}
80
81impl From<Vec<u8>> for CanonicalBytes {
82    fn from(value: Vec<u8>) -> Self {
83        Self(Arc::<[u8]>::from(value))
84    }
85}
86
87/// Shared execution family for method-backed neural backends.
88#[derive(Clone, Copy, Debug, Eq, PartialEq, Hash)]
89#[non_exhaustive]
90pub enum MethodBackendFamily {
91    /// Mamba family.
92    Mamba,
93    /// RWKV-7 family.
94    Rwkv7,
95}
96
97/// Trace-model execution strategy used by VM/AIXI adapters.
98#[derive(Clone, Copy, Debug, Eq, PartialEq, Hash)]
99#[non_exhaustive]
100pub enum RateBackendTraceStrategy {
101    /// Direct ROSA-specific strategy.
102    Rosa,
103    /// Direct CTW strategy.
104    Ctw,
105    /// Direct FAC-CTW strategy.
106    FacCtw,
107    /// Generic predictor-backed strategy.
108    PredictorBacked,
109    /// ZPAQ-specific rate model strategy.
110    Zpaq,
111    /// Mamba compressor-backed strategy.
112    Mamba,
113    /// RWKV compressor-backed strategy.
114    Rwkv7,
115}
116
117/// Shared capability metadata for rate backends.
118#[derive(Clone, Debug)]
119#[non_exhaustive]
120pub struct RateBackendCapabilities {
121    /// Canonical backend family name.
122    pub canonical_name: &'static str,
123    /// Human-readable backend label.
124    pub display_label: Arc<str>,
125    /// VM/AIXI trace strategy for this backend family.
126    pub trace_strategy: RateBackendTraceStrategy,
127    /// Whether biased/plugin entropy is supported.
128    pub supports_biased_entropy: bool,
129    /// Whether generic frozen-conditioning is supported.
130    pub supports_frozen_conditioning: bool,
131    /// Whether generic rate-coded compression wrapping is supported.
132    pub supports_rate_coded_compression: bool,
133    /// Whether this backend has a direct native bit predictor rather than a
134    /// byte-symbol adaptation over `{0,1}`.
135    pub supports_native_bit_prediction: bool,
136    /// Whether this backend can expose its byte PDF as a lazy binary prefix mass.
137    pub supports_byte_prefix_mass: bool,
138    /// Whether repeated byte-packed bit-session prefix queries remain practical
139    /// without a replay-heavy fallback.
140    pub supports_efficient_byte_packed_bit_sessions: bool,
141    /// Whether bit observations can be undone exactly after update.
142    pub supports_reversible_bit_updates: bool,
143    /// Whether the backend graph contains any ZPAQ component.
144    pub contains_zpaq: bool,
145    /// Whether this is a method-backed neural family.
146    pub method_family: Option<MethodBackendFamily>,
147}
148
149/// Shared capability metadata for compression backends.
150#[derive(Clone, Debug)]
151#[non_exhaustive]
152pub struct CompressionBackendCapabilities {
153    /// Canonical backend family name.
154    pub canonical_name: &'static str,
155    /// Human-readable backend label.
156    pub display_label: Arc<str>,
157    /// Whether this backend wraps a predictive rate backend.
158    pub uses_rate_backend: bool,
159    /// Whether this backend supports decompression.
160    pub supports_decompression: bool,
161}
162
163#[derive(Clone, Debug)]
164pub(crate) struct RateBackendPlanExpert {
165    pub name: Option<String>,
166    pub log_prior: f64,
167    pub backend: Arc<RateBackendPlan>,
168}
169
170#[derive(Clone, Debug)]
171pub(crate) enum RateBackendPlan {
172    RosaPlus {
173        max_order: i64,
174    },
175    Match {
176        hash_bits: usize,
177        min_len: usize,
178        max_len: usize,
179        base_mix: f64,
180        confidence_scale: f64,
181    },
182    SparseMatch {
183        hash_bits: usize,
184        min_len: usize,
185        max_len: usize,
186        gap_min: usize,
187        gap_max: usize,
188        base_mix: f64,
189        confidence_scale: f64,
190    },
191    Ppmd {
192        order: usize,
193        memory_mb: usize,
194    },
195    Sequitur {
196        context_bytes: usize,
197    },
198    Ctw {
199        depth: usize,
200    },
201    FacCtw {
202        base_depth: usize,
203        num_percept_bits: usize,
204        encoding_bits: usize,
205        msb_first: bool,
206    },
207    Zpaq {
208        method: String,
209    },
210    #[cfg(feature = "backend-mamba")]
211    Mamba {
212        method: String,
213        parsed_method: crate::mambazip::MethodSpec,
214        asset: Option<AssetRef>,
215    },
216    #[cfg(feature = "backend-rwkv")]
217    Rwkv7 {
218        method: String,
219        parsed_method: crate::rwkvzip::MethodSpec,
220        asset: Option<AssetRef>,
221    },
222    Mixture {
223        kind: MixtureKind,
224        schedule: MixtureScheduleMode,
225        alpha: f64,
226        decay: Option<f64>,
227        experts: Box<[RateBackendPlanExpert]>,
228    },
229    Particle {
230        spec: ParticleSpec,
231    },
232    Calibrated {
233        context: CalibrationContextKind,
234        bins: usize,
235        learning_rate: f64,
236        bias_clip: f64,
237        base: Arc<RateBackendPlan>,
238    },
239}
240
241impl RateBackendPlan {
242    pub(crate) fn kind(&self) -> crate::runtime::RateBackendKind {
243        match self {
244            RateBackendPlan::RosaPlus { .. } => crate::runtime::RateBackendKind::RosaPlus,
245            RateBackendPlan::Match { .. } => crate::runtime::RateBackendKind::Match,
246            RateBackendPlan::SparseMatch { .. } => crate::runtime::RateBackendKind::SparseMatch,
247            RateBackendPlan::Ppmd { .. } => crate::runtime::RateBackendKind::Ppmd,
248            RateBackendPlan::Sequitur { .. } => crate::runtime::RateBackendKind::Sequitur,
249            RateBackendPlan::Ctw { .. } => crate::runtime::RateBackendKind::Ctw,
250            RateBackendPlan::FacCtw { .. } => crate::runtime::RateBackendKind::FacCtw,
251            RateBackendPlan::Zpaq { .. } => crate::runtime::RateBackendKind::Zpaq,
252            #[cfg(feature = "backend-mamba")]
253            RateBackendPlan::Mamba { .. } => crate::runtime::RateBackendKind::Mamba,
254            #[cfg(feature = "backend-rwkv")]
255            RateBackendPlan::Rwkv7 { .. } => crate::runtime::RateBackendKind::Rwkv7,
256            RateBackendPlan::Mixture { .. } => crate::runtime::RateBackendKind::Mixture,
257            RateBackendPlan::Particle { .. } => crate::runtime::RateBackendKind::Particle,
258            RateBackendPlan::Calibrated { .. } => crate::runtime::RateBackendKind::Calibrated,
259        }
260    }
261}
262
263#[derive(Clone, Debug)]
264pub(crate) enum CompressionBackendPlan {
265    Zpaq {
266        method: String,
267        threads: usize,
268    },
269    #[cfg(feature = "backend-rwkv")]
270    Rwkv7 {
271        method: String,
272        parsed_method: crate::rwkvzip::MethodSpec,
273        asset: Option<AssetRef>,
274        coder: CoderType,
275    },
276    Rate {
277        rate_backend: Arc<RateBackendPlan>,
278        coder: CoderType,
279        framing: FramingMode,
280    },
281}
282
283impl CompressionBackendPlan {
284    pub(crate) fn kind(&self) -> crate::runtime::CompressionBackendKind {
285        match self {
286            CompressionBackendPlan::Zpaq { .. } => crate::runtime::CompressionBackendKind::Zpaq,
287            #[cfg(feature = "backend-rwkv")]
288            CompressionBackendPlan::Rwkv7 { .. } => crate::runtime::CompressionBackendKind::Rwkv7,
289            CompressionBackendPlan::Rate {
290                coder: CoderType::AC,
291                ..
292            } => crate::runtime::CompressionBackendKind::RateAc,
293            CompressionBackendPlan::Rate {
294                coder: CoderType::RANS,
295                ..
296            } => crate::runtime::CompressionBackendKind::RateRans,
297        }
298    }
299}
300
301/// Canonicalized and feature-validated rate backend wrapper.
302#[derive(Clone)]
303pub struct ValidatedRateBackend {
304    canonical_spec: Arc<RateBackend>,
305    canonical_bytes: CanonicalBytes,
306    capabilities: RateBackendCapabilities,
307    plan: Arc<RateBackendPlan>,
308}
309
310impl ValidatedRateBackend {
311    /// Canonical wrapper AST for this backend.
312    pub fn canonical_spec(&self) -> &RateBackend {
313        self.canonical_spec.as_ref()
314    }
315
316    /// Deterministic binary canonical code for this backend.
317    pub fn canonical_bytes(&self) -> &CanonicalBytes {
318        &self.canonical_bytes
319    }
320
321    /// Capability metadata derived from the canonicalized backend graph.
322    pub fn capabilities(&self) -> &RateBackendCapabilities {
323        &self.capabilities
324    }
325
326    /// Human-readable backend label derived from the compiled plan.
327    pub fn display_label(&self) -> String {
328        rate_backend_plan_display_label(self.plan.as_ref())
329    }
330
331    /// Short default backend name for logs, diagnostics, and model labels.
332    pub fn default_name(&self) -> String {
333        rate_backend_plan_default_name(self.plan.as_ref())
334    }
335
336    /// Compile the validated spec into an immutable runtime plan.
337    pub fn compile(&self) -> SpecResult<CompiledRateBackend> {
338        Ok(CompiledRateBackend {
339            canonical_spec: self.canonical_spec.clone(),
340            canonical_bytes: self.canonical_bytes.clone(),
341            capabilities: self.capabilities.clone(),
342            plan: self.plan.clone(),
343        })
344    }
345}
346
347/// Canonicalized and feature-validated compression backend wrapper.
348#[derive(Clone)]
349pub struct ValidatedCompressionBackend {
350    canonical_spec: Arc<CompressionBackend>,
351    canonical_bytes: CanonicalBytes,
352    capabilities: CompressionBackendCapabilities,
353    plan: Arc<CompressionBackendPlan>,
354}
355
356impl ValidatedCompressionBackend {
357    /// Canonical wrapper AST for this backend.
358    pub fn canonical_spec(&self) -> &CompressionBackend {
359        self.canonical_spec.as_ref()
360    }
361
362    /// Deterministic binary canonical code for this backend.
363    pub fn canonical_bytes(&self) -> &CanonicalBytes {
364        &self.canonical_bytes
365    }
366
367    /// Capability metadata derived from the canonicalized backend graph.
368    pub fn capabilities(&self) -> &CompressionBackendCapabilities {
369        &self.capabilities
370    }
371
372    /// Compile the validated spec into an immutable runtime plan.
373    pub fn compile(&self) -> SpecResult<CompiledCompressionBackend> {
374        Ok(CompiledCompressionBackend {
375            canonical_spec: self.canonical_spec.clone(),
376            canonical_bytes: self.canonical_bytes.clone(),
377            capabilities: self.capabilities.clone(),
378            plan: self.plan.clone(),
379        })
380    }
381}
382
383/// Compiled immutable rate-backend runtime plan.
384#[derive(Clone)]
385pub struct CompiledRateBackend {
386    canonical_spec: Arc<RateBackend>,
387    canonical_bytes: CanonicalBytes,
388    capabilities: RateBackendCapabilities,
389    plan: Arc<RateBackendPlan>,
390}
391
392impl CompiledRateBackend {
393    /// Canonical wrapper AST for this backend.
394    pub fn canonical_spec(&self) -> &RateBackend {
395        self.canonical_spec.as_ref()
396    }
397
398    /// Deterministic binary canonical code for this backend.
399    pub fn canonical_bytes(&self) -> &CanonicalBytes {
400        &self.canonical_bytes
401    }
402
403    /// Capability metadata derived from the compiled backend graph.
404    pub fn capabilities(&self) -> &RateBackendCapabilities {
405        &self.capabilities
406    }
407
408    /// Compile-friendly backend display label.
409    pub fn display_label(&self) -> String {
410        rate_backend_plan_display_label(self.plan.as_ref())
411    }
412
413    /// Short default backend name for logs, diagnostics, and model labels.
414    pub fn default_name(&self) -> String {
415        rate_backend_plan_default_name(self.plan.as_ref())
416    }
417
418    /// Canonical backend family name.
419    pub fn canonical_name(&self) -> &'static str {
420        self.capabilities.canonical_name
421    }
422
423    /// Whether the compiled backend graph contains any ZPAQ component.
424    pub fn contains_zpaq(&self) -> bool {
425        self.capabilities.contains_zpaq
426    }
427
428    /// Whether this backend supports generic frozen conditioning.
429    pub fn supports_frozen_conditioning(&self) -> bool {
430        self.capabilities.supports_frozen_conditioning
431    }
432
433    /// Whether this backend supports generic rate-coded compression wrapping.
434    pub fn supports_rate_coded_compression(&self) -> bool {
435        self.capabilities.supports_rate_coded_compression
436    }
437
438    /// Whether this backend has a direct native bit predictor rather than a
439    /// byte-symbol adaptation over `{0,1}`.
440    pub fn supports_native_bit_prediction(&self) -> bool {
441        self.capabilities.supports_native_bit_prediction
442    }
443
444    /// Whether this backend can expose its byte PDF as a lazy binary prefix mass.
445    pub fn supports_byte_prefix_mass(&self) -> bool {
446        self.capabilities.supports_byte_prefix_mass
447    }
448
449    /// Whether this backend can support repeated byte-packed bit-session prefix
450    /// queries without pathological replay-heavy fallback.
451    pub fn supports_efficient_byte_packed_bit_sessions(&self) -> bool {
452        self.capabilities
453            .supports_efficient_byte_packed_bit_sessions
454    }
455
456    /// Whether bit observations can be undone exactly after update.
457    pub fn supports_reversible_bit_updates(&self) -> bool {
458        self.capabilities.supports_reversible_bit_updates
459    }
460
461    pub(crate) fn plan(&self) -> &RateBackendPlan {
462        self.plan.as_ref()
463    }
464
465    #[allow(dead_code)]
466    pub(crate) fn method_string(&self) -> Option<&str> {
467        match self.plan() {
468            #[cfg(feature = "backend-mamba")]
469            RateBackendPlan::Mamba { method, .. } => Some(method.as_str()),
470            #[cfg(feature = "backend-rwkv")]
471            RateBackendPlan::Rwkv7 { method, .. } => Some(method.as_str()),
472            _ => None,
473        }
474    }
475}
476
477/// Compiled immutable compression-backend runtime plan.
478#[derive(Clone)]
479pub struct CompiledCompressionBackend {
480    canonical_spec: Arc<CompressionBackend>,
481    canonical_bytes: CanonicalBytes,
482    capabilities: CompressionBackendCapabilities,
483    plan: Arc<CompressionBackendPlan>,
484}
485
486impl CompiledCompressionBackend {
487    /// Canonical wrapper AST for this backend.
488    pub fn canonical_spec(&self) -> &CompressionBackend {
489        self.canonical_spec.as_ref()
490    }
491
492    /// Deterministic binary canonical code for this backend.
493    pub fn canonical_bytes(&self) -> &CanonicalBytes {
494        &self.canonical_bytes
495    }
496
497    /// Capability metadata derived from the compiled backend graph.
498    pub fn capabilities(&self) -> &CompressionBackendCapabilities {
499        &self.capabilities
500    }
501
502    pub(crate) fn plan(&self) -> &CompressionBackendPlan {
503        self.plan.as_ref()
504    }
505}
506
507pub(crate) fn validate_rate_backend_in(
508    backend: &RateBackend,
509    env: &SpecEnvironment,
510) -> SpecResult<ValidatedRateBackend> {
511    let plan = Arc::new(build_rate_plan(backend, env, MAX_MIXTURE_NESTING)?);
512    let canonical_spec = Arc::new(rate_plan_to_wrapper(plan.as_ref()));
513    crate::api::validate_rate_backend(canonical_spec.as_ref())
514        .map_err(|err| SpecError::new(err.to_string()))?;
515    Ok(ValidatedRateBackend {
516        canonical_spec,
517        canonical_bytes: encode_rate_backend_plan(plan.as_ref()),
518        capabilities: rate_backend_capabilities(plan.as_ref()),
519        plan,
520    })
521}
522
523pub(crate) fn validate_compression_backend_in(
524    backend: &CompressionBackend,
525    env: &SpecEnvironment,
526) -> SpecResult<ValidatedCompressionBackend> {
527    let plan = Arc::new(build_compression_plan(backend, env)?);
528    let canonical_spec = Arc::new(compression_plan_to_wrapper(plan.as_ref()));
529    crate::api::validate_compression_backend(canonical_spec.as_ref())
530        .map_err(|err| SpecError::new(err.to_string()))?;
531    Ok(ValidatedCompressionBackend {
532        canonical_spec,
533        canonical_bytes: encode_compression_backend_plan(plan.as_ref()),
534        capabilities: compression_backend_capabilities(plan.as_ref()),
535        plan,
536    })
537}
538
539pub(crate) fn compiled_rate_backend_from_plan(
540    plan: Arc<RateBackendPlan>,
541) -> SpecResult<CompiledRateBackend> {
542    let canonical_spec = Arc::new(rate_plan_to_wrapper(plan.as_ref()));
543    crate::api::validate_rate_backend(canonical_spec.as_ref())
544        .map_err(|err| SpecError::new(err.to_string()))?;
545    Ok(CompiledRateBackend {
546        canonical_bytes: encode_rate_backend_plan(plan.as_ref()),
547        capabilities: rate_backend_capabilities(plan.as_ref()),
548        canonical_spec,
549        plan,
550    })
551}
552
553pub(crate) fn compiled_compression_backend_from_plan(
554    plan: Arc<CompressionBackendPlan>,
555) -> SpecResult<CompiledCompressionBackend> {
556    let canonical_spec = Arc::new(compression_plan_to_wrapper(plan.as_ref()));
557    crate::api::validate_compression_backend(canonical_spec.as_ref())
558        .map_err(|err| SpecError::new(err.to_string()))?;
559    Ok(CompiledCompressionBackend {
560        canonical_bytes: encode_compression_backend_plan(plan.as_ref()),
561        capabilities: compression_backend_capabilities(plan.as_ref()),
562        canonical_spec,
563        plan,
564    })
565}
566
567fn build_rate_plan(
568    backend: &RateBackend,
569    env: &SpecEnvironment,
570    depth: usize,
571) -> SpecResult<RateBackendPlan> {
572    if depth == 0 {
573        return Err(SpecError::new("backend spec nesting too deep"));
574    }
575
576    let descriptor = backend
577        .descriptor()
578        .map_err(|err| SpecError::new(format!("{err} (while validating rate backend)")))?;
579    if !descriptor.enabled {
580        let Some(feature) = descriptor.feature else {
581            return Err(SpecError::new(format!(
582                "internal backend registry mismatch: disabled backend '{}' is missing required feature metadata",
583                descriptor.canonical
584            )));
585        };
586        return Err(SpecError::new(format!(
587            "backend '{}' requires infotheory feature '{}'",
588            descriptor.canonical, feature
589        )));
590    }
591
592    crate::runtime::compile_rate_backend_plan_via_kernel(descriptor.kind, backend, env, depth)
593}
594
595fn build_compression_plan(
596    backend: &CompressionBackend,
597    env: &SpecEnvironment,
598) -> SpecResult<CompressionBackendPlan> {
599    let descriptor = backend
600        .descriptor()
601        .map_err(|err| SpecError::new(format!("{err} (while validating compression backend)")))?;
602    if !descriptor.enabled {
603        let Some(feature) = descriptor.feature else {
604            return Err(SpecError::new(format!(
605                "internal backend registry mismatch: disabled compression backend '{}' is missing required feature metadata",
606                descriptor.canonical
607            )));
608        };
609        return Err(SpecError::new(format!(
610            "compression backend '{}' requires infotheory feature '{}'",
611            descriptor.canonical, feature
612        )));
613    }
614
615    crate::runtime::compile_compression_backend_plan_via_kernel(descriptor.kind, backend, env)
616}
617
618pub(crate) fn rate_plan_to_wrapper(plan: &RateBackendPlan) -> RateBackend {
619    crate::runtime::rate_backend_wrapper_via_kernel(plan)
620}
621
622fn compression_plan_to_wrapper(plan: &CompressionBackendPlan) -> CompressionBackend {
623    crate::runtime::compression_backend_wrapper_via_kernel(plan)
624}
625
626pub(crate) fn rate_backend_capabilities(plan: &RateBackendPlan) -> RateBackendCapabilities {
627    crate::runtime::rate_backend_capabilities_via_kernel(plan)
628}
629
630fn compression_backend_capabilities(
631    plan: &CompressionBackendPlan,
632) -> CompressionBackendCapabilities {
633    crate::runtime::compression_backend_capabilities_via_kernel(plan)
634}
635
636fn rate_plan_contains_zpaq(plan: &RateBackendPlan) -> bool {
637    (crate::runtime::rate_backend_kernel(plan.kind()).contains_zpaq)(plan)
638}
639
640fn rate_backend_plan_display_label(plan: &RateBackendPlan) -> String {
641    crate::runtime::rate_backend_display_label_via_kernel(plan)
642}
643
644fn rate_backend_plan_default_name(plan: &RateBackendPlan) -> String {
645    crate::runtime::rate_backend_default_name_via_kernel(plan)
646}
647
648#[cfg(feature = "backend-rwkv")]
649fn rwkv_asset_ref(spec: &crate::rwkvzip::MethodSpec) -> Option<AssetRef> {
650    match spec {
651        crate::rwkvzip::MethodSpec::File { path, .. } => Some(AssetRef::Filesystem(path.clone())),
652        crate::rwkvzip::MethodSpec::Online { .. } => None,
653    }
654}
655
656#[cfg(feature = "backend-mamba")]
657fn mamba_asset_ref(spec: &crate::mambazip::MethodSpec) -> Option<AssetRef> {
658    match spec {
659        crate::mambazip::MethodSpec::File { path, .. } => Some(AssetRef::Filesystem(path.clone())),
660        crate::mambazip::MethodSpec::Online { .. } => None,
661    }
662}
663
664pub(crate) fn encode_rate_backend_plan(plan: &RateBackendPlan) -> CanonicalBytes {
665    let mut payload = Vec::new();
666    payload.extend_from_slice(b"itrb");
667    payload.push(1);
668    encode_rate_backend_payload(plan, &mut payload);
669    frame_canonical_payload(payload)
670}
671
672fn encode_compression_backend_plan(plan: &CompressionBackendPlan) -> CanonicalBytes {
673    let mut payload = Vec::new();
674    payload.extend_from_slice(b"itcb");
675    payload.push(1);
676    encode_compression_backend_payload(plan, &mut payload);
677    frame_canonical_payload(payload)
678}
679
680fn frame_canonical_payload(payload: Vec<u8>) -> CanonicalBytes {
681    let mut framed = Vec::with_capacity(payload.len() + 10);
682    push_varint(&mut framed, payload.len() as u64);
683    framed.extend_from_slice(&payload);
684    CanonicalBytes::from(framed)
685}
686
687fn encode_rate_backend_payload(plan: &RateBackendPlan, out: &mut Vec<u8>) {
688    crate::runtime::encode_rate_backend_payload_via_kernel(plan, out)
689}
690
691fn encode_compression_backend_payload(plan: &CompressionBackendPlan, out: &mut Vec<u8>) {
692    crate::runtime::encode_compression_backend_payload_via_kernel(plan, out)
693}
694
695fn encode_particle_spec(spec: &ParticleSpec, out: &mut Vec<u8>) {
696    push_usize(out, spec.num_particles);
697    push_usize(out, spec.context_window);
698    push_usize(out, spec.unroll_steps);
699    push_usize(out, spec.num_cells);
700    push_usize(out, spec.cell_dim);
701    push_usize(out, spec.num_rules);
702    push_usize(out, spec.selector_hidden);
703    push_usize(out, spec.rule_hidden);
704    push_usize(out, spec.noise_dim);
705    push_bool(out, spec.deterministic);
706    push_bool(out, spec.enable_noise);
707    push_f64(out, spec.noise_scale);
708    push_usize(out, spec.noise_anneal_steps);
709    push_f64(out, spec.learning_rate_readout);
710    push_f64(out, spec.learning_rate_selector);
711    push_f64(out, spec.learning_rate_rule);
712    push_usize(out, spec.bptt_depth);
713    push_f64(out, spec.optimizer_momentum);
714    push_f64(out, spec.grad_clip);
715    push_f64(out, spec.state_clip);
716    push_f64(out, spec.forget_lambda);
717    push_f64(out, spec.resample_threshold);
718    push_f64(out, spec.mutate_fraction);
719    push_f64(out, spec.mutate_scale);
720    push_bool(out, spec.mutate_model_params);
721    push_usize(out, spec.diagnostics_interval);
722    push_f64(out, spec.min_prob);
723    push_varint(out, spec.seed);
724}
725
726fn push_varint(out: &mut Vec<u8>, mut value: u64) {
727    while value >= 0x80 {
728        out.push((value as u8) | 0x80);
729        value >>= 7;
730    }
731    out.push(value as u8);
732}
733
734fn push_i64(out: &mut Vec<u8>, value: i64) {
735    let zigzag = ((value << 1) ^ (value >> 63)) as u64;
736    push_varint(out, zigzag);
737}
738
739fn push_usize(out: &mut Vec<u8>, value: usize) {
740    push_varint(out, value as u64);
741}
742
743fn push_bool(out: &mut Vec<u8>, value: bool) {
744    out.push(u8::from(value));
745}
746
747fn push_f64(out: &mut Vec<u8>, value: f64) {
748    out.extend_from_slice(&value.to_bits().to_le_bytes());
749}
750
751fn push_string(out: &mut Vec<u8>, value: &str) {
752    push_varint(out, value.len() as u64);
753    out.extend_from_slice(value.as_bytes());
754}
755
756fn push_option_string(out: &mut Vec<u8>, value: Option<&str>) {
757    match value {
758        Some(value) => {
759            out.push(1);
760            push_string(out, value);
761        }
762        None => out.push(0),
763    }
764}
765
766fn push_option_f64(out: &mut Vec<u8>, value: Option<f64>) {
767    match value {
768        Some(value) => {
769            out.push(1);
770            push_f64(out, value);
771        }
772        None => out.push(0),
773    }
774}
775
776#[cfg(any(feature = "backend-mamba", feature = "backend-rwkv"))]
777fn push_asset_ref(out: &mut Vec<u8>, value: Option<&AssetRef>) {
778    match value {
779        Some(AssetRef::Filesystem(path)) => {
780            out.push(1);
781            push_string(out, &path.to_string_lossy());
782        }
783        None => out.push(0),
784    }
785}
786
787fn coder_tag(coder: CoderType) -> u8 {
788    match coder {
789        CoderType::AC => 0,
790        CoderType::RANS => 1,
791    }
792}
793
794fn framing_tag(framing: FramingMode) -> u8 {
795    match framing {
796        FramingMode::Raw => 0,
797        FramingMode::Framed => 1,
798    }
799}
800
801fn mixture_kind_tag(kind: MixtureKind) -> u8 {
802    match kind {
803        MixtureKind::Bayes => 0,
804        MixtureKind::FadingBayes => 1,
805        MixtureKind::Switching => 2,
806        MixtureKind::Convex => 3,
807        MixtureKind::Mdl => 4,
808        MixtureKind::Neural => 5,
809    }
810}
811
812fn mixture_schedule_tag(schedule: MixtureScheduleMode) -> u8 {
813    match schedule {
814        MixtureScheduleMode::Default => 0,
815        MixtureScheduleMode::Theorem => 1,
816    }
817}
818
819fn calibration_context_tag(context: CalibrationContextKind) -> u8 {
820    match context {
821        CalibrationContextKind::Global => 0,
822        CalibrationContextKind::ByteClass => 1,
823        CalibrationContextKind::Text => 2,
824        CalibrationContextKind::Repeat => 3,
825        CalibrationContextKind::TextRepeat => 4,
826    }
827}
828
829#[cfg(test)]
830mod tests {
831    use super::*;
832
833    fn assert_not_prefix(a: &CanonicalBytes, b: &CanonicalBytes) {
834        assert!(
835            !b.as_slice().starts_with(a.as_slice()),
836            "canonical encoding must be prefix-free"
837        );
838    }
839
840    fn read_varint_prefix(bytes: &[u8]) -> Result<(u64, usize), String> {
841        let mut shift: u32 = 0;
842        let mut out: u64 = 0;
843        let mut index: usize = 0;
844        loop {
845            let byte = *bytes.get(index).ok_or_else(|| {
846                "unexpected end while decoding canonical frame length".to_string()
847            })?;
848            out |= ((byte & 0x7f) as u64) << shift;
849            index = index.saturating_add(1);
850            if byte & 0x80 == 0 {
851                return Ok((out, index));
852            }
853            shift = shift.saturating_add(7);
854            if shift > 63 {
855                return Err("invalid canonical frame varint length".to_string());
856            }
857        }
858    }
859
860    fn parse_framed_canonical_payload(bytes: &[u8]) -> Result<&[u8], String> {
861        let (declared_len_u64, header_len) = read_varint_prefix(bytes)?;
862        let declared_len = usize::try_from(declared_len_u64)
863            .map_err(|_| "canonical frame length exceeds usize::MAX".to_string())?;
864        let payload_end = header_len
865            .checked_add(declared_len)
866            .ok_or_else(|| "canonical frame length overflow".to_string())?;
867        if bytes.len() < payload_end {
868            return Err("truncated canonical frame payload".to_string());
869        }
870        if bytes.len() > payload_end {
871            return Err("canonical frame has trailing bytes".to_string());
872        }
873        Ok(&bytes[header_len..payload_end])
874    }
875
876    fn sample_compression_backend_corpus() -> Vec<CompressionBackend> {
877        let mut out = Vec::<CompressionBackend>::new();
878        let leaf = crate::runtime::first_enabled_default_rate_backend_spec();
879        for descriptor in crate::runtime::COMPRESSION_BACKEND_REGISTRY {
880            if !descriptor.enabled {
881                continue;
882            }
883            match descriptor.kind {
884                crate::runtime::CompressionBackendKind::Zpaq => {
885                    out.push(CompressionBackend::zpaq("5"));
886                }
887                crate::runtime::CompressionBackendKind::RateAc => {
888                    if let Some(rate_backend) = leaf.clone() {
889                        out.push(CompressionBackend::Rate {
890                            rate_backend,
891                            coder: crate::coders::CoderType::AC,
892                            framing: crate::compression::FramingMode::Framed,
893                        });
894                    }
895                }
896                crate::runtime::CompressionBackendKind::RateRans => {
897                    if let Some(rate_backend) = leaf.clone() {
898                        out.push(CompressionBackend::Rate {
899                            rate_backend,
900                            coder: crate::coders::CoderType::RANS,
901                            framing: crate::compression::FramingMode::Raw,
902                        });
903                    }
904                }
905                #[cfg(feature = "backend-rwkv")]
906                crate::runtime::CompressionBackendKind::Rwkv7 => {
907                    let options = crate::spec::CompressionBackendShorthandOptions {
908                        default_framing: crate::compression::FramingMode::Raw,
909                        ..Default::default()
910                    };
911                    let candidate = crate::spec::parse_compression_backend_name_method(
912                        "rwkv7",
913                        Some(
914                            "cfg:hidden=64,intermediate=64,layers=1,train=sgd,lr=0.01;policy:schedule=0..100:infer",
915                        ),
916                        None,
917                        &options,
918                    )
919                    .expect("rwkv shorthand candidate");
920                    out.push(candidate);
921                }
922                #[cfg(not(feature = "backend-rwkv"))]
923                crate::runtime::CompressionBackendKind::Rwkv7 => {}
924            }
925        }
926        out
927    }
928
929    #[test]
930    fn canonical_compression_code_frame_is_self_delimiting_and_typed() {
931        let env = SpecEnvironment::default();
932        let corpus = sample_compression_backend_corpus();
933        if corpus.is_empty() {
934            return;
935        }
936        for candidate in corpus {
937            let validated =
938                validate_compression_backend_in(&candidate, &env).expect("validate candidate");
939            let bytes = validated.canonical_bytes.as_slice();
940            let payload = parse_framed_canonical_payload(bytes).expect("valid canonical frame");
941            assert!(
942                payload.len() >= 5,
943                "canonical payload must contain magic + version"
944            );
945            assert_eq!(&payload[..4], b"itcb");
946            assert_eq!(payload[4], 1_u8);
947        }
948    }
949
950    #[test]
951    fn canonical_compression_code_frame_rejects_trailing_bytes_in_contract_parser() {
952        let env = SpecEnvironment::default();
953        let Some(candidate) = sample_compression_backend_corpus().into_iter().next() else {
954            return;
955        };
956        let validated = validate_compression_backend_in(&candidate, &env).expect("validate");
957        let bytes = validated.canonical_bytes.as_slice().to_vec();
958        parse_framed_canonical_payload(&bytes).expect("valid canonical frame");
959
960        let mut extended = bytes.clone();
961        extended.extend_from_slice(&[0x00, 0x01]);
962        let err = parse_framed_canonical_payload(&extended)
963            .expect_err("trailing bytes must be rejected by canonical frame parser");
964        assert!(err.contains("trailing bytes"), "{err}");
965    }
966
967    #[test]
968    fn canonical_rate_plan_bytes_are_prefix_free_for_sample_corpus() {
969        let env = SpecEnvironment::default();
970        let samples: Vec<_> = crate::runtime::RATE_BACKEND_REGISTRY
971            .iter()
972            .filter(|descriptor| descriptor.enabled)
973            .filter_map(|descriptor| crate::runtime::default_rate_backend_spec(descriptor.kind))
974            .take(4)
975            .collect();
976        if samples.len() < 2 {
977            return;
978        }
979        let encodings: Vec<_> = samples
980            .iter()
981            .map(|backend| {
982                validate_rate_backend_in(backend, &env)
983                    .unwrap()
984                    .canonical_bytes
985                    .clone()
986            })
987            .collect();
988        for (i, left) in encodings.iter().enumerate() {
989            for (j, right) in encodings.iter().enumerate() {
990                if i != j {
991                    assert_not_prefix(left, right);
992                }
993            }
994        }
995    }
996
997    #[test]
998    fn compiled_rate_backend_clone_is_o1_arc_backed() {
999        let Some(default_backend) = crate::runtime::first_enabled_default_rate_backend_spec()
1000        else {
1001            return;
1002        };
1003        let compiled = validate_rate_backend_in(&default_backend, &SpecEnvironment::default())
1004            .unwrap()
1005            .compile()
1006            .unwrap();
1007        let cloned = compiled.clone();
1008        assert!(Arc::ptr_eq(&compiled.plan, &cloned.plan));
1009        assert!(Arc::ptr_eq(
1010            &compiled.canonical_spec,
1011            &cloned.canonical_spec
1012        ));
1013    }
1014
1015    #[cfg(feature = "backend-ctw")]
1016    #[test]
1017    fn compiled_ctw_preserves_family_identity() {
1018        let compiled =
1019            validate_rate_backend_in(&RateBackend::Ctw { depth: 7 }, &SpecEnvironment::default())
1020                .unwrap()
1021                .compile()
1022                .unwrap();
1023        assert!(matches!(
1024            compiled.canonical_spec(),
1025            RateBackend::Ctw { depth: 7 }
1026        ));
1027    }
1028
1029    #[cfg(feature = "backend-mixture")]
1030    #[cfg(feature = "backend-zpaq")]
1031    #[test]
1032    fn compiled_capabilities_track_nested_zpaq_components() {
1033        let backend = RateBackend::Mixture {
1034            spec: Arc::new(MixtureSpec::new(
1035                MixtureKind::Bayes,
1036                vec![
1037                    MixtureExpertSpec {
1038                        name: Some("ctw".to_string()),
1039                        log_prior: 0.0,
1040                        backend: RateBackend::Ctw { depth: 6 },
1041                    },
1042                    MixtureExpertSpec {
1043                        name: Some("zpaq".to_string()),
1044                        log_prior: -0.1,
1045                        backend: RateBackend::Zpaq {
1046                            method: crate::api::ZpaqMethodSpec::literal("1"),
1047                        },
1048                    },
1049                ],
1050            )),
1051        };
1052        let compiled = validate_rate_backend_in(&backend, &SpecEnvironment::default())
1053            .unwrap()
1054            .compile()
1055            .unwrap();
1056        assert!(compiled.contains_zpaq());
1057    }
1058}
infotheory/spec/core.rs

infotheory/spec/
core.rs
