rvoip-codec-core 0.3.8

//! ISF quantisation, encoder side: 3GPP TS 26.190 §5.2.6, in fixed point.
//!
//! Implements TS 26.173 `qpisf_2s.c`: `Qpisf_2s_46b` (every mode except
//! 6.60 kbit/s), `Qpisf_2s_36b` (6.60 kbit/s only), the two searches they rest
//! on — `VQ_stage1` and `Sub_VQ` — and the `bfi == 0` half of
//! `Dpisf_2s_46b` / `Dpisf_2s_36b`, which the reference encoder calls to
//! reconstruct what it just transmitted and to advance the predictor.
//!
//! Validated bit-exactly against `testdata/wb_enc_trace.txt`, produced by the
//! TS 26.173 encoder itself at 12.65 kbit/s: the `isf_unq46` rows go in, and
//! both the seven transmitted indices (`isf_indice46`) and the reconstructed
//! spectrum (`isf_q`) come out, replayed across all three committed frames from
//! the reset state.
//!
//! The committed trace covers 12.65 kbit/s, so the tests below reach the
//! 36-bit rate only through its reconstruction. Both rates were additionally
//! checked, off-tree, against full 50-frame traces regenerated by
//! `tools/trace-amrwb-encoder.sh 2` and `… 0` — every index and every ISF of
//! every frame, at 46 and 36 bits alike. Committing a 36-bit fixture would need
//! a second trace in `testdata`, which is the assembler's call, not this
//! module's.
//!
//! # Why the indices are the real test
//!
//! Everything here is a *decision*, not just an arithmetic result. Two codebook
//! entries can dequantise to spectra that differ by a handful of LSBs, so a
//! search that resolves a tie the other way still yields perfectly plausible
//! speech — and a bitstream no conformant decoder reproduces. The comparison
//! operators and visit orders below are therefore normative, and are the part
//! of this module most worth reading twice:
//!
//! * every scan runs over ascending codebook index, exhaustively, with no
//!   pruning and no early exit;
//! * every comparison is a strict `<`, so a tie keeps the incumbent, which is
//!   always the lower codebook index (or, among survivors, the one with the
//!   smaller first-stage distance);
//! * the first-stage distance is *not* part of the final objective. It only
//!   chooses which four candidates reach the second stage, and in what order.
//!
//! # Prediction, and therefore state
//!
//! The vector coded is not the ISF vector but its residual after removing a
//! long-term mean and a third of the previous frame's quantised residual. That
//! residual is this module's only state ([`IsfEncoder`]), it resets to zero —
//! not to the flat spectrum the decoder's *output* history resets to — and a
//! frame decoded in isolation is the wrong frame. It is also what the encoder
//! and decoder must agree on: they run the same update, from the same indices.
//!
//! # Q-formats
//!
//! ISFs and every codebook, mean and residual here are Q15 normalised
//! frequency, where 16384 is 6400 Hz. Distances are Q31-ish `Word32`
//! accumulators — twice the sum of squared Q15 errors — and are only ever
//! compared against each other, so the factor of two is immaterial except that
//! it is where saturation happens.

use crate::codecs::amr::wb::lp::autocorr::LP_ORDER;
use crate::codecs::amr::wb::lp::isf_codebooks::{
    DICO1, DICO2, DICO21, DICO21_36B, DICO22, DICO22_36B, DICO23, DICO23_36B, DICO24, DICO25,
    MEAN_ISF,
};
use crate::codecs::amr::wb::lp::isf_dequant::IsfQuantizer;
use crate::fixed_point::arith::{add, mult, sub};
use crate::fixed_point::arith32::{l_add, l_mac, l_sub};
use crate::fixed_point::types::{DspContext, Word16, Word32, MAX_32};

/// Prediction factor, 1/3 in Q15.
///
/// The header comment in the reference's `qpisf_2s.tab` calls this 0.25; the
/// code says 10923, which is a third. The comment is wrong.
const MU: Word16 = Word16(10923);

/// Minimum spacing between adjacent quantised ISFs, Q15 — 50 Hz.
pub(crate) const ISF_GAP: Word16 = Word16(128);

/// How many first-stage candidates are refined by the second stage.
///
/// `N_SURV_MAX` in the reference. Both encoder call sites pass four, and the
/// reference's own `dist_min[]` is sized for no more.
const SURVIVORS: usize = 4;

/// Widest first-stage vector, so one scratch buffer serves both splits.
const MAX_SPLIT_DIM: usize = 9;

