viterbi 0.2.0

// Author: BolivarTech
// Version: 1.0.0
// Date: 2026-07-05
//! Puncturing support: the puncture matrix as **runtime data** (RP1) plus its
//! structural validation.
//!
//! This module hosts the whole puncturing layer. Task 7 lands the data model —
//! [`PunctureMatrix`] (the keep-pattern), [`PunctureOrder`], [`PuncturedRate`],
//! and the [`PunctureError`] domain type — together with **structural**
//! validation only. The transmit-side `Puncturer` and receive-side
//! `DePuncturer` (which add the catastrophic-puncture guard, the `NMismatch`
//! guard, and the alignment checks) attach in later increments; their error
//! variants are declared here so the enum is stable.
//!
//! ## Design notes
//! - The keep-pattern is **data**, never a hardcoded constant in the core
//!   encoder/decoder (RP1). A target standard with a different convention is
//!   expressed by constructing a different [`PunctureMatrix`].
//! - Puncture order/labeling differs across DVB-S / 802.11 / CCSDS; it is
//!   captured as data via [`PunctureOrder`].

use crate::metric::BranchMetric;
use crate::packing::{get_bit, set_bit, CodedBlock};
use crate::params::CodeParams;
use crate::trellis::next_state;
use crate::MAX_SUPPORTED_INFO_BITS;

/// Maximum supported puncture period (columns of the keep-pattern).
///
/// Bounds construction cost and — for the later catastrophic check — the size of
/// the `(state, phase)` product graph, closing a construction-time denial of
/// service from an adversarial huge period.
pub const MAX_PERIOD: usize = 128;

/// Errors from constructing/validating the puncturing types.
///
/// Only the structural variants ([`EmptyPattern`](PunctureError::EmptyPattern),
/// [`RaggedRows`](PunctureError::RaggedRows),
/// [`PeriodTooLarge`](PunctureError::PeriodTooLarge),
/// [`AllVoidPeriod`](PunctureError::AllVoidPeriod),
/// [`UnsupportedRate`](PunctureError::UnsupportedRate)) are produced by
/// [`PunctureMatrix`]. The remaining variants
/// ([`Catastrophic`](PunctureError::Catastrophic),
/// [`RowCountMismatch`](PunctureError::RowCountMismatch),
/// [`NMismatch`](PunctureError::NMismatch),
/// [`MisalignedInput`](PunctureError::MisalignedInput),
/// [`AllocationFailed`](PunctureError::AllocationFailed)) are declared here for a
/// stable enum but are raised by the `Puncturer`/`DePuncturer` layers added in
/// later tasks.
#[derive(Debug, Clone, PartialEq, Eq)]
pub enum PunctureError {
    /// The keep-pattern had no rows, or at least one row was zero-length.
    EmptyPattern,
    /// The keep-pattern's row count did not match the mother code's `n`.
    ///
    /// Reserved: the `Puncturer`/`DePuncturer` layer reports an `n` mismatch as [`NMismatch`] instead;
    /// this variant is kept for a future increment and is not currently produced.
    ///
    /// [`NMismatch`]: PunctureError::NMismatch
    RowCountMismatch {
        /// The keep-pattern's row count.
        rows: usize,
        /// The mother code's number of outputs `n`.
        n: usize,
    },
    /// The keep-pattern's rows were not all the same length.
    RaggedRows,
    /// The period (row length) exceeded [`MAX_PERIOD`].
    PeriodTooLarge {
        /// The offending period.
        period: usize,
        /// The cap it was checked against ([`MAX_PERIOD`]).
        cap: usize,
    },
    /// A whole column of the period was punctured (no kept bit in any row).
    ///
    /// The degenerate "transmit nothing this step" pattern is rejected: every
    /// column must keep at least one bit.
    AllVoidPeriod,
    /// The requested `(n, rate)` has no supported preset (only `n = 2`).
    UnsupportedRate,
    /// The `(mother code + pattern)` combination is catastrophic.
    ///
    /// Raised only by `Puncturer::new`/`DePuncturer::new` (later tasks), which
    /// carry the generator polynomials; [`PunctureMatrix::new`] structurally
    /// cannot detect it.
    Catastrophic,
    /// The keep-pattern's `n` did not match the mother code's `n`.
    ///
    /// Raised by the `Puncturer`/`DePuncturer` layer.
    NMismatch {
        /// The keep-pattern's `n`.
        matrix_n: usize,
        /// The mother code's `n`.
        params_n: usize,
    },
    /// The de-punctured/punctured stream length did not match the expected value.
    ///
    /// Raised by the `Puncturer`/`DePuncturer` layer.
    MisalignedInput {
        /// The length that was supplied.
        got: usize,
        /// The length that was expected for the given `nbits`.
        expected: usize,
    },
    /// A fallible allocation (`try_reserve`) failed.
    ///
    /// Raised by the `DePuncturer` layer.
    AllocationFailed,
    /// The payload `nbits` exceeded the supported cap [`MAX_SUPPORTED_INFO_BITS`] (R18).
    ///
    /// Raised by [`Puncturer::puncture`]/[`DePuncturer::depuncture`], symmetric with the
    /// encoder/decoder payload cap. Distinct from [`MisalignedInput`](PunctureError::MisalignedInput),
    /// which reports a length *inconsistency* for an in-range `nbits`.
    PayloadTooLarge {
        /// The requested payload bit count.
        got: usize,
        /// The hard cap ([`MAX_SUPPORTED_INFO_BITS`]).
        cap: usize,
    },
}
impl core::fmt::Display for PunctureError {
    fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
        write!(f, "{self:?}")
    }
}
impl std::error::Error for PunctureError {}

