viterbi 0.0.1

// Author: BolivarTech
// Version: 1.0.0
// Date: 2026-07-02
//! Hard-decision block Viterbi decoder (ACS + traceback).
//!
//! Orchestrates ACS + traceback and enforces the memory/length limits (R3, R11, R18).
pub(crate) mod acs;
pub(crate) mod traceback;

use crate::metric::{BranchMetric, HardHamming};
use crate::packing::{pack_into, unpack_into, CodedBlock, DecodedBlock};
use crate::params::{CodeParams, ParamError};
use crate::trellis::Trellis;
use crate::{K, M, MAX_SUPPORTED_INFO_BITS, N};
use acs::{acs_stage, SENTINEL};

/// Errors from [`ViterbiDecoder::new`].
#[derive(Debug, Clone, PartialEq, Eq)]
pub enum ConfigError {
    /// A [`ParamError`] surfaced while validating the code parameters or the state count `S`.
    Param(ParamError),
    /// `max_info_bits` was zero (a decoder must admit at least one info bit).
    MaxBlockZero,
    /// `max_info_bits` exceeded the hard cap [`MAX_SUPPORTED_INFO_BITS`] (R18).
    MaxBlockTooLarge {
        /// The hard cap.
        cap: usize,
        /// The requested value.
        got: usize,
    },
    /// A fallible scratch preallocation failed (out of memory) — reported, never a panic (R11).
    AllocationFailed,
}
/// Errors from decoding (boundary/allocation checks only — the ACS correction loop is infallible, R11).
#[derive(Debug, Clone, PartialEq, Eq)]
pub enum DecodeError {
    /// The implied payload length exceeded the decoder's configured `max_info_bits` (R18).
    InputTooLong {
        /// The configured cap.
        max_bits: usize,
        /// The rejected payload length in bits.
        got_bits: usize,
    },
    /// The coded input length was inconsistent (not a multiple of `N`, outside the byte buffer,
    /// too short for the `M` tail, or `received.len()` inconsistent with `nbits`).
    LengthMismatch,
    /// The recovered-bytes output allocation failed (out of memory) — reported, never a panic (R11).
    AllocationFailed,
}
impl core::fmt::Display for ConfigError {
    fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
        write!(f, "{self:?}")
    }
}
impl core::fmt::Display for DecodeError {
    fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
        write!(f, "{self:?}")
    }
}
impl std::error::Error for ConfigError {
    fn source(&self) -> Option<&(dyn std::error::Error + 'static)> {
        match self {
            ConfigError::Param(e) => Some(e),
            _ => None,
        }
    }
}
impl std::error::Error for DecodeError {}

/// Hard-decision block Viterbi decoder over `S` states, generic over the branch metric `Me` (R7, R12).
///
/// All scratch (survivor, traceback, sample buffers) is preallocated in [`new`](Self::new), so the
/// ACS correction loop allocates nothing and cannot OOM-panic (R11, R18).
pub struct ViterbiDecoder<const S: usize, Me: BranchMetric> {
    trellis: Trellis<S>,
    max_info_bits: usize,
    survivors: Vec<u64>, // preallocated (max_info_bits + M) · ⌈S/64⌉ decision words
    inputs: Vec<u8>,     // preallocated traceback scratch (max_info_bits + M)
    samples: Vec<Me::Sample>, // preallocated unpack scratch ((max_info_bits + M) * N)
    wps: usize,          // decision words per stage = ⌈S/64⌉ (cached; computed once in new)
    terminal_metric: u32, // metric of state 0 after the last decode (diagnostic)
    _metric: core::marker::PhantomData<Me>,
}

/// First-class CCSDS hard-decision profile.
pub type CcsdsViterbiDecoder = ViterbiDecoder<64, HardHamming>;