/// Most refinement books either rate puts on one split.
const MAX_REFINEMENTS: usize = 3;

/// Most indices either rate transmits.
const MAX_INDICES: usize = 7;

/// How the sixteen ISFs are split between the two first-stage codebooks.
const LOW_DIM: usize = 9;

/// One half of the ISF vector, and the codebooks that code it.
///
/// The first stage codes `dim` consecutive ISF residuals as a single vector;
/// the second stage refines that choice in narrower pieces, whose dimensions
/// sum back to `dim`. Search and reconstruction are both driven from this one
/// description, so the two cannot disagree about which book covers which ISFs —
/// a mismatch there is invisible in the encoder and catastrophic at the decoder.
#[derive(Clone, Copy, Debug)]
struct Split {
    /// First-stage codebook, flat and row-major: entry `n` is at `n * dim`.
    stage1: &'static [i16],
    /// First-stage vector dimension.
    dim: usize,
    /// Refinement books with their dimensions, in ISF order.
    stage2: &'static [(&'static [i16], usize)],
}

/// 46-bit low split: ISFs 0..9, refined in three groups of three.
const LOW_46: Split = Split {
    stage1: &DICO1,
    dim: LOW_DIM,
    stage2: &[(&DICO21, 3), (&DICO22, 3), (&DICO23, 3)],
};

/// 46-bit high split: ISFs 9..16, refined as three then four.
const HIGH_46: Split = Split {
    stage1: &DICO2,
    dim: LP_ORDER - LOW_DIM,
    stage2: &[(&DICO24, 3), (&DICO25, 4)],
};

/// 36-bit low split: the same first stage as 46 bits, refined as five then four.
const LOW_36: Split = Split {
    stage1: &DICO1,
    dim: LOW_DIM,
    stage2: &[(&DICO21_36B, 5), (&DICO22_36B, 4)],
};

/// 36-bit high split: one un-split seven-dimensional refinement.
const HIGH_36: Split = Split {
    stage1: &DICO2,
    dim: LP_ORDER - LOW_DIM,
    stage2: &[(&DICO23_36B, 7)],
};

/// The two splits a rate uses, low first.
const fn splits(rate: IsfQuantizer) -> (Split, Split) {
    match rate {
        IsfQuantizer::Bits46 => (LOW_46, HIGH_46),
        IsfQuantizer::Bits36 => (LOW_36, HIGH_36),
    }
}

/// One frame's quantiser output: what to transmit, and what was transmitted.
///
/// The reconstruction is not a convenience — the encoder needs it, because
/// every stage downstream (the stability measure, the interpolated synthesis
/// filter, the DTX history) must see the ISFs the *decoder* will see, not the
/// ones the analysis produced.
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub struct QuantizedIsf {
    /// Reconstructed ISFs, Q15, minimum spacing enforced.
    isf: [Word16; LP_ORDER],
    /// Transmitted indices; only the first [`Self::indices`] entries are live.
    indices: [u16; MAX_INDICES],
    /// How many of `indices` this rate transmits.
    count: usize,
}

impl QuantizedIsf {
    /// The reconstructed ISFs, Q15 — identical to what the decoder will produce
    /// from [`Self::indices`].
    #[must_use]
    pub const fn isf(&self) -> &[Word16; LP_ORDER] {
        &self.isf
    }

    /// The quantiser indices, in transmission order: first-stage low,
    /// first-stage high, then the refinements low-to-high. Seven entries at
    /// 46 bits, five at 36.
    #[must_use]
    pub fn indices(&self) -> &[u16] {
        &self.indices[..self.count]
    }
}

/// The ISF quantiser's carried state: one predictor memory.
#[derive(Clone, Debug)]
pub struct IsfEncoder {
    /// Previous frame's quantised residual, Q15 — the sum of the chosen
    /// codevectors, before the mean and the prediction were added back.
    past: [Word16; LP_ORDER],
}

impl Default for IsfEncoder {
    fn default() -> Self {
        Self::new()
    }
}

impl IsfEncoder {
    /// A quantiser in its reset state.
    ///
    /// Zero, not the flat spectrum `ISF_INIT`: this memory holds a *residual*,
    /// and a zero residual means "the mean, as predicted". The reference resets
    /// it at initialisation, after every SID frame, and after an encoder homing
    /// frame — all three go through this.
    #[must_use]
    pub const fn new() -> Self {
        Self {
            past: [Word16(0); LP_ORDER],
        }
    }