/// Order/labeling of the punctured output bits (data, not hardcoded — RP1).
///
/// Captures the DVB-S vs 802.11 vs CCSDS labeling divergence:
/// - [`ByColumn`](PunctureOrder::ByColumn): all `C1` of the period, then all
///   `C2`, …, then all `Cn` (DVB/802.11 style).
/// - [`Interleaved`](PunctureOrder::Interleaved): per time step emit
///   `C1(t), C2(t), …, Cn(t)` then advance `t`. For `n = 2` this is exactly the
///   CCSDS 131.0-B-3 `XY` order `C1(1), C2(1), C1(2), C2(2), …`.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum PunctureOrder {
    /// Emit every kept `C1` of the period, then every kept `C2`, … (DVB/802.11).
    ByColumn,
    /// Emit `C1(t), C2(t), …, Cn(t)` per time step (CCSDS 131.0-B-3 for `n = 2`).
    Interleaved,
}

/// Preconfigured punctured code rates from the standard Yasuda/DVB-S/IESS-308 set
/// that CCSDS references.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum PuncturedRate {
    /// Rate 2/3 (period 2).
    R2_3,
    /// Rate 3/4 (period 3).
    R3_4,
    /// Rate 5/6 (period 5).
    R5_6,
    /// Rate 7/8 (period 7).
    R7_8,
}

/// A puncture keep-pattern: `n` rows × `period` columns of keep (`true`) / puncture
/// (`false`) flags, plus its labeling [`PunctureOrder`].
///
/// The pattern is validated and its metadata cached at construction so the hot
/// path only reads it. Construct once and reuse across frames.
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct PunctureMatrix {
    /// Keep-pattern rows `[C1; C2; …; Cn]`; `true` = transmit, `false` = puncture.
    keep: Vec<Vec<bool>>,
    /// Number of rows (= mother code output count `n`).
    n: usize,
    /// Number of columns (the puncture period).
    period: usize,
    /// Count of `true` entries over one period.
    kept_per_period: usize,
    /// Output labeling.
    order: PunctureOrder,
}

impl PunctureMatrix {
    /// Build a puncture matrix from a keep-pattern, validating it **structurally**.
    ///
    /// # Arguments
    /// - `keep`: `n` rows `[C1; C2; …; Cn]`, each of length `period`; `true`
    ///   transmits the corresponding coded bit, `false` punctures it.
    /// - `order`: the output labeling ([`PunctureOrder`]).
    ///
    /// # Validation order
    /// 1. no rows, or any zero-length row → [`PunctureError::EmptyPattern`]
    ///    (checked first, before any column indexing);
    /// 2. rows not all the same length → [`PunctureError::RaggedRows`];
    /// 3. `period` (the row length) outside `1..=`[`MAX_PERIOD`] →
    ///    [`PunctureError::PeriodTooLarge`];
    /// 4. any column with no kept bit → [`PunctureError::AllVoidPeriod`].
    ///
    /// # Errors
    /// Returns the [`PunctureError`] above on a structurally invalid pattern.
    /// This constructor **never** returns [`PunctureError::Catastrophic`]:
    /// catastrophicity is a property of `(mother code + pattern)` and needs the
    /// generator polynomials, which this constructor does not have — that check
    /// is performed by `Puncturer::new`/`DePuncturer::new` (later tasks).
    ///
    /// # Examples
    /// ```
    /// use viterbi::{PunctureMatrix, PunctureOrder};
    ///
    /// let m = PunctureMatrix::new(
    ///     vec![vec![true, true], vec![true, false]],
    ///     PunctureOrder::ByColumn,
    /// )
    /// .unwrap();
    /// assert_eq!(m.n(), 2);
    /// assert_eq!(m.period(), 2);
    /// assert_eq!(m.kept_per_period(), 3);
    /// ```
    pub fn new(keep: Vec<Vec<bool>>, order: PunctureOrder) -> Result<Self, PunctureError> {
        // (1) EmptyPattern — no rows, or any zero-length row. Checked first so no
        // later step indexes into an empty row.
        if keep.is_empty() || keep.iter().any(Vec::is_empty) {
            return Err(PunctureError::EmptyPattern);
        }
        // (1b) RowCountMismatch — cap the output count `n` (row count). The engine supports only
        // n ∈ {2,3}, and the catastrophic-check keepmask uses `1u32 << (n-1-j)`, which would
        // shift-overflow for n > 32. Rejecting n > 3 here closes that path structurally (defense in
        // depth), independent of the `Puncturer`/`DePuncturer::new` `NMismatch` guard against params.
        if keep.len() > 3 {
            return Err(PunctureError::RowCountMismatch {
                rows: keep.len(),
                n: 3,
            });
        }
        // (2) RaggedRows — all rows must share one period.
        let period = keep[0].len();
        if keep.iter().any(|row| row.len() != period) {
            return Err(PunctureError::RaggedRows);
        }
        // (3) PeriodTooLarge — bound construction/graph cost. period >= 1 here.
        if period > MAX_PERIOD {
            return Err(PunctureError::PeriodTooLarge {
                period,
                cap: MAX_PERIOD,
            });
        }
        // (4) AllVoidPeriod — every column must keep at least one bit.
        for col in 0..period {
            if keep.iter().all(|row| !row[col]) {
                return Err(PunctureError::AllVoidPeriod);
            }
        }
        let n = keep.len();
        let kept_per_period = keep.iter().flatten().filter(|&&b| b).count();
        Ok(Self {
            keep,
            n,
            period,
            kept_per_period,
            order,
        })
    }