impl<const S: usize, Me: BranchMetric> ViterbiDecoder<S, Me> {
    /// Build the decoder; validate params and `S`, preallocate all scratch via `try_reserve` (R11, R18).
    ///
    /// Never panics: `S != 2^(k-1)` is returned as `ConfigError::Param(StateCountMismatch)`.
    ///
    /// # Errors
    /// - [`ConfigError::Param`] — invalid/unsupported/catastrophic params, or `S != 2^(k-1)`.
    /// - [`ConfigError::MaxBlockZero`] / [`ConfigError::MaxBlockTooLarge`] — `max_info_bits` out of range.
    /// - [`ConfigError::AllocationFailed`] — a scratch preallocation failed.
    pub fn new(params: CodeParams, max_info_bits: usize) -> Result<Self, ConfigError> {
        params.validate().map_err(ConfigError::Param)?; // defense-in-depth: re-check the invariant (R19)
        if params.k != K {
            return Err(ConfigError::Param(ParamError::KUnsupported { k: params.k }));
        }
        if params.generators.len() != N {
            return Err(ConfigError::Param(ParamError::RateUnsupported {
                n: params.generators.len(),
            }));
        }
        if max_info_bits == 0 {
            return Err(ConfigError::MaxBlockZero);
        }
        if max_info_bits > MAX_SUPPORTED_INFO_BITS {
            return Err(ConfigError::MaxBlockTooLarge {
                cap: MAX_SUPPORTED_INFO_BITS,
                got: max_info_bits,
            });
        }
        let trellis = Trellis::from_params(&params).map_err(ConfigError::Param)?; // validates S; no panic
        let stages = max_info_bits + M as usize;
        let wps = S.div_ceil(64); // decision words per stage (scales with S — R12)
        let mut survivors = Vec::new();
        let mut inputs = Vec::new();
        let mut samples = Vec::new();
        // saturating_mul on the size products for defense-in-depth consistency with encode_bits
        // (these are already bounded: stages ≤ max_info_bits + M ≤ 1M+6, wps ≤ 4, N = 2 → no overflow).
        survivors
            .try_reserve_exact(stages.saturating_mul(wps))
            .map_err(|_| ConfigError::AllocationFailed)?;
        inputs
            .try_reserve_exact(stages)
            .map_err(|_| ConfigError::AllocationFailed)?;
        samples
            .try_reserve_exact(stages.saturating_mul(N))
            .map_err(|_| ConfigError::AllocationFailed)?;
        // `survivors` is indexed by stage, so it is resized (zero-filled) up front. `inputs` and
        // `samples` are reused push-buffers: kept at len 0 here and refilled within their reserved
        // capacity each decode (traceback / unpack), so they are intentionally not resized.
        survivors.resize(stages.saturating_mul(wps), 0);
        Ok(Self {
            trellis,
            max_info_bits,
            survivors,
            inputs,
            samples,
            wps,
            terminal_metric: 0,
            _metric: core::marker::PhantomData,
        })
    }

    /// Metric of the terminal state 0 after the last `decode` — a soft-confidence diagnostic
    /// (higher ⇒ the received block was farther from any valid terminated codeword). Never a pass/fail.
    /// Reflects the last `decode` that reached traceback; **unspecified** after a call that returned `Err`.
    #[must_use]
    pub fn last_terminal_metric(&self) -> u32 {
        self.terminal_metric
    }