    /// Return to the reset state, as the reference's `Reset_encoder` does.
    pub const fn reset(&mut self) {
        self.past = [Word16(0); LP_ORDER];
    }

    /// Quantise one frame's ISFs.
    ///
    /// `isf` is the unquantised analysis ISF vector, Q15, sixteen elements —
    /// element 15 carries half the usual scale and is not part of the monotone
    /// ladder, which is why the spacing rule leaves it alone. The result holds
    /// both the indices to transmit and the ISFs the decoder will reconstruct,
    /// Q15. The predictor state advances by one frame.
    pub fn quantize(&mut self, isf: &[Word16; LP_ORDER], rate: IsfQuantizer) -> QuantizedIsf {
        let mut ctx = DspContext::default();
        let (low, high) = splits(rate);

        let residual = self.residual(&mut ctx, isf);

        // The two halves are searched independently and sequentially; there is
        // no joint optimisation across them, and each starts from a fresh best.
        let mut low_picks = [0u16; MAX_REFINEMENTS];
        let mut high_picks = [0u16; MAX_REFINEMENTS];
        let low_stage1 = search_split(&mut ctx, &residual[..LOW_DIM], low, &mut low_picks);
        let high_stage1 = search_split(&mut ctx, &residual[LOW_DIM..], high, &mut high_picks);

        let mut indices = [0u16; MAX_INDICES];
        indices[0] = low_stage1;
        indices[1] = high_stage1;
        let refinements = low.stage2.len() + high.stage2.len();
        indices[2..2 + low.stage2.len()].copy_from_slice(&low_picks[..low.stage2.len()]);
        indices[2 + low.stage2.len()..2 + refinements]
            .copy_from_slice(&high_picks[..high.stage2.len()]);

        let isf = self.reconstruct(&mut ctx, rate, &indices);

        QuantizedIsf {
            isf,
            indices,
            count: 2 + refinements,
        }
    }

    /// Remove the mean and the prediction, giving the vector actually coded.
    ///
    /// Q15 in, Q15 out. The two subtractions saturate separately: folding them
    /// into one wide subtract changes the result whenever the intermediate
    /// leaves 16 bits. `mult` floors, so a negative predictor memory predicts
    /// one LSB lower than a rounding implementation would — which is most
    /// frames, on most coefficients.
    fn residual(&self, ctx: &mut DspContext, isf: &[Word16; LP_ORDER]) -> [Word16; LP_ORDER] {
        let mut residual = [Word16(0); LP_ORDER];
        for (i, slot) in residual.iter_mut().enumerate() {
            let centred = sub(ctx, isf[i], Word16(MEAN_ISF[i]));
            let predicted = mult(ctx, MU, self.past[i]);
            *slot = sub(ctx, centred, predicted);
        }
        residual
    }

    /// Rebuild the ISFs from the chosen indices and advance the predictor.
    ///
    /// This is the reference's `Dpisf_2s_*` with `bfi = 0` and `enc_dec = 0`:
    /// the encoder reconstructs by running the decoder's own arithmetic, which
    /// is the only way the two stay in step. Q15 throughout.
    fn reconstruct(
        &mut self,
        ctx: &mut DspContext,
        rate: IsfQuantizer,
        indices: &[u16; MAX_INDICES],
    ) -> [Word16; LP_ORDER] {
        let (low, high) = splits(rate);
        let mut isf = [Word16(0); LP_ORDER];
        let refinements = low.stage2.len();

        place_codevectors(ctx, &mut isf[..LOW_DIM], low, indices[0], &indices[2..]);
        place_codevectors(
            ctx,
            &mut isf[LOW_DIM..],
            high,
            indices[1],
            &indices[2 + refinements..],
        );

        // Add back the mean and a third of the last residual, then keep this
        // frame's residual for the next one. The read of `past[i]` must precede
        // the write: hoisting the write out of the loop, or vectorising it the
        // other way round, gives a wrong but entirely plausible spectrum.
        for (i, slot) in isf.iter_mut().enumerate() {
            let residual = *slot;
            let centred = add(ctx, residual, Word16(MEAN_ISF[i]));
            let predicted = mult(ctx, MU, self.past[i]);
            *slot = add(ctx, centred, predicted);
            self.past[i] = residual;
        }

        // Spacing is enforced after the predictor update, so it never feeds
        // back into the memory the decoder also keeps.
        enforce_min_spacing(ctx, &mut isf);

        isf
    }
}