    /// Build a preconfigured CCSDS/Yasuda punctured matrix for the given rate.
    ///
    /// The matrix uses [`PunctureOrder::Interleaved`] — the CCSDS 131.0-B-3 `XY`
    /// order for `n = 2`.
    ///
    /// # Arguments
    /// - `n`: the mother code's output count. Only `n = 2` (rate-1/2 mother) has
    ///   a standard punctured family here; there is no widely standardized CCSDS
    ///   rate-1/3 punctured set.
    /// - `rate`: the target [`PuncturedRate`].
    ///
    /// # Errors
    /// Returns [`PunctureError::UnsupportedRate`] for any `n != 2`.
    ///
    /// # Examples
    /// ```
    /// use viterbi::{PunctureMatrix, PunctureOrder, PuncturedRate};
    ///
    /// let m = PunctureMatrix::ccsds_rate(2, PuncturedRate::R3_4).unwrap();
    /// assert_eq!(m.code_rate(), (3, 4));
    /// assert_eq!(m.order(), PunctureOrder::Interleaved);
    /// ```
    pub fn ccsds_rate(n: usize, rate: PuncturedRate) -> Result<Self, PunctureError> {
        if n != 2 {
            return Err(PunctureError::UnsupportedRate);
        }
        Self::new(yasuda_pattern(rate), PunctureOrder::Interleaved)
    }

    /// Number of rows of the keep-pattern (the mother code output count `n`).
    #[must_use]
    pub fn n(&self) -> usize {
        self.n
    }

    /// Number of columns of the keep-pattern (the puncture period).
    #[must_use]
    pub fn period(&self) -> usize {
        self.period
    }

    /// Count of kept (`true`) bits over one period.
    #[must_use]
    pub fn kept_per_period(&self) -> usize {
        self.kept_per_period
    }

    /// The output labeling ([`PunctureOrder`]).
    ///
    /// Exposed so a TX/RX integration can assert both sides agree
    /// (`assert_eq!(tx.order(), rx.order())`), guarding the order-mismatch interop
    /// hazard.
    #[must_use]
    pub fn order(&self) -> PunctureOrder {
        self.order
    }

    /// The punctured code rate as `(period, kept_per_period)`.
    ///
    /// One input bit is consumed per column (per trellis stage), so `period`
    /// input bits map to `kept_per_period` transmitted bits — e.g. `(2, 3)` for
    /// rate 2/3.
    #[must_use]
    pub fn code_rate(&self) -> (usize, usize) {
        (self.period, self.kept_per_period)
    }
}

/// The Yasuda/DVB-S/IESS-308 keep-matrices over the rate-1/2 mother code that
/// CCSDS references. Rows are `[C1; C2]`; `true` = transmit, `false` = puncture.
fn yasuda_pattern(rate: PuncturedRate) -> Vec<Vec<bool>> {
    // Encoded as 1/0 rows for readability, then mapped to booleans.
    let rows: &[&[u8]] = match rate {
        PuncturedRate::R2_3 => &[&[1, 1], &[1, 0]],
        PuncturedRate::R3_4 => &[&[1, 0, 1], &[1, 1, 0]],
        PuncturedRate::R5_6 => &[&[1, 0, 1, 0, 1], &[1, 1, 0, 1, 0]],
        PuncturedRate::R7_8 => &[&[1, 0, 0, 0, 1, 0, 1], &[1, 1, 1, 1, 0, 1, 0]],
    };
    rows.iter()
        .map(|row| row.iter().map(|&b| b == 1).collect())
        .collect()
}

/// The **raw, un-inverted** `n`-bit output symbol for `(state, input)` (first generator / C1 in the
/// high bit, row `j` at bit `n-1-j`).
///
/// Unlike [`crate::trellis::output_symbol`], the per-generator inversion is **deliberately not**
/// applied: catastrophicity is a linear property of the generator matrix over GF(2), and output
/// inversion is only an affine transmission convention that does not change it. Using the linear
/// (un-inverted) output is also what makes the trivial zero-state / zero-input branch an all-zero
/// self-cycle — the exact branch [`punctured_is_catastrophic`] excludes.
fn raw_output_symbol(state: usize, input: u8, params: &CodeParams) -> u32 {
    let m = (params.k() - 1) as u32;
    let g = ((input as u32) << m) | state as u32;
    let n = params.generators().len();
    let mut sym = 0u32;
    for (i, &gen) in params.generators().iter().enumerate() {
        let bit = (g & gen).count_ones() & 1;
        sym |= bit << (n - 1 - i); // first generator → high bit
    }
    sym
}

/// Popcount of the output symbol restricted to the bits kept at `phase` (`keep_mask[phase]`).
///
/// `keep_mask[phase]` sets bit `n-1-j` iff row `j` is kept at column `phase`, aligning with the
/// high-bit-first packing of [`raw_output_symbol`]. A weight of `0` means the puncture pattern
/// discards every coded bit this branch would emit at that phase.
fn output_weight_after_puncture(symbol: u32, keep_mask: u32) -> u32 {
    (symbol & keep_mask).count_ones()
}