    /// Decode `(nbits + M)` stages of `received` (N samples per stage) into the payload bits.
    ///
    /// The ACS correction loop allocates nothing (preallocated survivor buffer + stack arrays);
    /// only the recovered-bytes output uses a fallible `try_reserve` (R11).
    ///
    /// # Errors
    /// - [`DecodeError::InputTooLong`] — `nbits` exceeds the configured `max_info_bits`.
    /// - [`DecodeError::LengthMismatch`] — `received.len()` is inconsistent with `nbits`.
    /// - [`DecodeError::AllocationFailed`] — the recovered-bytes output allocation failed.
    pub fn decode(
        &mut self,
        received: &[Me::Sample],
        nbits: usize,
    ) -> Result<DecodedBlock, DecodeError> {
        if nbits > self.max_info_bits {
            return Err(DecodeError::InputTooLong {
                max_bits: self.max_info_bits,
                got_bits: nbits,
            });
        }
        let stages = nbits + M as usize;
        if received.len() != stages * N {
            return Err(DecodeError::LengthMismatch);
        }
        let wps = self.wps; // cached ⌈S/64⌉ (computed once in `new`)
        let mut cur = [SENTINEL; S];
        cur[0] = 0;
        let mut next = [0u32; S];
        for stage in 0..stages {
            let slice = &received[stage * N..stage * N + N];
            acs_stage::<S, Me>(
                &cur,
                &mut next,
                &mut self.survivors[stage * wps..(stage + 1) * wps],
                &self.trellis,
                slice,
            );
            core::mem::swap(&mut cur, &mut next);
        }
        self.terminal_metric = cur[0]; // diagnostic: metric of the assumed terminal state 0
                                       // traceback into the preallocated scratch (disjoint field borrows; no allocation)
        traceback::traceback::<S>(
            &self.trellis,
            &self.survivors[..stages * wps],
            stages,
            &mut self.inputs,
        );
        // Pack the recovered info bits via the packing module (sole owner of the MSB-first
        // convention, R16). Allocation stays fallible: reserve exactly the needed capacity, then let
        // `pack_into` clear+resize into it — since the capacity already matches, no reallocation occurs.
        let mut bytes: Vec<u8> = Vec::new();
        bytes
            .try_reserve_exact(nbits.div_ceil(8))
            .map_err(|_| DecodeError::AllocationFailed)?;
        pack_into(&self.inputs[..nbits], nbits, &mut bytes);
        Ok(DecodedBlock { bytes, nbits })
    }
}