/// Write one split's first-stage codevector, then add its refinements.
///
/// `stage2` may be longer than the split needs; the extra entries belong to the
/// other split and are ignored.
fn place_codevectors(
    ctx: &mut DspContext,
    dst: &mut [Word16],
    split: Split,
    stage1: u16,
    stage2: &[u16],
) {
    let base = usize::from(stage1) * split.dim;
    for (slot, &value) in dst.iter_mut().zip(&split.stage1[base..base + split.dim]) {
        *slot = Word16(value);
    }

    let mut offset = 0;
    for (&(book, dim), &index) in split.stage2.iter().zip(stage2) {
        let base = usize::from(index) * dim;
        for (slot, &value) in dst[offset..offset + dim]
            .iter_mut()
            .zip(&book[base..base + dim])
        {
            *slot = add(ctx, *slot, Word16(value));
        }
        offset += dim;
    }
}

/// Force a minimum spacing between adjacent ISFs, Q15.
///
/// A forward-only push: nothing is ever lowered and nothing is exchanged,
/// despite the reference calling this a reorder. The floor starts at the gap
/// itself, so ISF 0 is clamped too, and the last ISF is deliberately never
/// examined — it carries half the usual scale and is not part of the ladder.
pub(crate) fn enforce_min_spacing(ctx: &mut DspContext, isf: &mut [Word16; LP_ORDER]) {
    let mut floor = ISF_GAP;
    for slot in isf.iter_mut().take(LP_ORDER - 1) {
        if slot.0 < floor.0 {
            *slot = floor;
        }
        // Unconditional, and off the possibly-just-clamped value.
        floor = add(ctx, *slot, ISF_GAP);
    }
}

/// Search one half of the ISF vector: N-best first stage, exhaustive second.
///
/// `residual` is that half of the mean- and prediction-removed vector, Q15.
/// Returns the first-stage index and writes the refinement indices into
/// `picks`.
///
/// The objective minimised across survivors is the *second-stage* error alone.
/// The first-stage distance chose the survivors and ordered them; adding it
/// into the total would pick a lower-error codevector than the reference does,
/// and fail conformance while sounding better.
fn search_split(
    ctx: &mut DspContext,
    residual: &[Word16],
    split: Split,
    picks: &mut [u16; MAX_REFINEMENTS],
) -> u16 {
    let survivors = n_best_entries(ctx, residual, split.stage1, split.dim);

    let mut best_error = Word32(MAX_32);
    // The reference writes its indices only inside the `if`, so a split whose
    // every total saturated would leave them at whatever C had there before.
    // Unreachable with real residuals; zero here, and the strict `<` stays.
    let mut best_stage1 = 0;
    let mut refined = [Word16(0); MAX_SPLIT_DIM];

    // Survivors in ascending first-stage distance — the order `n_best_entries`
    // produced. Ties below resolve by position in this list, so it must not be
    // re-sorted, and every survivor is evaluated however good an earlier one was.
    for &stage1 in &survivors {
        let base = stage1 * split.dim;
        let stage1_vector = &split.stage1[base..base + split.dim];
        for (slot, (&x, &c)) in refined
            .iter_mut()
            .zip(residual.iter().zip(stage1_vector.iter()))
        {
            *slot = sub(ctx, x, Word16(c));
        }

        let mut total = Word32(0);
        let mut candidate = [0u16; MAX_REFINEMENTS];
        let mut offset = 0;
        for (n, (&(book, dim), pick)) in split.stage2.iter().zip(candidate.iter_mut()).enumerate() {
            let (index, error) = nearest_entry(ctx, &refined[offset..offset + dim], book);
            *pick = index;
            // Assigned for the first split and saturating-accumulated after —
            // the saturation can manufacture a tie a wider accumulator would
            // have broken, and the tie-break below is then load-bearing.
            total = if n == 0 {
                error
            } else {
                l_add(ctx, total, error)
            };
            offset += dim;
        }

        // Strict: an equal total keeps the earlier survivor, i.e. the one with
        // the smaller first-stage distance. This is the single likeliest source
        // of a silent conformance failure in the whole quantiser.
        if l_sub(ctx, total, best_error).0 < 0 {
            best_error = total;
            best_stage1 = stage1;
            *picks = candidate;
        }
    }

    u16::try_from(best_stage1).expect("first-stage codebooks hold at most 256 entries")
}