/// Decide whether the `(mother code + puncture pattern)` combination is **catastrophic**.
///
/// This is **distinct** from [`crate::params`]'s `gcd_gf2`/`is_catastrophic`, which check the
/// *unpunctured* mother code at [`CodeParams`] construction. Here the property is evaluated on the
/// code as seen *through the puncture period* via an exhaustive zero-output-weight cycle search on
/// the `(state, phase)` **product graph** (spec §4.4):
///
/// - Nodes: `(state, phase)` for `state ∈ [0, 2^(k-1))`, `phase ∈ [0, period)` —
///   `S · period ≤ 256 · MAX_PERIOD` at the worst supported case.
/// - For each `input ∈ {0, 1}`, add the edge `(s, p) → (next, (p+1) mod period)` **iff** the branch's
///   punctured output weight at phase `p` is `0`, **excluding** the trivial zero-state / zero-input
///   self-branch (`s == 0 && input == 0`), which is the all-zero path of any code.
/// - The code is catastrophic **iff** that zero-weight sub-graph contains a cycle: a cycle is a
///   nonzero-input sequence yielding bounded output ⟺ a finite-weight output for an infinite-weight
///   input.
///
/// Cycle detection is an **iterative** three-colour DFS with an explicit heap stack (never recursion,
/// so the `≤ 256 · MAX_PERIOD`-node graph cannot overflow the call stack); a back-edge to a node still
/// on the DFS stack is a cycle. Complexity is `O(S · period)` in nodes and edges — a **one-time**
/// construction cost, never touched in the decode hot path.
///
/// Both `Puncturer::new` and (later) `DePuncturer::new` call this single fn, so the transmit and
/// receive sides can never diverge on the catastrophicity verdict.
fn punctured_is_catastrophic(
    matrix: &PunctureMatrix,
    params: &CodeParams,
) -> Result<bool, PunctureError> {
    let period = matrix.period;
    let n = matrix.n;
    let k = params.k();
    let states = 1usize << (k - 1);

    // keep_mask[phase]: bit `n-1-j` set iff row `j` is kept at that column (high-bit-first).
    let masks: Vec<u32> = (0..period)
        .map(|phase| {
            let mut mask = 0u32;
            for (j, row) in matrix.keep.iter().enumerate() {
                if row[phase] {
                    mask |= 1u32 << (n - 1 - j);
                }
            }
            mask
        })
        .collect();

    // Adjacency of the zero-punctured-weight product sub-graph (≤ 2 out-edges per node).
    // Node count is provably bounded: states ≤ 2^(9-1) = 256 and period ≤ MAX_PERIOD = 128, so
    // `states * period ≤ 256 * 128 = 32768` — far below `usize::MAX`, no multiply overflow.
    // The allocation is **fallible** (`try_reserve`) so a constrained target that cannot fit the graph
    // returns `AllocationFailed` rather than aborting — the same no-abort invariant (R11) that
    // `Puncturer::puncture`/`depuncture` already honour.
    let num_nodes = states * period;
    let mut adj: Vec<Vec<usize>> = Vec::new();
    adj.try_reserve_exact(num_nodes)
        .map_err(|_| PunctureError::AllocationFailed)?;
    adj.resize_with(num_nodes, Vec::new);
    for state in 0..states {
        for (phase, &mask) in masks.iter().enumerate() {
            let node = state * period + phase;
            // Reserve the ≤ 2 out-edges up front so the inner `push`es never realloc (hence never abort).
            adj[node]
                .try_reserve(2)
                .map_err(|_| PunctureError::AllocationFailed)?;
            for input in 0u8..2 {
                if state == 0 && input == 0 {
                    continue; // exclude the trivial all-zero self-branch
                }
                let sym = raw_output_symbol(state, input, params);
                if output_weight_after_puncture(sym, mask) == 0 {
                    let next = next_state(state, input, k);
                    let next_phase = (phase + 1) % period;
                    adj[node].push(next * period + next_phase);
                }
            }
        }
    }
    product_graph_has_cycle(&adj)
}

/// Iterative (explicit-stack) three-colour DFS back-edge cycle detection.
///
/// Returns `true` as soon as any edge reaches a node still on the DFS stack (grey) — a cycle.
/// `White` = unvisited, `Grey` = on the current DFS path, `Black` = fully explored. No recursion, so
/// a deep graph cannot overflow the call stack (spec §4.4).
fn product_graph_has_cycle(adj: &[Vec<usize>]) -> Result<bool, PunctureError> {
    #[derive(Clone, Copy, PartialEq, Eq)]
    enum Color {
        White,
        Grey,
        Black,
    }
    // Fallible allocation (R11 no-abort): a constrained target returns `AllocationFailed` rather than
    // aborting on OOM.
    let mut color: Vec<Color> = Vec::new();
    color
        .try_reserve_exact(adj.len())
        .map_err(|_| PunctureError::AllocationFailed)?;
    color.resize(adj.len(), Color::White);
    for start in 0..adj.len() {
        if color[start] != Color::White {
            continue;
        }
        color[start] = Color::Grey;
        // Stack frames hold `(node, next child index)` so exploration resumes where it left off. Every
        // growth goes through `try_reserve` so a push can never abort on OOM.
        let mut stack: Vec<(usize, usize)> = Vec::new();
        stack
            .try_reserve(1)
            .map_err(|_| PunctureError::AllocationFailed)?;
        stack.push((start, 0));
        while let Some(&(node, child)) = stack.last() {
            if child < adj[node].len() {
                stack.last_mut().unwrap().1 += 1;
                let neighbor = adj[node][child];
                match color[neighbor] {
                    Color::Grey => return Ok(true), // back-edge → cycle
                    Color::White => {
                        color[neighbor] = Color::Grey;
                        stack
                            .try_reserve(1)
                            .map_err(|_| PunctureError::AllocationFailed)?;
                        stack.push((neighbor, 0));
                    }
                    Color::Black => {}
                }
            } else {
                color[node] = Color::Black;
                stack.pop();
            }
        }
    }
    Ok(false)
}