impl<const S: usize> ViterbiDecoder<S, HardHamming> {
    /// Hard-decision block decode — reads (by reference) the `CodedBlock` (symmetric with `encode`).
    ///
    /// **Validates `coded.nbits` fully before unpacking** (R16, no OOB): must be a multiple of `N`,
    /// within the byte buffer, and long enough for the `M` tail.
    ///
    /// # Errors
    /// - [`DecodeError::LengthMismatch`] — `nbits` not a multiple of `N`, outside the byte buffer,
    ///   or too short for the `M` tail.
    /// - [`DecodeError::InputTooLong`] — the implied payload exceeds `max_info_bits`.
    /// - [`DecodeError::AllocationFailed`] — the recovered-bytes output allocation failed.
    pub fn decode_block(&mut self, coded: &CodedBlock) -> Result<DecodedBlock, DecodeError> {
        if coded.nbits % N != 0 {
            return Err(DecodeError::LengthMismatch);
        }
        // `div_ceil` (no multiplication) is overflow-proof and provably correct: unpacking `coded.nbits`
        // bits needs `⌈nbits/8⌉` bytes, so reject when that exceeds the buffer. Avoids any saturating-mul
        // edge case on an attacker-controlled `bytes.len()`.
        if coded.nbits.div_ceil(8) > coded.bytes.len() {
            return Err(DecodeError::LengthMismatch);
        }
        let stages = coded.nbits / N;
        if stages < M as usize {
            return Err(DecodeError::LengthMismatch);
        }
        let nbits = stages - M as usize;
        if nbits > self.max_info_bits {
            return Err(DecodeError::InputTooLong {
                max_bits: self.max_info_bits,
                got_bits: nbits,
            });
        }
        // Unpack into preallocated scratch via the packing module (sole owner, R16). The buffer is
        // moved out (borrow checker) and restored below; safe because `decode` never panics, so the
        // restore always runs. Even in a hypothetical unwind the empty scratch **self-heals** on the
        // next call (`unpack_into` refills it) — no correctness impact. No allocation (preallocated).
        let mut samples = core::mem::take(&mut self.samples);
        samples.clear();
        // The buffer is preallocated in `new` to `(max_info_bits + M)·N ≥ coded.nbits` (guaranteed by the
        // checks above), so this reserve is a no-op in practice; making it fallible removes any OOM-panic
        // path (R11) — an allocation failure returns `Err`, never aborts.
        if samples.try_reserve(coded.nbits).is_err() {
            self.samples = samples; // restore before bailing
            return Err(DecodeError::AllocationFailed);
        }
        unpack_into(&coded.bytes, coded.nbits, &mut samples);
        let result = self.decode(&samples, nbits);
        self.samples = samples; // restore the scratch buffer
        result
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::encoder::ViterbiEncoder;
    use crate::params::CodeParams;
    #[test]
    fn clean_round_trip_recovers_payload() {
        let enc = ViterbiEncoder::new(CodeParams::ccsds_r1_2()).unwrap();
        let mut dec = CcsdsViterbiDecoder::new(CodeParams::ccsds_r1_2(), 4096).unwrap();
        let raw = b"Viterbi works!";
        let coded = enc.encode(raw).unwrap();
        let out = dec.decode_block(&coded).unwrap();
        assert_eq!(&out.bytes, raw);
        assert_eq!(out.nbits, raw.len() * 8);
    }
    #[test]
    fn new_rejects_zero_and_oversized_max_bits() {
        assert_eq!(
            CcsdsViterbiDecoder::new(CodeParams::ccsds_r1_2(), 0)
                .err()
                .unwrap(),
            ConfigError::MaxBlockZero
        );
        assert!(matches!(
            CcsdsViterbiDecoder::new(CodeParams::ccsds_r1_2(), 2_000_000),
            Err(ConfigError::MaxBlockTooLarge { .. })
        ));
    }
    #[test]
    fn decode_rejects_input_longer_than_max() {
        let enc = ViterbiEncoder::new(CodeParams::ccsds_r1_2()).unwrap();
        let mut dec = CcsdsViterbiDecoder::new(CodeParams::ccsds_r1_2(), 8).unwrap();
        let coded = enc.encode(&[0u8; 4]).unwrap(); // 32 info bits > 8 cap
        assert!(matches!(
            dec.decode_block(&coded),
            Err(DecodeError::InputTooLong { .. })
        ));
    }
    #[test]
    fn decode_block_rejects_malformed_lengths() {
        use crate::packing::CodedBlock;
        let mut dec = CcsdsViterbiDecoder::new(CodeParams::ccsds_r1_2(), 4096).unwrap();
        // nbits not a multiple of N=2
        assert!(matches!(
            dec.decode_block(&CodedBlock {
                bytes: vec![0; 4],
                nbits: 15
            }),
            Err(DecodeError::LengthMismatch)
        ));
        // nbits exceeds the byte buffer (100 > 2*8)
        assert!(matches!(
            dec.decode_block(&CodedBlock {
                bytes: vec![0; 2],
                nbits: 100
            }),
            Err(DecodeError::LengthMismatch)
        ));
        // too short for the M=6 tail (stages = 8/2 = 4 < 6)
        assert!(matches!(
            dec.decode_block(&CodedBlock {
                bytes: vec![0; 1],
                nbits: 8
            }),
            Err(DecodeError::LengthMismatch)
        ));
    }
    #[test]
    fn new_rejects_unsupported_k_and_rate() {
        // rate 1/3 (3 generators) at K=7 → RateUnsupported
        let r13 = CodeParams::new(7, vec![0o133, 0o171, 0o145], vec![false, false, false]).unwrap();
        assert!(matches!(
            CcsdsViterbiDecoder::new(r13, 64),
            Err(ConfigError::Param(ParamError::RateUnsupported { n: 3 }))
        ));
        // K=3 (unsupported by the K=7 engine) → KUnsupported (checked before the S/trellis validation)
        let k3 = CodeParams::new(3, vec![0b101, 0b111], vec![false, false]).unwrap();
        assert!(matches!(
            CcsdsViterbiDecoder::new(k3, 64),
            Err(ConfigError::Param(ParamError::KUnsupported { k: 3 }))
        ));
    }
    #[test]
    fn new_rejects_catastrophic_params_defensively() {
        // Defense-in-depth (release-active): even an in-crate struct literal bypassing `CodeParams::new`
        // (fields are pub(crate)) with catastrophic generators [0o12, 0o12] (gcd = 0o12, not a monomial)
        // is rejected by `new`'s `validate()` re-check — the engine never trusts params blindly.
        let bad = CodeParams {
            k: 7,
            generators: vec![0o12, 0o12],
            invert_outputs: vec![false, false],
        };
        assert!(matches!(
            CcsdsViterbiDecoder::new(bad, 64),
            Err(ConfigError::Param(ParamError::Catastrophic))
        ));
    }
}