/// The `SURVIVORS` nearest entries of `book`, nearest first.
///
/// `target` is Q15 and `dim` long; `book` is flat and row-major. The returned
/// list is ordered by (distance ascending, codebook index ascending).
fn n_best_entries(
    ctx: &mut DspContext,
    target: &[Word16],
    book: &[i16],
    dim: usize,
) -> [usize; SURVIVORS] {
    let mut best = [Word32(MAX_32); SURVIVORS];

    // Seeded with 0,1,2,3 rather than left empty. A slot that no *accepted*
    // candidate ever fills keeps its seed and is still handed to the second
    // stage — only reachable when distances saturate, but the decision must not
    // depend on that never happening.
    let mut index = [0usize; SURVIVORS];
    for (slot, seed) in index.iter_mut().zip(0..) {
        *slot = seed;
    }

    for (i, entry) in book.chunks_exact(dim).enumerate() {
        let dist = squared_distance(ctx, target, entry);

        // Strict `<`, so an equal distance does not displace an incumbent and
        // is tried against the next slot instead: among equals, the lower
        // codebook index ranks first. A distance saturated to MAX_32 compares
        // equal to an empty slot and is therefore rejected outright.
        let slot = best.iter().position(|&d| l_sub(ctx, dist, d).0 < 0);
        if let Some(k) = slot {
            // Ordered insert with truncation: the worst survivor falls off.
            best.copy_within(k..SURVIVORS - 1, k + 1);
            index.copy_within(k..SURVIVORS - 1, k + 1);
            best[k] = dist;
            index[k] = i;
        }
    }

    index
}

/// The single nearest entry of `book` to `target`, with its distance.
///
/// `target` is Q15 and its length is the codebook's dimension. Exhaustive,
/// ascending, strict `<` — so a tie goes to the lower index — and the index is
/// seeded to 0, which is what a book whose every distance saturates returns.
///
/// The reference also copies the winning codevector back over `target`. That
/// write is dead at every call site here (each survivor recomputes the
/// refinement target from scratch, and the splits within a survivor are
/// disjoint), so this returns the index instead of mutating.
pub(crate) fn nearest_entry(
    ctx: &mut DspContext,
    target: &[Word16],
    book: &[i16],
) -> (u16, Word32) {
    let mut best = Word32(MAX_32);
    let mut index = 0usize;

    for (i, entry) in book.chunks_exact(target.len()).enumerate() {
        let dist = squared_distance(ctx, target, entry);
        if l_sub(ctx, dist, best).0 < 0 {
            best = dist;
            index = i;
        }
    }

    (
        u16::try_from(index).expect("refinement codebooks hold at most 128 entries"),
        best,
    )
}