/// The **single** cyclic-prefix body map shared by the puncturer (TX) and de-puncturer (RX).
///
/// Yields one `(coded_index, kept)` per coded bit of the first `nbits` trellis stages. The
/// `coded_index = t * n + j` is always the SOURCE position in the full interleaved rate-1/n stream
/// (the encoder emits `C1(t), …, Cn(t)` per stage), and `kept` is the keep-pattern's decision at that
/// `(stage, output)` under the cyclic-with-prefix rule (`phase = t % period`). The *sequence* in which
/// pairs are yielded is the **on-air emission order**, and it follows [`PunctureMatrix::order`]:
/// - [`PunctureOrder::Interleaved`]: iterate `t` outer, `j` inner → the emission order equals
///   `coded_index` order (`C1(0), C2(0), C1(1), …`, CCSDS XY for `n = 2`).
/// - [`PunctureOrder::ByColumn`]: iterate `j` outer, `t` inner → all of output `C1` across the body,
///   then all of `C2`, … (DVB-S / 802.11 labeling). The `coded_index` is still `t * n + j`, so it is
///   no longer monotonic in emission order — RX must place kept samples at `coded_index`, not append.
///
/// Factoring it here guarantees TX deletion and RX reinsertion replay the *identical* map for **both**
/// orders, so `depuncture(puncture(x))` is exact by construction (RP4). The `m·n` verbatim tail
/// (`m = k-1`) is handled separately by each caller (it is exempt from puncturing, §4.4).
///
/// Both orders are produced by a **single concrete `Map`** over the emission index `e ∈ [0, nbits·n)`
/// (no `Box`, no heap allocation): `e` is decomposed into `(t, j)` per [`PunctureMatrix::order`] —
/// [`Interleaved`](PunctureOrder::Interleaved) walks `t` outer / `j` inner (`t = e/n, j = e%n`),
/// [`ByColumn`](PunctureOrder::ByColumn) walks `j` outer / `t` inner (`j = e/nbits, t = e%nbits`) —
/// then mapped to the source `coded_index = t·n + j` and its keep decision `keep[j][t % period]`.
/// This reproduces **both** orders' emission sequences exactly while staying one non-boxed iterator.
fn body_keep_map(
    matrix: &PunctureMatrix,
    n: usize,
    nbits: usize,
) -> impl Iterator<Item = (usize, bool)> + '_ {
    let period = matrix.period;
    let order = matrix.order;
    // ByColumn divides by `nbits`; use a provably-non-zero divisor so the decomposition can never
    // divide by zero even if the iterator were forced. The range `0..(nbits*n)` is empty when
    // `nbits == 0`, so the closure never runs and the result is unaffected; when `nbits > 0`, `d == nbits`.
    let d = nbits.max(1);
    (0..(nbits * n)).map(move |e| {
        // Decompose the emission index into (time step, output row) per the on-air order. The range
        // `0..(nbits*n)` is empty when `nbits == 0`, so the closure never runs and the ByColumn arm's
        // `j = e / nbits` (and `t = e % nbits`) never divides by zero — the division is unreachable
        // rather than guarded.
        let (t, j) = match order {
            PunctureOrder::Interleaved => (e / n, e % n),
            PunctureOrder::ByColumn => (e % d, e / d),
        };
        (t * n + j, matrix.keep[j][t % period])
    })
}

/// Transmit-side puncturer: deletes coded bits per a [`PunctureMatrix`] to reach a higher code rate.
///
/// A separable wrapping layer over the clean rate-1/n core (D6): the encoder is **not** modified.
/// Construct **once per `(code, matrix)`** and reuse across frames — [`Puncturer::new`] runs the
/// one-time `O(S · period)` catastrophic-puncture check, so it must not be rebuilt per
/// [`Puncturer::puncture`] call.
///
/// # Interop hazard
/// The [`PunctureOrder`] and the payload `nbits` are **cross-boundary frame contracts**: the receive
/// side must agree on both. A global mismatch is caught by the BER sanity gate, but no single block can
/// detect it — see the module docs.
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct Puncturer {
    matrix: PunctureMatrix,
    n: usize,
    /// Mother-code memory `m = k - 1` (the verbatim-tail length in stages), captured from `params`
    /// at construction so the tail split is the code's own `k-1`, not the crate default `M`.
    m: usize,
}

impl Puncturer {
    /// Build a puncturer for `matrix` over the mother code `params`.
    ///
    /// # Errors
    /// - [`PunctureError::NMismatch`] if `matrix.n()` differs from `params.n()` (a rate-1/2 matrix can
    ///   never be applied to a rate-1/3 stream by accident).
    /// - [`PunctureError::Catastrophic`] if the `(mother code + pattern)` combination is catastrophic
    ///   (a hard gate — the punctured custom-matrix path has no other safety net). This check needs the
    ///   generator polynomials, which is why it lives here and not in [`PunctureMatrix::new`].
    pub fn new(matrix: PunctureMatrix, params: &CodeParams) -> Result<Self, PunctureError> {
        let matrix_n = matrix.n();
        let params_n = params.n();
        if matrix_n != params_n {
            return Err(PunctureError::NMismatch { matrix_n, params_n });
        }
        if punctured_is_catastrophic(&matrix, params)? {
            return Err(PunctureError::Catastrophic);
        }
        Ok(Self {
            matrix,
            n: matrix_n,
            m: params.k() as usize - 1,
        })
    }

    /// Puncture `coded` (the full rate-1/n coded block for a `nbits`-bit payload).
    ///
    /// The keep-pattern is applied **cyclically** over the coded bits of the first `nbits` trellis
    /// stages (the body); if the body is not a whole number of periods, the final partial period uses
    /// the matching prefix of the pattern (no data padding). The last `m·n` coded bits (the `m = k-1`
    /// zero-tail stages) are transmitted **verbatim** (tail exemption, spec §4.4) so the decoder
    /// converges cleanly to state 0. The tail length `m` is the mother code's own `k-1` (6 for K=7,
    /// 8 for K=9), captured at construction. `nbits` here is the **payload** bit count — the same value
    /// `encode`/`decode` take — not `coded.nbits` (the coded length).
    ///
    /// # Errors
    /// - [`PunctureError::PayloadTooLarge`] if `nbits` exceeds [`MAX_SUPPORTED_INFO_BITS`] (the payload
    ///   cap, symmetric with the encoder/decoder — R18).
    /// - [`PunctureError::MisalignedInput`] if `coded.nbits != (nbits + m) * n` (computed
    ///   with checked arithmetic, so an adversarial near-`usize::MAX` `nbits` fails gracefully rather
    ///   than overflow-panicking), or if `coded.bytes` is too short to hold `coded.nbits` bits (a
    ///   hand-built block whose `nbits` passes the length check but whose buffer is undersized). A
    ///   malformed or hand-built [`CodedBlock`] therefore never index-panics.
    /// - [`PunctureError::AllocationFailed`] if the fallible `try_reserve` of the output byte buffer
    ///   fails — the punctured output is packed through `try_reserve`, so an OOM returns an error rather
    ///   than aborting (R11 no-abort invariant).
    pub fn puncture(&self, coded: &CodedBlock, nbits: usize) -> Result<CodedBlock, PunctureError> {
        let n = self.n;

        // Payload cap symmetric with encode/decode (R18): reject an over-large `nbits` up front. This
        // also keeps every product below (`nbits * n`, the body kept-count) provably within `usize`, so
        // none can overflow. A dedicated `PayloadTooLarge` variant carries the clean semantics ("payload
        // over the cap") rather than overloading `MisalignedInput` with a `usize::MAX` sentinel.
        if nbits > MAX_SUPPORTED_INFO_BITS {
            return Err(PunctureError::PayloadTooLarge {
                got: nbits,
                cap: MAX_SUPPORTED_INFO_BITS,
            });
        }

        // Checked expected coded length: both the `+m` add and the `*n` mul are checked, so a huge
        // `nbits` yields `None` → MisalignedInput, never a wrapping/overflow panic.
        match nbits
            .checked_add(self.m)
            .and_then(|stages| stages.checked_mul(n))
        {
            Some(e) if e == coded.nbits => {}
            Some(e) => {
                return Err(PunctureError::MisalignedInput {
                    got: coded.nbits,
                    expected: e,
                })
            }
            None => {
                return Err(PunctureError::MisalignedInput {
                    got: coded.nbits,
                    expected: usize::MAX,
                })
            }
        }

        // Guard the byte buffer length against `nbits` BEFORE any `get_bit`: a hand-built
        // [`CodedBlock`] whose `nbits` matches the expected coded length but whose `bytes` is too
        // short would index out of bounds in the loop below → panic. Mirror `decode_block`'s
        // `div_ceil` check so a malformed block returns [`PunctureError::MisalignedInput`] rather than
        // index-panicking (crate no-panic invariant). `saturating_mul` keeps the reported `expected`
        // from overflowing for a near-`usize::MAX` `coded.nbits`.
        if coded.nbits.div_ceil(8) > coded.bytes.len() {
            return Err(PunctureError::MisalignedInput {
                got: coded.nbits,
                expected: coded.nbits.div_ceil(8).saturating_mul(8),
            });
        }

        // Compute the punctured bit length up front so the output can be written **packed** directly
        // (no per-bit `Vec<u8>` intermediate, which cost 8× the memory). `nbits ≤ MAX_SUPPORTED_INFO_BITS`
        // and the validated `coded.nbits == (nbits + m) * n` bound every product here, so none overflow.
        let period = self.matrix.period;
        let full_periods = nbits / period; // period ≥ 1, never divides by zero
        let rem = nbits % period;
        let prefix_kept: usize = (0..rem)
            .map(|col| self.matrix.keep.iter().filter(|row| row[col]).count())
            .sum();
        let kept_body = full_periods * self.matrix.kept_per_period + prefix_kept;
        // `coded.nbits == (nbits + m) * n` was validated above, so `coded.nbits >= nbits * n` and the
        // subtraction is provably non-underflowing (`= m·n`, the verbatim tail). The `debug_assert`
        // pins the invariant for any future edit; it can never fire in release under the length check.
        debug_assert!(
            coded.nbits >= nbits * n,
            "tail underflow: coded.nbits < nbits*n"
        );
        let tail = coded.nbits - nbits * n; // = m·n verbatim tail; in range by the validated length
        let out_nbits = kept_body + tail;

        // Fallible allocation of the packed output: OOM → `AllocationFailed`, never an abort (R11).
        let mut bytes: Vec<u8> = Vec::new();
        bytes
            .try_reserve(out_nbits.div_ceil(8))
            .map_err(|_| PunctureError::AllocationFailed)?;
        bytes.resize(out_nbits.div_ceil(8), 0);

        // Write the kept body bits (via the shared cyclic keep-map) then the verbatim tail directly as
        // packed bits (MSB-first via `set_bit`); `oi` is the running output bit index.
        let mut oi = 0usize;
        for (idx, kept) in body_keep_map(&self.matrix, n, nbits) {
            if kept {
                if get_bit(&coded.bytes, idx) == 1 {
                    set_bit(&mut bytes, oi);
                }
                oi += 1;
            }
        }
        for idx in (nbits * n)..coded.nbits {
            if get_bit(&coded.bytes, idx) == 1 {
                set_bit(&mut bytes, oi);
            }
            oi += 1;
        }
        debug_assert_eq!(oi, out_nbits);

        Ok(CodedBlock {
            bytes,
            nbits: out_nbits,
        })
    }
}