/// Twice the squared Euclidean distance between a Q15 target and one codevector.
///
/// Plain and unweighted — AMR-WB's ISF search has no perceptual weighting
/// anywhere, unlike AMR-NB's, which weights every error term before squaring.
/// The factor of two comes from `L_mult` and is uniform, so it cannot change an
/// ordering; the saturation around it can, which is why it stays.
pub(crate) fn squared_distance(ctx: &mut DspContext, target: &[Word16], entry: &[i16]) -> Word32 {
    let mut acc = Word32(0);
    for (&x, &c) in target.iter().zip(entry) {
        let error = sub(ctx, x, Word16(c));
        acc = l_mac(ctx, acc, error, error);
    }
    acc
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::codecs::amr::wb::lp::isf_dequant::IsfDecoder;

    /// The reference encoder's own intermediates at 12.65 kbit/s, three frames.
    const TRACE: &str = include_str!("../../testdata/wb_enc_trace.txt");

    /// One frame-level trace row, as integers.
    fn frame_row(frame: usize, name: &str) -> Vec<i32> {
        let prefix = format!("T {frame} -1 {name} ");
        let line = TRACE
            .lines()
            .find(|l| l.starts_with(&prefix))
            .unwrap_or_else(|| panic!("the trace has no frame {frame} row {name:?}"));
        line[prefix.len()..]
            .split_whitespace()
            .map(|v| v.parse().expect("trace values are integers"))
            .collect()
    }

    /// How many frames the trace holds a given row for.
    fn traced_frames(name: &str) -> usize {
        let tail = format!(" -1 {name} ");
        TRACE
            .lines()
            .filter(|l| l.starts_with("T ") && l.contains(&tail))
            .count()
    }

    fn isf_vector(frame: usize, name: &str) -> [Word16; LP_ORDER] {
        let row = frame_row(frame, name);
        assert_eq!(row.len(), LP_ORDER, "frame {frame} row {name:?} length");
        let mut isf = [Word16(0); LP_ORDER];
        for (slot, &v) in isf.iter_mut().zip(row.iter()) {
            *slot = Word16(i16::try_from(v).expect("ISFs are Word16"));
        }
        isf
    }

    /// Replay every traced frame through one encoder, keeping the outputs.
    ///
    /// The predictor carries between frames, so the frames have to be fed in
    /// order and from the reset state; quantising frame 1 alone gives a
    /// different — and wrong — answer.
    fn replay_trace() -> Vec<QuantizedIsf> {
        let frames = traced_frames("isf_unq46");
        assert!(
            frames >= 3,
            "the trace holds {frames} frames; fewer than two leaves the predictor untested"
        );

        let mut enc = IsfEncoder::new();
        (0..frames)
            .map(|f| enc.quantize(&isf_vector(f, "isf_unq46"), IsfQuantizer::Bits46))
            .collect()
    }

    #[test]
    fn the_46_bit_search_picks_the_reference_indices() {
        // The strongest available check: two codebook entries can dequantise to
        // near-identical spectra, so an index that agrees rules out a tie
        // resolved the other way in a way a vector comparison cannot.
        let got = replay_trace();
        let mut compared = 0;

        for (f, frame) in got.iter().enumerate() {
            let want = frame_row(f, "isf_indice46");
            assert_eq!(want.len(), 7, "frame {f}: the 46-bit rate transmits seven");
            assert_eq!(frame.indices().len(), 7, "frame {f}: indices produced");
            for (i, (&mine, &theirs)) in frame.indices().iter().zip(want.iter()).enumerate() {
                assert_eq!(
                    i32::from(mine),
                    theirs,
                    "frame {f}: index {i} = {mine} but the reference chose {theirs}"
                );
                compared += 1;
            }
        }

        assert_eq!(compared, 21, "compared {compared} indices, expected 3 x 7");
    }

    #[test]
    fn the_46_bit_quantizer_reconstructs_the_reference_isfs() {
        let got = replay_trace();
        let mut compared = 0;

        for (f, frame) in got.iter().enumerate() {
            let want = frame_row(f, "isf_q");
            assert_eq!(want.len(), LP_ORDER, "frame {f}: isf_q length");
            for (i, (&mine, &theirs)) in frame.isf().iter().zip(want.iter()).enumerate() {
                assert_eq!(
                    i32::from(mine.0),
                    theirs,
                    "frame {f}: isf[{i}] = {} but the reference gives {theirs}",
                    mine.0
                );
                compared += 1;
            }
        }

        assert_eq!(
            compared,
            3 * LP_ORDER,
            "compared {compared} ISFs, expected 48"
        );
    }

    #[test]
    fn the_predictor_carries_between_frames() {
        // Guards the replay above against a state update that does nothing: if
        // the memory were dropped, frame 1 in sequence would equal frame 1
        // quantised cold, and every trace comparison would still be run — just
        // against a quantiser with no prediction in it.
        let sequence = replay_trace();

        let mut cold = IsfEncoder::new();
        let alone = cold.quantize(&isf_vector(1, "isf_unq46"), IsfQuantizer::Bits46);

        assert_ne!(
            alone.indices(),
            sequence[1].indices(),
            "frame 1 quantised cold chose the same indices as frame 1 in sequence, \
             so the predictor memory is not reaching the search"
        );
    }

    #[test]
    fn a_reset_encoder_repeats_itself() {
        let first = replay_trace();

        let mut enc = IsfEncoder::new();
        enc.quantize(&isf_vector(0, "isf_unq46"), IsfQuantizer::Bits46);
        enc.quantize(&isf_vector(1, "isf_unq46"), IsfQuantizer::Bits46);
        enc.reset();

        let mut compared = 0;
        for (f, want) in first.iter().enumerate() {
            let got = enc.quantize(&isf_vector(f, "isf_unq46"), IsfQuantizer::Bits46);
            assert_eq!(got, *want, "frame {f} differs after a reset");
            compared += 1;
        }
        assert_eq!(compared, first.len(), "compared {compared} frames");
    }

    #[test]
    fn both_rates_reconstruct_what_the_dequantiser_decodes() {
        // The committed trace is 12.65 kbit/s only, so there is no reference row
        // for the 36-bit rate. What there is, is a dequantiser already proved
        // bit-exact against TS 26.173: the encoder must reconstruct exactly what
        // it will produce, or encoder and decoder drift apart within a frame.
        // This covers the 36-bit index layout, whose slot assignment differs.
        for rate in [IsfQuantizer::Bits46, IsfQuantizer::Bits36] {
            let mut enc = IsfEncoder::new();
            let mut dec = IsfDecoder::new();
            let mut compared = 0;

            for f in 0..traced_frames("isf_unq46") {
                let out = enc.quantize(&isf_vector(f, "isf_unq46"), rate);
                assert_eq!(
                    out.indices().len(),
                    rate.index_count(),
                    "{rate:?} frame {f}: index count"
                );
                let decoded = dec.decode(rate, out.indices(), false);
                for (i, (&mine, &theirs)) in out.isf().iter().zip(decoded.iter()).enumerate() {
                    assert_eq!(
                        mine.0, theirs.0,
                        "{rate:?} frame {f}: encoder has isf[{i}] = {} where the \
                         dequantiser gives {}",
                        mine.0, theirs.0
                    );
                    compared += 1;
                }
            }

            assert_eq!(compared, 3 * LP_ORDER, "{rate:?}: compared {compared} ISFs");
        }
    }

    /// Four two-dimensional entries, the first two equidistant from `[1, 0]`.
    const TIED_STAGE1: [i16; 8] = [0, 0, 2, 0, 40, 0, 60, 0];

    /// One entry at the origin, so refinement error is the survivor's own error.
    const FLAT_STAGE2: [i16; 2] = [0, 0];

    /// Two entries equidistant from `[1, 0]`, again.
    const TIED_STAGE2: [i16; 4] = [0, 0, 2, 0];

    /// Four entries at spread-out distances from the origin: 32, 800, and two
    /// far enough away to fill the survivor list without being chosen.
    const SPREAD_STAGE1: [i16; 8] = [4, 0, 20, 0, 1000, 0, 2000, 0];

    /// A lone refinement entry that survivor 1 of [`SPREAD_STAGE1`] reaches
    /// exactly and survivor 0 does not.
    const DISTANT_STAGE2: [i16; 2] = [-20, 0];

    #[test]
    fn a_first_stage_tie_ranks_the_lower_codebook_index_first() {
        let mut ctx = DspContext::default();
        let target = [Word16(1), Word16(0)];

        // Entries 0 and 1 are both one away; 2 and 3 are far. Strict `<` means
        // the later of the two equals cannot displace the earlier.
        let survivors = n_best_entries(&mut ctx, &target, &TIED_STAGE1, 2);
        assert_eq!(
            survivors,
            [0, 1, 2, 3],
            "survivors should be distance-ascending then index-ascending"
        );

        // Mirrored: with the target past entry 1, entry 1 wins outright, which
        // shows the list above was ordered by distance and not by luck.
        let survivors = n_best_entries(&mut ctx, &[Word16(2), Word16(0)], &TIED_STAGE1, 2);
        assert_eq!(survivors, [1, 0, 2, 3]);
    }

    #[test]
    fn a_refinement_tie_goes_to_the_lower_codebook_index() {
        let mut ctx = DspContext::default();
        let (index, error) = nearest_entry(&mut ctx, &[Word16(1), Word16(0)], &TIED_STAGE2);
        assert_eq!(
            index, 0,
            "an equal distance must not displace the incumbent"
        );
        assert_eq!(error.0, 2, "twice the squared error of one LSB");

        let (index, _) = nearest_entry(&mut ctx, &[Word16(2), Word16(0)], &TIED_STAGE2);
        assert_eq!(index, 1, "a strictly nearer later entry must win");
    }

    #[test]
    fn a_saturating_search_keeps_its_seeds() {
        // Every distance saturates to MAX_32, which compares equal to an empty
        // slot and is rejected by the strict `<`. The survivor list is then the
        // seed 0,1,2,3, and the refinement search returns entry 0 — not
        // undefined, and not "the least bad".
        let mut ctx = DspContext::default();
        let target = [Word16(i16::MAX); MAX_SPLIT_DIM];
        let book = [i16::MIN; MAX_SPLIT_DIM * SURVIVORS];

        let survivors = n_best_entries(&mut ctx, &target, &book, MAX_SPLIT_DIM);
        assert_eq!(survivors, [0, 1, 2, 3]);

        let (index, error) = nearest_entry(&mut ctx, &target, &book);
        assert_eq!(index, 0);
        assert_eq!(
            error.0, MAX_32,
            "the accumulator saturates rather than wraps"
        );
    }

    #[test]
    fn the_survivor_loop_keeps_the_earlier_survivor_on_a_tie() {
        // Survivors 0 and 1 sit one LSB either side of the target, so their
        // refinement errors against a single origin entry are equal. The
        // reference keeps the earlier survivor — the one the first stage ranked
        // higher — and an implementation using `<=` would silently take the
        // later one and encode a different, equally good, frame.
        let mut ctx = DspContext::default();
        let split = Split {
            stage1: &TIED_STAGE1,
            dim: 2,
            stage2: &[(&FLAT_STAGE2, 2)],
        };
        let mut picks = [0u16; MAX_REFINEMENTS];

        let chosen = search_split(&mut ctx, &[Word16(1), Word16(0)], split, &mut picks);
        assert_eq!(chosen, 0, "an equal refinement error must keep survivor 0");
        assert_eq!(picks[0], 0);

        // And a survivor with a strictly smaller refinement error does win, so
        // the assertion above is about the tie and not about the loop stalling.
        let chosen = search_split(&mut ctx, &[Word16(2), Word16(0)], split, &mut picks);
        assert_eq!(chosen, 1);
    }

    #[test]
    fn the_first_stage_distance_stays_out_of_the_objective() {
        // Survivor 0 is much nearer in the first stage; survivor 1 refines
        // exactly. The reference decides on refinement error alone, so
        // survivor 1 wins. The numbers are chosen so that summing both stages
        // reverses that: 32 + 512 = 544 for survivor 0 against 800 + 0 = 800
        // for survivor 1. An implementation that "improves" the objective by
        // including the first-stage distance therefore fails here and nowhere
        // else — the spectra it produces are, if anything, slightly better.
        let mut ctx = DspContext::default();
        let split = Split {
            stage1: &SPREAD_STAGE1,
            dim: 2,
            stage2: &[(&DISTANT_STAGE2, 2)],
        };
        let mut picks = [0u16; MAX_REFINEMENTS];

        let chosen = search_split(&mut ctx, &[Word16(0), Word16(0)], split, &mut picks);
        assert_eq!(
            chosen, 1,
            "the survivor with the larger first-stage distance refines better and must win"
        );
    }

    #[test]
    fn the_spacing_rule_clamps_the_first_isf_and_spares_the_last() {
        let mut ctx = DspContext::default();
        let mut isf = [Word16(0); LP_ORDER];
        isf[15] = Word16(3);

        enforce_min_spacing(&mut ctx, &mut isf);

        assert_eq!(
            isf[0].0, ISF_GAP.0,
            "the floor starts at the gap, not at isf[0]"
        );
        for (i, slot) in isf.iter().enumerate().take(LP_ORDER - 1).skip(1) {
            assert_eq!(
                slot.0,
                ISF_GAP.0 * i16::try_from(i + 1).expect("fifteen fits"),
                "isf[{i}] should have been pushed one gap above its predecessor"
            );
        }
        assert_eq!(
            isf[15].0, 3,
            "the last ISF carries half the usual scale and is never clamped"
        );
    }

    #[test]
    fn the_two_rates_share_a_first_stage_but_not_its_winner() {
        // Both rates run the same first stage over the same codebooks, so both
        // must choose from the same four survivors — but the winner is decided
        // by refinement error, and the refinement books differ. On frame 0 the
        // high split lands on a different survivor at each rate. Caching a
        // first-stage *decision* across rates, rather than the survivor list,
        // would be wrong in exactly the way that still sounds fine.
        let isf = isf_vector(0, "isf_unq46");
        let wide = IsfEncoder::new().quantize(&isf, IsfQuantizer::Bits46);
        let narrow = IsfEncoder::new().quantize(&isf, IsfQuantizer::Bits36);
        assert_eq!(wide.indices().len(), 7);
        assert_eq!(narrow.indices().len(), 5);

        let mut ctx = DspContext::default();
        let residual = IsfEncoder::new().residual(&mut ctx, &isf);
        let low = n_best_entries(&mut ctx, &residual[..LOW_DIM], &DICO1, LOW_DIM);
        let high = n_best_entries(&mut ctx, &residual[LOW_DIM..], &DICO2, LP_ORDER - LOW_DIM);

        let mut checked = 0;
        for picked in [wide.indices(), narrow.indices()] {
            assert!(
                low.contains(&usize::from(picked[0])),
                "low-split index {} is not among the survivors {low:?}",
                picked[0]
            );
            assert!(
                high.contains(&usize::from(picked[1])),
                "high-split index {} is not among the survivors {high:?}",
                picked[1]
            );
            checked += 2;
        }
        assert_eq!(checked, 4, "checked {checked} first-stage indices");
    }
}