/// Receive-side de-puncturer: the exact inverse of [`Puncturer::puncture`], reconstructing the full
/// rate-1/n sample stream so the decoder core runs unchanged (D6, RP3).
///
/// It reinserts the metric's neutral erasure sample [`BranchMetric::erasure`] at every punctured
/// **body** position and copies the `m·n` verbatim **tail** samples unchanged (`m = k-1`), so a decode sees a
/// full-length stream in which punctured positions cost zero (spec §3.3/§4.4). The layer is generic
/// over the branch metric `M` and **binds the erasure sentinel to `M::erasure()`** — a wrong sentinel
/// is impossible because the caller never supplies one.
///
/// Construct **once per `(code, matrix)`** and reuse across frames — [`DePuncturer::new`] runs the same
/// one-time `O(S · period)` catastrophic-puncture check as [`Puncturer::new`] (the *shared*
/// module-private `punctured_is_catastrophic`, so TX and RX can never diverge on the verdict).
///
/// # Interop hazard
/// The [`PunctureOrder`] and the payload `nbits` are **cross-boundary frame contracts**: the transmit
/// side must agree on both. A global mismatch is caught by the BER sanity gate, but no single block can
/// detect it — see the module docs.
pub struct DePuncturer<Me: BranchMetric> {
    matrix: PunctureMatrix,
    n: usize,
    /// Mother-code memory `m = k - 1` (the verbatim-tail length in stages), captured from `params`
    /// at construction so the tail split matches the transmit side's code, not the crate default `M`.
    m: usize,
    _marker: core::marker::PhantomData<Me>,
}

// Manual trait impls (not `#[derive]`): the only `Me`-typed field is a zero-sized
// `PhantomData<Me>`, so these hold for **any** metric `Me` regardless of whether `Me` itself is
// `Debug`/`Clone`/`Eq`. A `#[derive]` would wrongly add `Me: Debug` etc. bounds that a unit metric
// struct (e.g. `SoftLlr`) need not satisfy.
impl<Me: BranchMetric> core::fmt::Debug for DePuncturer<Me> {
    fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
        f.debug_struct("DePuncturer")
            .field("matrix", &self.matrix)
            .field("n", &self.n)
            .field("m", &self.m)
            .finish()
    }
}
impl<Me: BranchMetric> Clone for DePuncturer<Me> {
    fn clone(&self) -> Self {
        Self {
            matrix: self.matrix.clone(),
            n: self.n,
            m: self.m,
            _marker: core::marker::PhantomData,
        }
    }
}
impl<Me: BranchMetric> PartialEq for DePuncturer<Me> {
    fn eq(&self, other: &Self) -> bool {
        self.matrix == other.matrix && self.n == other.n && self.m == other.m
    }
}
impl<Me: BranchMetric> Eq for DePuncturer<Me> {}

impl<Me: BranchMetric> DePuncturer<Me> {
    /// Build a de-puncturer for `matrix` over the mother code `params`.
    ///
    /// Applies the **same guards as [`Puncturer::new`]** — TX and RX must agree exactly:
    ///
    /// # Errors
    /// - [`PunctureError::NMismatch`] if `matrix.n()` differs from `params.n()` (a rate-1/2 matrix can
    ///   never be applied to a rate-1/3 stream by accident).
    /// - [`PunctureError::Catastrophic`] if the `(mother code + pattern)` combination is catastrophic —
    ///   evaluated by the **shared** module-private `punctured_is_catastrophic`, so the receive side can
    ///   never accept a pattern the transmit side rejects (or vice versa).
    pub fn new(matrix: PunctureMatrix, params: &CodeParams) -> Result<Self, PunctureError> {
        let matrix_n = matrix.n();
        let params_n = params.n();
        if matrix_n != params_n {
            return Err(PunctureError::NMismatch { matrix_n, params_n });
        }
        if punctured_is_catastrophic(&matrix, params)? {
            return Err(PunctureError::Catastrophic);
        }
        Ok(Self {
            matrix,
            n: matrix_n,
            m: params.k() as usize - 1,
            _marker: core::marker::PhantomData,
        })
    }

    /// Deterministic on-air sample length for a `nbits`-bit payload: kept **body** bits (the cyclic
    /// keep-pattern over the first `nbits` trellis stages) plus the full `m·n` verbatim **tail**
    /// (`m = k-1`, the mother code's own memory — 6 for K=7, 8 for K=9).
    ///
    /// Because the tail is exempt from puncturing (§4.4) the stream is *not* uniformly periodic, so the
    /// length is computed from `nbits` rather than a modulo. All products are evaluated with
    /// `checked_add`/`checked_mul`, so an adversarial `nbits` whose `(nbits + m) · n` layout would
    /// overflow `usize` yields **`None`** rather than wrapping — [`depuncture`](Self::depuncture) maps
    /// `None` to [`PunctureError::MisalignedInput`].
    #[must_use]
    pub fn expected_punctured_len(&self, nbits: usize) -> Option<usize> {
        let period = self.matrix.period;
        let n = self.n;
        // Body: full periods contribute `kept_per_period` each; the final partial period contributes the
        // count of kept bits over its prefix of columns (cyclic-with-prefix — identical to `puncture`).
        let full_periods = nbits / period; // period >= 1, never divides by zero
        let rem = nbits % period;
        let prefix_kept: usize = (0..rem).map(|col| self.column_kept(col)).sum();
        let body = full_periods
            .checked_mul(self.matrix.kept_per_period)?
            .checked_add(prefix_kept)?;
        // Tail: the last `m = k-1` stages are transmitted verbatim → `m·n` samples.
        let tail = self.m.checked_mul(n)?;
        body.checked_add(tail)
    }

    /// Reconstruct the full `(nbits + m)·n` sample stream from the received on-air samples (`m = k-1`).
    ///
    /// Inserts [`BranchMetric::erasure`] at each punctured **body** position (cyclic keep-pattern,
    /// identical to [`Puncturer::puncture`]'s split) and copies the `m·n` **tail** samples verbatim, so
    /// the decoder core runs on a full-length stream where punctured positions cost zero.
    ///
    /// # Errors
    /// - [`PunctureError::PayloadTooLarge`] if `nbits` exceeds [`MAX_SUPPORTED_INFO_BITS`] (the payload
    ///   cap, symmetric with [`Puncturer::puncture`] and the encoder/decoder — R18).
    /// - [`PunctureError::MisalignedInput`] if `received.len()` differs from
    ///   [`expected_punctured_len`](Self::expected_punctured_len)`(nbits)` — including the overflow case
    ///   (`None` → `expected = usize::MAX`), so an adversarial `nbits` fails gracefully rather than
    ///   overflow- or index-panicking.
    /// - [`PunctureError::AllocationFailed`] if the fallible `try_reserve` of the output buffer fails.
    ///
    /// Never panics on any input (spec §3.3/RX2).
    pub fn depuncture(
        &self,
        received: &[Me::Sample],
        nbits: usize,
    ) -> Result<Vec<Me::Sample>, PunctureError> {
        let n = self.n;
        // Payload cap symmetric with `Puncturer::puncture` and encode/decode (R18): reject an over-large
        // `nbits` up front with the dedicated `PayloadTooLarge` variant, closing the asymmetric-cap gap.
        if nbits > MAX_SUPPORTED_INFO_BITS {
            return Err(PunctureError::PayloadTooLarge {
                got: nbits,
                cap: MAX_SUPPORTED_INFO_BITS,
            });
        }
        // Validate the received length against the checked expected length. `None` (overflow) is treated
        // as a mismatch with `expected = usize::MAX`, mirroring `Puncturer::puncture`.
        let expected = self.expected_punctured_len(nbits);
        match expected {
            Some(e) if e == received.len() => {}
            Some(e) => {
                return Err(PunctureError::MisalignedInput {
                    got: received.len(),
                    expected: e,
                })
            }
            None => {
                return Err(PunctureError::MisalignedInput {
                    got: received.len(),
                    expected: usize::MAX,
                })
            }
        }

        // Full reconstructed length `(nbits + m)·n`, checked (defensive — reachable only for a validated
        // `received`, but computed without any unchecked product so it can never overflow-panic).
        let full_len = match nbits
            .checked_add(self.m)
            .and_then(|stages| stages.checked_mul(n))
        {
            Some(len) => len,
            None => {
                return Err(PunctureError::MisalignedInput {
                    got: received.len(),
                    expected: usize::MAX,
                })
            }
        };

        let mut out: Vec<Me::Sample> = Vec::new();
        out.try_reserve(full_len)
            .map_err(|_| PunctureError::AllocationFailed)?;

        // The reconstructed stream is ALWAYS interleaved (`out[t*n+j]`) — the decoder contract —
        // regardless of the on-air `PunctureOrder`. Under `ByColumn` the shared keep-map yields
        // `coded_index` non-monotonically, so kept samples must be PLACED at their `coded_index`, not
        // appended in emission order. Prefill the body with `Me::erasure()`, then overwrite kept
        // positions. (For `Interleaved` the placement order equals the emission order, so this is
        // identical to the previous append — no behaviour change.)
        // `body_len = nbits*n = full_len - m*n`; the subtraction is in range because `m*n` is the
        // tail portion of the already-checked `full_len`.
        let body_len = full_len - self.m * n;
        out.resize(body_len, Me::erasure());
        let mut ri = 0usize;
        for (idx, kept) in body_keep_map(&self.matrix, n, nbits) {
            if kept {
                // `idx = t*n+j < nbits*n = body_len` by construction → always in bounds (no panic).
                out[idx] = received[ri];
                ri += 1;
            }
        }
        // Tail: the remaining `m·n` received samples, copied verbatim (order-exempt, appended after
        // the body in both TX and RX so `received[ri..]` is exactly the tail).
        for &s in &received[ri..] {
            out.push(s);
        }
        debug_assert_eq!(out.len(), full_len);
        Ok(out)
    }

    /// Count of kept (`true`) bits in column `col` of the keep-pattern (one trellis stage's contribution
    /// to the body length). Callers pass `col < period`, so the row indexing is in-bounds.
    fn column_kept(&self, col: usize) -> usize {
        // Documented invariant: callers always pass `col = t % period < period`, so the row indexing is
        // in-bounds. The `debug_assert` catches a future caller that violates it (never hit in release).
        debug_assert!(
            col < self.matrix.period,
            "column_kept: col must be < period"
        );
        self.matrix.keep.iter().filter(|row| row[col]).count()
    }
}

#[cfg(test)]
mod tests {
    use super::raw_output_symbol;
    use crate::params::CodeParams;
    use crate::trellis::output_symbol;

    /// [`raw_output_symbol`] deliberately duplicates [`crate::trellis::output_symbol`] **minus** the
    /// per-generator CCSDS inversion (catastrophicity is inversion-invariant, §4.4). The two therefore
    /// have a maintenance coupling: a future edit to either that drifts their symbol packing apart would
    /// silently corrupt the catastrophic-puncture check. This cross-check pins the exact relation —
    /// `output_symbol == raw_output_symbol ^ inversion_mask` — over the full `(state, input)` domain of
    /// both mother codes, so any drift fails loudly here.
    #[test]
    fn raw_output_symbol_matches_trellis_output_symbol_minus_inversion() {
        for params in [CodeParams::ccsds_r1_2(), CodeParams::ccsds_r1_3()] {
            let n = params.generators().len();
            // inversion_mask: bit `n-1-i` set iff generator `i` is inverted (high-bit-first packing,
            // matching both `raw_output_symbol` and `output_symbol`).
            let inv_mask: u32 = params
                .invert_outputs()
                .iter()
                .enumerate()
                .filter(|(_, &inv)| inv)
                .fold(0u32, |mask, (i, _)| mask | (1u32 << (n - 1 - i)));
            let states = 1usize << (params.k() - 1);
            for state in 0..states {
                for input in 0u8..2 {
                    let raw = raw_output_symbol(state, input, &params);
                    let inverted = output_symbol(state, input, &params);
                    assert_eq!(
                        inverted,
                        raw ^ inv_mask,
                        "raw_output_symbol drifted from output_symbol at \
                         state={state} input={input} (k={})",
                        params.k()
                    );
                }
            }
        }
    }
}