oxideav-av1 0.1.21

//! §7.17 Loop restoration — the post-CDEF restoration pass that sits
//! **after** the §7.15 CDEF de-ringing pass per av1-spec p.327-335.
//!
//! ## Coverage (round 198 — close-out push)
//!
//! This module covers the §7.17 loop-restoration top-level driver, the
//! §7.17.4 Wiener arm, and the §7.17.2 / §7.17.3 self-guided
//! projection arm end-to-end at the sample-filtering level:
//!
//! * §7.17 — [`loop_restoration_frame`]: top-level walk over `(y, x)`
//!   in `MI_SIZE` steps, dispatching per plane / per
//!   `FrameRestorationType` to the per-block driver.
//! * §7.17.1 — [`loop_restore_block`]: per-`MI_SIZE` block driver that
//!   derives `(unitRow, unitCol, x, y, w, h, StripeStartY, StripeEndY,
//!   PlaneEndX, PlaneEndY)` and dispatches by `rType` to the Wiener,
//!   self-guided, or no-op arm.
//! * §7.17.2 — [`self_guided_filter`]: invokes the §7.17.3 box filter
//!   twice (`pass = 0` with `(r0, eps0)` and `pass = 1` with
//!   `(r1, eps1)` from `Sgr_Params[set]`) and combines `flt0` / `flt1`
//!   with the projection weights `(w0, w1, w2)` against the
//!   `UpscaledCdefFrame << SGRPROJ_RST_BITS` reference.
//! * §7.17.3 — [`box_filter`]: builds the `A[]` / `B[]` integral-style
//!   per-sample mean / variance arrays over a `(2r + 1)²` window with
//!   the `Sgr_Params[set][pass * 2 + 1]` eps-driven `m`-table
//!   nonlinearity, then convolves `A` / `B` with the §7.17.3 weighted
//!   3×3 footprint (pass-0 uses only odd-row neighbours; pass-1 uses
//!   the centre-plus-cross 4 / corner-3 weighting).
//! * §7.17.4 — [`wiener_filter`]: 7-tap horizontal × 7-tap vertical
//!   separable convolution with the §7.11.3.2 `(InterRound0,
//!   InterRound1)` rounding-shift schedule.
//! * §7.17.5 — [`wiener_coefficients`]: reconstruct the symmetric 7-tap
//!   filter (`filter[3] = 128 - 2 * Σ coeff[i]`; `filter[i] =
//!   filter[6 - i] = coeff[i]`).
//! * §7.17.6 — [`get_source_sample`]: per-sample fetch that snaps to
//!   `[0, PlaneEndX] × [0, PlaneEndY]` and routes to either the
//!   `UpscaledCdefFrame` (inside the current stripe) or the
//!   `UpscaledCurrFrame` (above / below the stripe, with the cropped
//!   `StripeStartY - 2` / `StripeEndY + 2` neighbour-line reach).
//!
//! ## Constant tables (av1-spec p.107 / p.332)
//!
//! * [`WIENER_TAPS_MID`] / [`WIENER_TAPS_MIN`] / [`WIENER_TAPS_MAX`] /
//!   [`WIENER_TAPS_K`] — §5.11.57 / §5.11.58 reference midpoint, range,
//!   and subexp `k` for the three transmitted Wiener coefficients.
//! * [`SGRPROJ_XQD_MID`] / [`SGRPROJ_XQD_MIN`] / [`SGRPROJ_XQD_MAX`] —
//!   §5.11.57 reference midpoint and range for the two transmitted
//!   self-guided projection weights.
//! * [`SGR_PARAMS`] — §7.17.3 `Sgr_Params[16][4]` per av1-spec p.332
//!   lines 18395-18400. Each `set` row is `(r0, eps0, r1, eps1)`.
//!
//! ## Standalone-friendly surface
//!
//! Like the §7.14 / §7.15 drivers, the top-level driver takes a small
//! [`LoopRestorationFrameContext`] bundling the frame-level inputs plus
//! closures for the §5.11.x decode state the driver reads:
//!
//! * Per-unit `LrType[plane][unitRow][unitCol]` (§5.11.58).
//! * Per-unit `LrWiener[plane][unitRow][unitCol][pass][i]` (§5.11.58 —
//!   only consulted when `LrType == RESTORE_WIENER`).
//! * Per-unit `LrSgrSet[plane][unitRow][unitCol]` (§5.11.58 — only
//!   consulted when `LrType == RESTORE_SGRPROJ`; the index into
//!   [`SGR_PARAMS`]).
//! * Per-unit `LrSgrXqd[plane][unitRow][unitCol][i]` (§5.11.58 —
//!   `i ∈ {0, 1}`; only consulted when `LrType == RESTORE_SGRPROJ`).
//! * Per-plane subsampling, bit depth, MI dimensions, plane sizes, and
//!   the [`crate::LrParams`] schedule.
//!
//! ## Bitstream-conformance gates not enforced here
//!
//! * The caller is responsible for the §5.11.58 `Wiener_Taps_{Min,Max}`
//!   clipping (the bit-reader applies it).
//! * The caller has already populated `UpscaledCurrFrame[plane]` (pre-
//!   CDEF samples) and `UpscaledCdefFrame[plane]` (post-CDEF samples)
//!   per §7.15.
//! * Plane buffers are sized at `Round2(FrameHeight, subY) ×
//!   Round2(UpscaledWidth, subX)` per av1-spec §5.11.34.
//!
//! ## Out-of-scope for this arc
//!
//! * §7.18 output process — separate pass that runs after loop
//!   restoration.
//! * Cross-plane SIMD / cache-friendly batched filtering — the
//!   reference loop here mirrors the spec's per-sample formulation.

use crate::cdf::MI_SIZE;
use crate::loop_filter::PlaneBuffer;
use crate::uncompressed_header_tail::{FrameRestorationType, LrParams};

// =====================================================================
// §7.17 constant lookup tables — av1-spec p.107 / p.332.
// =====================================================================

/// `FILTER_BITS = 7` per av1-spec p.32 line 1312 — bit precision of the
/// Wiener filter coefficients.
pub const FILTER_BITS: u32 = 7;

/// `WIENER_COEFFS = 3` per av1-spec p.32 line 1314 — number of
/// transmitted Wiener filter coefficients (the symmetric 7-tap filter
/// has 4 independent values; the centre tap is derived).
pub const WIENER_COEFFS: usize = 3;

/// `SGRPROJ_PARAMS_BITS = 4` per av1-spec p.32 line 1316 — bit count of
/// the §5.11.58 `lr_sgr_set` selector.
pub const SGRPROJ_PARAMS_BITS: u32 = 4;

/// `SGRPROJ_PRJ_SUBEXP_K = 4` per av1-spec p.32 line 1319 — `k`
/// parameter of the §5.11.58 subexp decoder for self-guided projection
/// weight deltas.
pub const SGRPROJ_PRJ_SUBEXP_K: u32 = 4;

/// `SGRPROJ_PRJ_BITS = 7` per av1-spec p.32 line 1321 — precision of
/// the self-guided projection weights.
pub const SGRPROJ_PRJ_BITS: u32 = 7;

/// `SGRPROJ_RST_BITS = 4` per av1-spec p.32 line 1323 — extra
/// restoration precision bits the box-filter output carries above the
/// 8-bit sample range.
pub const SGRPROJ_RST_BITS: u32 = 4;

/// `SGRPROJ_MTABLE_BITS = 20` per av1-spec p.32 line 1326 — precision
/// of the §7.17.3 `m`-table division shift.
pub const SGRPROJ_MTABLE_BITS: u32 = 20;

/// `SGRPROJ_RECIP_BITS = 12` per av1-spec p.32 line 1328 — precision of
/// the `1/n` division-by-area reciprocal.
pub const SGRPROJ_RECIP_BITS: u32 = 12;

/// `SGRPROJ_SGR_BITS = 8` per av1-spec p.32 line 1330 — internal
/// precision of the box-filter core.
pub const SGRPROJ_SGR_BITS: u32 = 8;

/// `Wiener_Taps_Mid[3]` per av1-spec p.74 line 3884 — reference centre
/// of the three transmitted Wiener coefficients (used by §5.11.57's
/// frame-state reset).
pub const WIENER_TAPS_MID: [i32; WIENER_COEFFS] = [3, -7, 15];

/// `Wiener_Taps_Min[3]` per av1-spec p.107 line 6484 — minimum value of
/// the three transmitted Wiener coefficients.
pub const WIENER_TAPS_MIN: [i32; WIENER_COEFFS] = [-5, -23, -17];

/// `Wiener_Taps_Max[3]` per av1-spec p.107 line 6485 — maximum value of
/// the three transmitted Wiener coefficients.
pub const WIENER_TAPS_MAX: [i32; WIENER_COEFFS] = [10, 8, 46];

/// `Wiener_Taps_K[3]` per av1-spec p.107 line 6486 — subexp `k`
/// parameter for the per-tap §5.11.58 delta read.
pub const WIENER_TAPS_K: [u32; WIENER_COEFFS] = [1, 2, 3];

/// `Sgrproj_Xqd_Mid[2]` per av1-spec p.74 line 3886 — reference centre
/// of the two transmitted self-guided projection weights.
pub const SGRPROJ_XQD_MID: [i32; 2] = [-32, 31];

/// `Sgrproj_Xqd_Min[2]` per av1-spec p.107 line 6488 — minimum value of
/// the two transmitted self-guided projection weights.
pub const SGRPROJ_XQD_MIN: [i32; 2] = [-96, -32];

/// `Sgrproj_Xqd_Max[2]` per av1-spec p.107 line 6489 — maximum value of
/// the two transmitted self-guided projection weights.
pub const SGRPROJ_XQD_MAX: [i32; 2] = [31, 95];

/// `Sgr_Params[16][4]` per av1-spec p.332 lines 18395-18400 — the
/// `(1 << SGRPROJ_PARAMS_BITS)` self-guided parameter sets, each
/// `(r0, eps0, r1, eps1)`. `r0` / `r1` are the box-filter radii and
/// `eps0` / `eps1` are the variance-scaling thresholds.
pub const SGR_PARAMS: [[i32; 4]; 16] = [
    [2, 12, 1, 4],
    [2, 15, 1, 6],
    [2, 18, 1, 8],
    [2, 21, 1, 9],
    [2, 24, 1, 10],
    [2, 29, 1, 11],
    [2, 36, 1, 12],
    [2, 45, 1, 13],
    [2, 56, 1, 14],
    [2, 68, 1, 15],
    [0, 0, 1, 5],
    [0, 0, 1, 8],
    [0, 0, 1, 11],
    [0, 0, 1, 14],
    [2, 30, 0, 0],
    [2, 75, 0, 0],
];

// =====================================================================
// §7.17 standalone-driver surface — av1-spec p.327.
// =====================================================================

/// Caller-supplied frame-level inputs the §7.17 driver consults.
///
/// All §5.5 / §5.9.5 / §5.9.20 fields come straight from the parsed
/// frame / sequence header. Per-unit predicates are closures so the
/// driver can be exercised in isolation without wiring the full
/// §5.11.58 decode state.
pub struct LoopRestorationFrameContext<'a> {
    /// `MiRows` per §5.9.5.
    pub mi_rows: u32,
    /// `MiCols` per §5.9.5.
    pub mi_cols: u32,
    /// `NumPlanes` per §5.5.2 — `1` for monochrome, `3` otherwise.
    pub num_planes: u8,
    /// `BitDepth` per §5.5.2 (`8`, `10`, or `12`).
    pub bit_depth: u8,
    /// `subsampling_x` per §5.5.2 (`0` or `1`).
    pub subsampling_x: u8,
    /// `subsampling_y` per §5.5.2 (`0` or `1`).
    pub subsampling_y: u8,
    /// `FrameHeight` per §5.9.5.
    pub frame_height: u32,
    /// `UpscaledWidth` per §5.9.8 (pre-superres-downscale).
    pub upscaled_width: u32,
    /// `LrParams` per §5.9.20 — `FrameRestorationType[0..3]`, `UsesLr`,
    /// and `LoopRestorationSize[0..3]`. The §7.17 walker indexes the
    /// `loop_restoration_size[plane]` array.
    pub lr_params: &'a LrParams,
    /// `LrType[plane][unitRow][unitCol]` per §5.11.58 — `RESTORE_NONE`
    /// / `RESTORE_WIENER` / `RESTORE_SGRPROJ`.
    pub lr_type: &'a dyn Fn(u8, u32, u32) -> FrameRestorationType,
    /// `LrWiener[plane][unitRow][unitCol][pass][i]` per §5.11.58 —
    /// `pass ∈ {0, 1}` selects vertical / horizontal coefficient sets;
    /// `i ∈ {0, 1, 2}` selects one of the three transmitted Wiener
    /// coefficients. Only consulted when the per-unit `LrType` is
    /// `RESTORE_WIENER`.
    pub lr_wiener: &'a dyn Fn(u8, u32, u32, u8, usize) -> i32,
    /// `LrSgrSet[plane][unitRow][unitCol]` per §5.11.58 — index in
    /// `[0, 16)` into [`SGR_PARAMS`]. Only consulted when the per-unit
    /// `LrType` is `RESTORE_SGRPROJ`.
    pub lr_sgr_set: &'a dyn Fn(u8, u32, u32) -> u8,
    /// `LrSgrXqd[plane][unitRow][unitCol][i]` per §5.11.58 —
    /// `i ∈ {0, 1}` selects one of the two transmitted self-guided
    /// projection weights (`w0` / `w1`); `w2` is derived as
    /// `(1 << SGRPROJ_PRJ_BITS) - w0 - w1`. Only consulted when the
    /// per-unit `LrType` is `RESTORE_SGRPROJ`.
    pub lr_sgr_xqd: &'a dyn Fn(u8, u32, u32, usize) -> i32,
}

impl std::fmt::Debug for LoopRestorationFrameContext<'_> {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        f.debug_struct("LoopRestorationFrameContext")
            .field("mi_rows", &self.mi_rows)
            .field("mi_cols", &self.mi_cols)
            .field("num_planes", &self.num_planes)
            .field("bit_depth", &self.bit_depth)
            .field("subsampling_x", &self.subsampling_x)
            .field("subsampling_y", &self.subsampling_y)
            .field("frame_height", &self.frame_height)
            .field("upscaled_width", &self.upscaled_width)
            .field("lr_params", &self.lr_params)
            .finish_non_exhaustive()
    }
}

// =====================================================================
// §7.17 top-level driver — av1-spec p.327 lines 18088-18106.
// =====================================================================

/// §7.17 top-level loop-restoration driver — av1-spec p.327 lines
/// 18088-18106.
///
/// Iterates the frame in `MI_SIZE`-sized blocks, dispatches each
/// `(plane, row, col)` triple whose `FrameRestorationType[plane] !=
/// RESTORE_NONE` to [`loop_restore_block`].
///
/// Reads pre-CDEF samples from `curr_planes` and post-CDEF samples from
/// `cdef_planes`; writes restored samples into `lr_planes`. The §7.17
/// "copy `LrFrame = UpscaledCdefFrame`" prelude is materialised by
/// copying `cdef_planes ↦ lr_planes` up-front; on `UsesLr == 0` the
/// driver returns immediately after the copy.
pub fn loop_restoration_frame(
    ctx: &LoopRestorationFrameContext<'_>,
    curr_planes: &[PlaneBuffer<'_>],
    cdef_planes: &[PlaneBuffer<'_>],
    lr_planes: &mut [PlaneBuffer<'_>],
) {
    let num_planes = ctx
        .num_planes
        .min(curr_planes.len() as u8)
        .min(cdef_planes.len() as u8)
        .min(lr_planes.len() as u8);
    // av1-spec p.327 line 18088: LrFrame = UpscaledCdefFrame.
    for plane in 0..num_planes {
        copy_plane(cdef_planes, lr_planes, plane);
    }
    // av1-spec p.327 line 18091: UsesLr == 0 ⇒ return.
    if !ctx.lr_params.uses_lr {
        return;
    }
    // av1-spec p.327 lines 18096-18106: the walk is specified in
    // `MI_SIZE` steps over `(y, x) ∈ [0, FrameHeight) × [0,
    // UpscaledWidth)`. r464: every MI block inside one (stripe, unit)
    // rectangle derives the same `(unitRow, unitCol, StripeStartY,
    // StripeEndY, PlaneEndX, PlaneEndY)` and the two filters are
    // pointwise functions of absolute sample position under that
    // geometry (the pass-0 row-parity weights ride an even rectangle
    // origin exactly as they ride the even MI-block origin), so the
    // frame is restored one rectangle at a time — identical output,
    // a fraction of the per-block overhead.
    for plane in 0..num_planes {
        let r_type = ctx.lr_params.frame_restoration_type[plane as usize];
        if r_type == FrameRestorationType::None {
            continue;
        }
        for geom in stripe_unit_rects(ctx, plane) {
            loop_restore_rect(ctx, curr_planes, cdef_planes, lr_planes, plane, &geom);
        }
    }
}

/// r464 — the (stripe, restoration-unit) rectangles of `plane`, each
/// carrying the §7.17.1 geometry every `MI_SIZE` block inside it
/// derives: stripes are the §7.17.1 `(lumaY + 8) / 64` bands (the
/// first starts at row 0, every later one at `(64·n - 8) >> subY`),
/// columns are the §5.9.20 `LoopRestorationSize[plane]` units with
/// the last unit absorbing the remainder, and `unitRow` / `unitCol`
/// are pegged to the unit grid like §7.17.1's.
#[must_use]
pub fn stripe_unit_rects(ctx: &LoopRestorationFrameContext<'_>, plane: u8) -> Vec<LrBlockGeometry> {
    let (sub_x, sub_y) = subsampling_for_plane(plane, ctx.subsampling_x, ctx.subsampling_y);
    let plane_w = round2(ctx.upscaled_width, sub_x);
    let plane_h = round2(ctx.frame_height, sub_y);
    let mut out = Vec::new();
    if plane_w == 0 || plane_h == 0 {
        return out;
    }
    let unit_size = ctx
        .lr_params
        .loop_restoration_size
        .get(plane as usize)
        .copied()
        .unwrap_or(0)
        .max(1);
    let unit_rows = count_units_in_frame(unit_size, plane_h);
    let unit_cols = count_units_in_frame(unit_size, plane_w);
    let plane_end_x = plane_w as i32 - 1;
    let plane_end_y = plane_h as i32 - 1;
    let mut stripe = 0i32;
    loop {
        let stripe_start_y = (-8 + stripe * 64) >> sub_y;
        let stripe_end_y = stripe_start_y + (64 >> sub_y) - 1;
        let y0 = stripe_start_y.max(0);
        if y0 > plane_end_y {
            break;
        }
        let y1 = stripe_end_y.min(plane_end_y);
        // `unitRow` from any luma row of the stripe: `((lumaY + 8) >>
        // subY) / unitSize` with `lumaY = 64·stripe - 8`.
        let unit_row = (((64 * stripe) as u32) >> sub_y) / unit_size;
        let unit_row = unit_row.min(unit_rows.saturating_sub(1));
        for unit_col in 0..unit_cols {
            let x0 = unit_col * unit_size;
            if x0 as i32 > plane_end_x {
                break;
            }
            let x1 = if unit_col + 1 == unit_cols {
                plane_end_x
            } else {
                (x0 + unit_size - 1).min(plane_end_x as u32) as i32
            };
            out.push(LrBlockGeometry {
                unit_row,
                unit_col,
                x: x0,
                y: y0 as u32,
                w: (x1 - x0 as i32 + 1) as u32,
                h: (y1 - y0 + 1) as u32,
                stripe_start_y,
                stripe_end_y,
                plane_end_x,
                plane_end_y,
            });
        }
        stripe += 1;
    }
    out
}

/// r464 — restore one rectangle under an explicit §7.17.1 geometry
/// (see [`stripe_unit_rects`]): `rType = LrType[plane][unitRow]
/// [unitCol]` dispatches to the Wiener / self-guided / no-op arm
/// exactly like [`loop_restore_block`].
pub fn loop_restore_rect(
    ctx: &LoopRestorationFrameContext<'_>,
    curr_planes: &[PlaneBuffer<'_>],
    cdef_planes: &[PlaneBuffer<'_>],
    lr_planes: &mut [PlaneBuffer<'_>],
    plane: u8,
    geom: &LrBlockGeometry,
) {
    if geom.w == 0 || geom.h == 0 {
        return;
    }
    let r_type = (ctx.lr_type)(plane, geom.unit_row, geom.unit_col);
    match r_type {
        FrameRestorationType::Wiener => {
            wiener_filter(ctx, curr_planes, cdef_planes, lr_planes, plane, geom);
        }
        FrameRestorationType::SgrProj => {
            self_guided_filter(ctx, curr_planes, cdef_planes, lr_planes, plane, geom);
        }
        FrameRestorationType::None | FrameRestorationType::Switchable => {}
    }
}

/// Copy `src_planes[plane]` into `dst_planes[plane]` in full — used by
/// the §7.17 driver's `LrFrame = UpscaledCdefFrame` prelude. Both
/// buffers carry the same `(rows, cols)` shape; on a mismatch the copy
/// honours the smaller extent on each axis.
fn copy_plane(src_planes: &[PlaneBuffer<'_>], dst_planes: &mut [PlaneBuffer<'_>], plane: u8) {
    let p = plane as usize;
    let Some(src) = src_planes.get(p) else { return };
    let Some(dst) = dst_planes.get_mut(p) else {
        return;
    };
    let rows = src.rows.min(dst.rows) as usize;
    let cols = src.cols.min(dst.cols) as usize;
    let src_cols = src.cols as usize;
    let dst_cols = dst.cols as usize;
    for y in 0..rows {
        for x in 0..cols {
            dst.samples[y * dst_cols + x] = src.samples[y * src_cols + x];
        }
    }
}

// =====================================================================
// §7.17.1 loop-restore block driver — av1-spec p.328-329 lines
// 18113-18215.
// =====================================================================

/// Per-block derived geometry returned by [`derive_block_geometry`].
///
/// The fields map directly to the §7.17.1 named locals
/// `(unitRow, unitCol, x, y, w, h, StripeStartY, StripeEndY,
/// PlaneEndX, PlaneEndY)` and feed both the Wiener and (future)
/// self-guided arms.
#[derive(Debug, Clone, Copy)]
pub struct LrBlockGeometry {
    /// `unitRow` per av1-spec p.328 line 18168.
    pub unit_row: u32,
    /// `unitCol` per av1-spec p.328 line 18171.
    pub unit_col: u32,
    /// `x` per av1-spec p.329 line 18182.
    pub x: u32,
    /// `y` per av1-spec p.329 line 18184.
    pub y: u32,
    /// `w` per av1-spec p.329 line 18186.
    pub w: u32,
    /// `h` per av1-spec p.329 line 18195.
    pub h: u32,
    /// `StripeStartY` per av1-spec p.328 line 18144.
    pub stripe_start_y: i32,
    /// `StripeEndY` per av1-spec p.328 line 18147.
    pub stripe_end_y: i32,
    /// `PlaneEndX` per av1-spec p.329 line 18176.
    pub plane_end_x: i32,
    /// `PlaneEndY` per av1-spec p.329 line 18179.
    pub plane_end_y: i32,
}

/// §7.17.1 geometry derivation — av1-spec p.328-329 lines 18122-18198.
///
/// Computes `(unitRow, unitCol, x, y, w, h, StripeStartY, StripeEndY,
/// PlaneEndX, PlaneEndY)` from the input `(plane, row, col)` triple
/// and the [`LoopRestorationFrameContext`].
pub fn derive_block_geometry(
    ctx: &LoopRestorationFrameContext<'_>,
    plane: u8,
    row: u32,
    col: u32,
) -> LrBlockGeometry {
    // av1-spec p.328 lines 18122-18124: stripeNum = (lumaY + 8) / 64.
    let luma_y = (row as i32) * (MI_SIZE as i32);
    let stripe_num = (luma_y + 8) / 64;
    // av1-spec p.328 lines 18138-18142: subX / subY from the plane.
    let (sub_x, sub_y) = subsampling_for_plane(plane, ctx.subsampling_x, ctx.subsampling_y);
    // av1-spec p.328 line 18144: StripeStartY = (-8 + stripeNum * 64) >> subY.
    let stripe_start_y = (-8 + stripe_num * 64) >> sub_y;
    // av1-spec p.328 line 18147: StripeEndY = StripeStartY + (64 >> subY) - 1.
    let stripe_end_y = stripe_start_y + (64 >> sub_y) - 1;
    // av1-spec p.328 line 18155: unitSize = LoopRestorationSize[plane].
    let unit_size = ctx
        .lr_params
        .loop_restoration_size
        .get(plane as usize)
        .copied()
        .unwrap_or(0)
        .max(1);
    // av1-spec p.328 lines 18158-18162: unitRows, unitCols via
    // count_units_in_frame.
    let unit_rows = count_units_in_frame(unit_size, round2(ctx.frame_height, sub_y));
    let unit_cols = count_units_in_frame(unit_size, round2(ctx.upscaled_width, sub_x));
    // av1-spec p.328 lines 18168-18172: unitRow / unitCol pegged to
    // `unitRows - 1` / `unitCols - 1`.
    let unit_row_raw = ((row * (MI_SIZE as u32) + 8) >> sub_y) / unit_size;
    let unit_row = unit_row_raw.min(unit_rows.saturating_sub(1));
    let unit_col_raw = ((col * (MI_SIZE as u32)) >> sub_x) / unit_size;
    let unit_col = unit_col_raw.min(unit_cols.saturating_sub(1));
    // av1-spec p.329 lines 18176-18179: PlaneEndX / PlaneEndY.
    let plane_end_x = round2(ctx.upscaled_width, sub_x) as i32 - 1;
    let plane_end_y = round2(ctx.frame_height, sub_y) as i32 - 1;
    // av1-spec p.329 lines 18182-18195: x / y / w / h.
    let x = (col * (MI_SIZE as u32)) >> sub_x;
    let y = (row * (MI_SIZE as u32)) >> sub_y;
    let w = ((MI_SIZE as u32) >> sub_x).min((plane_end_x + 1 - x as i32).max(0) as u32);
    let h = ((MI_SIZE as u32) >> sub_y).min((plane_end_y + 1 - y as i32).max(0) as u32);
    LrBlockGeometry {
        unit_row,
        unit_col,
        x,
        y,
        w,
        h,
        stripe_start_y,
        stripe_end_y,
        plane_end_x,
        plane_end_y,
    }
}

/// §7.17.1 per-block driver — av1-spec p.328-329 lines 18113-18215.
///
/// Computes the block's `(unitRow, unitCol, x, y, w, h, ...)`
/// geometry, fetches `rType = LrType[plane][unitRow][unitCol]`, and
/// dispatches to the Wiener / self-guided / no-op arm.
pub fn loop_restore_block(
    ctx: &LoopRestorationFrameContext<'_>,
    curr_planes: &[PlaneBuffer<'_>],
    cdef_planes: &[PlaneBuffer<'_>],
    lr_planes: &mut [PlaneBuffer<'_>],
    plane: u8,
    row: u32,
    col: u32,
) {
    let geom = derive_block_geometry(ctx, plane, row, col);
    if geom.w == 0 || geom.h == 0 {
        return;
    }
    // av1-spec p.329 line 18205: rType = LrType[plane][unitRow][unitCol].
    let r_type = (ctx.lr_type)(plane, geom.unit_row, geom.unit_col);
    match r_type {
        FrameRestorationType::Wiener => {
            wiener_filter(ctx, curr_planes, cdef_planes, lr_planes, plane, &geom);
        }
        FrameRestorationType::SgrProj => {
            self_guided_filter(ctx, curr_planes, cdef_planes, lr_planes, plane, &geom);
        }
        FrameRestorationType::None | FrameRestorationType::Switchable => {
            // RESTORE_NONE — `LrFrame` already holds `UpscaledCdefFrame`
            // from the §7.17 prelude. Switchable is a frame-level
            // setting; per-unit `rType` materialises to one of the
            // three primitive types, so this arm is unreachable on a
            // conformant decode but kept exhaustive.
        }
    }
}

// =====================================================================
// §7.17.2 Self-guided filter — av1-spec p.330 lines 18218-18273.
// §7.17.3 Box filter      — av1-spec p.331-332 lines 18277-18389.
// =====================================================================

/// §7.17.2 self-guided projection filter — av1-spec p.330 lines
/// 18218-18273.
///
/// Reads the per-unit selector `set = LrSgrSet[plane][unitRow][unitCol]`
/// and the two projection weights
/// `(w0, w1) = LrSgrXqd[plane][unitRow][unitCol][0..2]`. Runs the
/// §7.17.3 box filter twice — once with `(r0, eps0) =
/// (Sgr_Params[set][0], Sgr_Params[set][1])` (`pass = 0`) and once with
/// `(r1, eps1) = (Sgr_Params[set][2], Sgr_Params[set][3])`
/// (`pass = 1`) — then combines `flt0[i][j]` / `flt1[i][j]` against
/// `u = UpscaledCdefFrame[plane][y + i][x + j] << SGRPROJ_RST_BITS`
/// per the spec's `(w0, w1, w2)` projection:
/// `v = w1 * u + w0 * (flt0 or u) + w2 * (flt1 or u)`,
/// `s = Round2(v, SGRPROJ_RST_BITS + SGRPROJ_PRJ_BITS)`,
/// `LrFrame[plane][y + i][x + j] = Clip1(s)`.
///
/// `w2` is derived as `(1 << SGRPROJ_PRJ_BITS) - w0 - w1` per av1-spec
/// p.330 line 18255.
pub fn self_guided_filter(
    ctx: &LoopRestorationFrameContext<'_>,
    curr_planes: &[PlaneBuffer<'_>],
    cdef_planes: &[PlaneBuffer<'_>],
    lr_planes: &mut [PlaneBuffer<'_>],
    plane: u8,
    geom: &LrBlockGeometry,
) {
    let Some(lr) = lr_planes.get_mut(plane as usize) else {
        return;
    };
    let src = PlaneSource {
        curr: curr_planes,
        cdef: cdef_planes,
        plane,
    };
    self_guided_filter_src(ctx, &src, lr, 0, plane, geom);
}

/// [`self_guided_filter`] from an [`LrSource`] into one plane buffer
/// that may be a row window of the plane (`out_row0` = the plane row
/// at `lr.samples[ 0 ]`).
pub(crate) fn self_guided_filter_src<S: LrSource>(
    ctx: &LoopRestorationFrameContext<'_>,
    src: &S,
    lr: &mut PlaneBuffer<'_>,
    out_row0: i32,
    plane: u8,
    geom: &LrBlockGeometry,
) {
    if geom.w == 0 || geom.h == 0 {
        return;
    }
    // av1-spec p.330 line 18231: set = LrSgrSet[plane][unitRow][unitCol].
    let set = (ctx.lr_sgr_set)(plane, geom.unit_row, geom.unit_col) as usize;
    let set = set.min(SGR_PARAMS.len() - 1);
    let r0 = SGR_PARAMS[set][0];
    let eps0 = SGR_PARAMS[set][1];
    let r1 = SGR_PARAMS[set][2];
    let eps1 = SGR_PARAMS[set][3];
    // av1-spec p.330 lines 18253-18255: w0 / w1 read, w2 = (1 << PRJ) -
    // w0 - w1.
    let w0 = (ctx.lr_sgr_xqd)(plane, geom.unit_row, geom.unit_col, 0);
    let w1 = (ctx.lr_sgr_xqd)(plane, geom.unit_row, geom.unit_col, 1);
    // av1-spec p.330 lines 18233-18241: box_filter twice, with the
    // per-pass `(r, eps)` from Sgr_Params[set].
    let flt0 = box_filter_src(src, geom, ctx.bit_depth, r0, eps0, 0);
    let flt1 = box_filter_src(src, geom, ctx.bit_depth, r1, eps1, 1);
    // av1-spec p.330 lines 18258-18272: per-sample projection.
    let h = geom.h as usize;
    let w = geom.w as usize;
    let lr_rows = lr.rows as i32;
    let lr_cols = lr.cols as i32;
    let lr_cols_usize = lr.cols as usize;
    for i in 0..h {
        for j in 0..w {
            let py = geom.y as i32 + i as i32;
            let px = geom.x as i32 + j as i32;
            // av1-spec p.330 line 18260: u = CDEF[y + i][x + j] << RST_BITS.
            let u = src.cdef_at(px, py) << SGRPROJ_RST_BITS;
            let out = sgr_project(
                u,
                flt0[i * w + j],
                flt1[i * w + j],
                w0,
                w1,
                r0,
                r1,
                ctx.bit_depth,
            );
            let wy = py - out_row0;
            if wy < 0 || wy >= lr_rows || px < 0 || px >= lr_cols {
                continue;
            }
            lr.samples[(wy as usize) * lr_cols_usize + px as usize] = out;
        }
    }
}

/// §7.17.2 per-sample projection — av1-spec p.330 lines 18261-18272:
/// `v = w1·u + w0·(flt0 or u) + w2·(flt1 or u)` (a zero-radius pass
/// substitutes `u`), `s = Round2(v, RST_BITS + PRJ_BITS)`, `Clip1(s)`.
/// `u` is the CDEF sample already shifted by `SGRPROJ_RST_BITS`.
#[inline]
#[must_use]
#[allow(clippy::too_many_arguments)]
pub fn sgr_project(
    u: i32,
    flt0: i32,
    flt1: i32,
    w0: i32,
    w1: i32,
    r0: i32,
    r1: i32,
    bit_depth: u8,
) -> i32 {
    let w2 = (1i32 << SGRPROJ_PRJ_BITS) - w0 - w1;
    let max_sample = (1i32 << bit_depth) - 1;
    let shift = SGRPROJ_RST_BITS + SGRPROJ_PRJ_BITS;
    let mut v = w1 * u;
    v += w0 * if r0 != 0 { flt0 } else { u };
    v += w2 * if r1 != 0 { flt1 } else { u };
    let s = round2_i32(v, shift);
    clip3(0, max_sample, s)
}

/// §7.17.3 box filter — av1-spec p.331-332 lines 18277-18389.
///
/// Returns `F[h][w]` (row-major) for one `pass` of the self-guided
/// filter. `r == 0` returns a zero-filled buffer (the caller
/// substitutes `u` per av1-spec p.330 line 18264 / 18268).
///
/// r464 — the `(2r + 1)²` window sums behind every `A[i][j]` /
/// `B[i][j]` are built from one gathered source window and separable
/// running sums (exact integer arithmetic, so the result is the
/// spec's sample for sample) instead of re-fetching the window per
/// output position.
#[allow(clippy::too_many_arguments)]
pub(crate) fn box_filter(
    curr_planes: &[PlaneBuffer<'_>],
    cdef_planes: &[PlaneBuffer<'_>],
    plane: u8,
    geom: &LrBlockGeometry,
    bit_depth: u8,
    r: i32,
    eps: i32,
    pass: u8,
) -> Vec<i32> {
    let src = PlaneSource {
        curr: curr_planes,
        cdef: cdef_planes,
        plane,
    };
    box_filter_src(&src, geom, bit_depth, r, eps, pass)
}

/// [`box_filter`] from an [`LrSource`].
pub(crate) fn box_filter_src<S: LrSource>(
    src_view: &S,
    geom: &LrBlockGeometry,
    bit_depth: u8,
    r: i32,
    eps: i32,
    pass: u8,
) -> Vec<i32> {
    let h = geom.h as usize;
    let w = geom.w as usize;
    let mut f = vec![0i32; h * w];
    // av1-spec p.331 line 18294: r == 0 ⇒ early-terminate (caller uses u
    // in place of flt, so the zero-filled buffer is a safe sentinel).
    if r == 0 {
        return f;
    }
    let ru = r as usize;
    // av1-spec p.332 lines 18309-18311: n = (2r + 1)²; n2e = n²·eps;
    // s = ((1 << MTABLE) + n2e/2) / n2e.
    let n = (2 * r + 1) * (2 * r + 1);
    let n2e = (n as i64) * (n as i64) * (eps as i64);
    let s_div: i64 = ((1i64 << SGRPROJ_MTABLE_BITS) + n2e / 2) / n2e.max(1);
    // av1-spec p.332 line 18333: oneOverN = ((1 << RECIP_BITS) + n/2) / n.
    let one_over_n = ((1i32 << SGRPROJ_RECIP_BITS) + n / 2) / n;
    // A / B span [-1, h + 1] × [-1, w + 1] per av1-spec p.331 line 18298.
    let aw = w + 2;
    let ah = h + 2;
    let mut a_arr = vec![0i32; ah * aw];
    let mut b_arr = vec![0i32; ah * aw];
    let bd_shift_a = 2 * (bit_depth as i32 - 8).max(0);
    let bd_shift_b = (bit_depth as i32 - 8).max(0);
    // Gather the source window once: rows `-1 - r ..= h + r`, columns
    // `-1 - r ..= w + r` relative to the rectangle origin (§7.17.6
    // routing per sample).
    let sw = aw + 2 * ru;
    let sh = ah + 2 * ru;
    let mut src = vec![0i64; sh * sw];
    let mut row_buf = vec![0i32; sw];
    for (sy, row) in src.chunks_mut(sw).enumerate() {
        let y = geom.y as i32 + sy as i32 - 1 - r;
        src_view.source_row(geom.x as i32 - 1 - r, y, geom, &mut row_buf);
        for (v, &c) in row.iter_mut().zip(row_buf.iter()) {
            *v = i64::from(c);
        }
    }
    // Horizontal running sums over `2r + 1` columns, then vertical
    // over `2r + 1` rows: `b_acc = Σ c`, `a_acc = Σ c²`.
    let win = 2 * ru + 1;
    let mut row_b = vec![0i64; sh * aw];
    let mut row_a = vec![0i64; sh * aw];
    for sy in 0..sh {
        let srow = &src[sy * sw..(sy + 1) * sw];
        let mut sb: i64 = 0;
        let mut sa: i64 = 0;
        for &c in &srow[..win] {
            sb += c;
            sa += c * c;
        }
        row_b[sy * aw] = sb;
        row_a[sy * aw] = sa;
        for j in 1..aw {
            let out = srow[j - 1];
            let inc = srow[j - 1 + win];
            sb += inc - out;
            sa += inc * inc - out * out;
            row_b[sy * aw + j] = sb;
            row_a[sy * aw + j] = sa;
        }
    }
    let mut col_b = vec![0i64; aw];
    let mut col_a = vec![0i64; aw];
    for sy in 0..win {
        for j in 0..aw {
            col_b[j] += row_b[sy * aw + j];
            col_a[j] += row_a[sy * aw + j];
        }
    }
    for ai in 0..ah {
        if ai > 0 {
            let out_row = ai - 1;
            let in_row = ai - 1 + win;
            for j in 0..aw {
                col_b[j] += row_b[in_row * aw + j] - row_b[out_row * aw + j];
                col_a[j] += row_a[in_row * aw + j] - row_a[out_row * aw + j];
            }
        }
        for aj in 0..aw {
            let a_acc = col_a[aj];
            let b_acc = col_b[aj];
            // av1-spec p.332 lines 18323-18324: Round2 by bit-depth shift.
            let a_r = round2_i64(a_acc, bd_shift_a as u32) as i32;
            let d_r = round2_i64(b_acc, bd_shift_b as u32) as i32;
            // av1-spec p.332 line 18325: p = Max(0, a * n - d * d).
            let p = ((a_r as i64) * (n as i64) - (d_r as i64) * (d_r as i64)).max(0);
            // av1-spec p.332 line 18326: z = Round2(p * s, MTABLE_BITS).
            let z = round2_i64(p.saturating_mul(s_div), SGRPROJ_MTABLE_BITS) as i32;
            // av1-spec p.332 lines 18327-18332: piecewise a2.
            let a2 = if z >= 255 {
                256
            } else if z == 0 {
                1
            } else {
                ((z << SGRPROJ_SGR_BITS) + (z / 2)) / (z + 1)
            };
            // av1-spec p.332 line 18334: b2 = ((1 << SGR_BITS) - a2) *
            //                                 b * oneOverN. The §7.17.3
            // spec's `b` is the post-Round2 scalar `d`, not the raw
            // accumulator — we clamp into [-i32::MAX, i32::MAX] for the
            // 10/12-bit profiles before the `i64` multiply.
            let b_eff = b_acc.clamp(i32::MIN as i64, i32::MAX as i64);
            let b2 = i64::from((1i32 << SGRPROJ_SGR_BITS) - a2) * b_eff * i64::from(one_over_n);
            a_arr[ai * aw + aj] = a2;
            b_arr[ai * aw + aj] = round2_i64(b2, SGRPROJ_RECIP_BITS) as i32;
        }
    }
    // av1-spec p.332 lines 18359-18385: F[i][j] from weighted 3×3 of
    // (A, B) and the centre `u = UpscaledCdefFrame[plane][y + i][x + j]`.
    for i in 0..h {
        // av1-spec p.332 lines 18360-18362: shift = 5 by default, 4 on
        // odd rows when pass == 0.
        let shift = if pass == 0 && (i & 1) == 1 {
            4u32
        } else {
            5u32
        };
        let py = geom.y as i32 + i as i32;
        let row_m = (i) * aw;
        let row_c = (i + 1) * aw;
        let row_p = (i + 2) * aw;
        for j in 0..w {
            let (a_sum, b_sum): (i64, i64) = if pass == 0 {
                // Odd-row neighbours only: weights 6 at dx = 0, 5 at
                // dx = ±1, on the rows whose block-local index is odd.
                let mut a: i64 = 0;
                let mut b: i64 = 0;
                for (dy, base) in [(-1i32, row_m), (0, row_c), (1, row_p)] {
                    if ((i as i32 + dy) & 1) != 0 {
                        a += 6 * a_arr[base + j + 1] as i64
                            + 5 * (a_arr[base + j] as i64 + a_arr[base + j + 2] as i64);
                        b += 6 * b_arr[base + j + 1] as i64
                            + 5 * (b_arr[base + j] as i64 + b_arr[base + j + 2] as i64);
                    }
                }
                (a, b)
            } else {
                let cross = |arr: &[i32]| -> i64 {
                    4 * (arr[row_c + j + 1] as i64
                        + arr[row_c + j] as i64
                        + arr[row_c + j + 2] as i64
                        + arr[row_m + j + 1] as i64
                        + arr[row_p + j + 1] as i64)
                        + 3 * (arr[row_m + j] as i64
                            + arr[row_m + j + 2] as i64
                            + arr[row_p + j] as i64
                            + arr[row_p + j + 2] as i64)
                };
                (cross(&a_arr), cross(&b_arr))
            };
            // av1-spec p.332 line 18382: v = a * UpscaledCdefFrame[..] + b.
            let px = geom.x as i32 + j as i32;
            let u = i64::from(src_view.cdef_at(px, py));
            let v = a_sum * u + b_sum;
            // av1-spec p.332 line 18383: F = Round2(v, SGR_BITS + shift -
            //                                            RST_BITS).
            let total_shift = SGRPROJ_SGR_BITS + shift - SGRPROJ_RST_BITS;
            f[i * w + j] = round2_i64(v, total_shift) as i32;
        }
    }
    f
}

// =====================================================================
// §7.17.4 Wiener filter — av1-spec p.333-334 lines 18404-18484.
// =====================================================================

/// §7.17.4 Wiener filter — av1-spec p.333-334 lines 18404-18484.
///
/// Two-pass separable convolution: builds an `intermediate[h+6][w]`
/// array via the 7-tap horizontal filter, then writes
/// `LrFrame[plane][y + r][x + c]` via the 7-tap vertical filter.
/// `(InterRound0, InterRound1)` follow the §7.11.3.2 single-ref
/// schedule (`isCompound = 0`).
pub fn wiener_filter(
    ctx: &LoopRestorationFrameContext<'_>,
    curr_planes: &[PlaneBuffer<'_>],
    cdef_planes: &[PlaneBuffer<'_>],
    lr_planes: &mut [PlaneBuffer<'_>],
    plane: u8,
    geom: &LrBlockGeometry,
) {
    let Some(_curr) = curr_planes.get(plane as usize) else {
        return;
    };
    let Some(_cdef) = cdef_planes.get(plane as usize) else {
        return;
    };
    let Some(lr) = lr_planes.get_mut(plane as usize) else {
        return;
    };
    let src = PlaneSource {
        curr: curr_planes,
        cdef: cdef_planes,
        plane,
    };
    wiener_filter_src(ctx, &src, lr, 0, plane, geom);
}

/// [`wiener_filter`] from an [`LrSource`] into one plane buffer that
/// may be a row window of the plane (`out_row0` = the plane row at
/// `lr.samples[ 0 ]`).
pub(crate) fn wiener_filter_src<S: LrSource>(
    ctx: &LoopRestorationFrameContext<'_>,
    src: &S,
    lr: &mut PlaneBuffer<'_>,
    out_row0: i32,
    plane: u8,
    geom: &LrBlockGeometry,
) {
    // av1-spec p.333 lines 18429-18433: vfilter = wiener_coefficients(
    // LrWiener[plane][unitRow][unitCol][0]); hfilter = wiener_coefficients(
    // LrWiener[plane][unitRow][unitCol][1]).
    let vcoeff: [i32; WIENER_COEFFS] = [
        (ctx.lr_wiener)(plane, geom.unit_row, geom.unit_col, 0, 0),
        (ctx.lr_wiener)(plane, geom.unit_row, geom.unit_col, 0, 1),
        (ctx.lr_wiener)(plane, geom.unit_row, geom.unit_col, 0, 2),
    ];
    let hcoeff: [i32; WIENER_COEFFS] = [
        (ctx.lr_wiener)(plane, geom.unit_row, geom.unit_col, 1, 0),
        (ctx.lr_wiener)(plane, geom.unit_row, geom.unit_col, 1, 1),
        (ctx.lr_wiener)(plane, geom.unit_row, geom.unit_col, 1, 2),
    ];
    let vfilter = wiener_coefficients(&vcoeff);
    let hfilter = wiener_coefficients(&hcoeff);
    // av1-spec p.333 line 18426: §7.11.3.2 isCompound = 0 schedule.
    let inter_round0 = inter_round0(ctx.bit_depth);
    let inter_round1 = inter_round1(ctx.bit_depth);
    let bit_depth = ctx.bit_depth as i32;
    let offset = 1i32 << (bit_depth + FILTER_BITS as i32 - inter_round0 as i32 - 1);
    let limit = (1i32 << (bit_depth + 1 + FILTER_BITS as i32 - inter_round0 as i32)) - 1;
    let h = geom.h as usize;
    let w = geom.w as usize;
    // av1-spec p.333 lines 18447-18455: horizontal pass into
    // `intermediate[h + 6][w]`. r464: the `(h + 6) × (w + 6)` source
    // window is fetched once through §7.17.6 and the 7 taps read it.
    let sw = w + 6;
    let mut window = vec![0i32; (h + 6) * sw];
    for (r, row) in window.chunks_mut(sw).enumerate() {
        let sy = geom.y as i32 + r as i32 - 3;
        src.source_row(geom.x as i32 - 3, sy, geom, row);
    }
    let mut intermediate = vec![0i32; (h + 6) * w];
    for r in 0..(h + 6) {
        let wrow = &window[r * sw..(r + 1) * sw];
        for c in 0..w {
            let mut s = 0i32;
            for (t, &hcoef) in hfilter.iter().enumerate() {
                s += hcoef * wrow[c + t];
            }
            let v = round2_i32(s, inter_round0);
            let clipped = clip3(-offset, limit - offset, v);
            intermediate[r * w + c] = clipped;
        }
    }
    // av1-spec p.334 lines 18476-18484: vertical pass writes
    // `LrFrame[plane][y + r][x + c]`.
    let max_sample = (1i32 << bit_depth) - 1;
    let lr_cols = lr.cols as i32;
    let lr_rows = lr.rows as i32;
    for r in 0..h {
        for c in 0..w {
            let mut s = 0i32;
            for (t, &vcoef) in vfilter.iter().enumerate() {
                s += vcoef * intermediate[(r + t) * w + c];
            }
            let v = round2_i32(s, inter_round1);
            let out = clip3(0, max_sample, v);
            let dst_y = geom.y as i32 + r as i32 - out_row0;
            let dst_x = geom.x as i32 + c as i32;
            if dst_y < 0 || dst_y >= lr_rows || dst_x < 0 || dst_x >= lr_cols {
                continue;
            }
            lr.samples[(dst_y as usize) * (lr.cols as usize) + dst_x as usize] = out;
        }
    }
}

// =====================================================================
// §7.17.5 Wiener coefficient process — av1-spec p.334 lines 18488-18508.
// =====================================================================

/// §7.17.5 Wiener coefficient process — av1-spec p.334 lines
/// 18488-18508.
///
/// Reconstructs the symmetric 7-tap Wiener filter from three
/// transmitted coefficients: `filter[i] = coeff[i]`, `filter[6 - i] =
/// coeff[i]`, `filter[3] = 128 - 2 * Σ coeff[i]`. The DC-gain
/// constraint forces `Σ filter == 128` — the §7.17.4 horizontal pass
/// then rounds by `InterRound0` and the vertical pass rounds by
/// `InterRound1`, with the two together giving unity overall gain.
///
/// Per the spec's `Note`, for chroma planes `coeff[0]` is always 0, so
/// `filter[0]` and `filter[6]` are 0 and the effective filter is a
/// 5-tap convolution.
pub fn wiener_coefficients(coeff: &[i32; WIENER_COEFFS]) -> [i32; 7] {
    let mut filter = [0i32; 7];
    filter[3] = 128;
    for i in 0..3usize {
        let c = coeff[i];
        filter[i] = c;
        filter[6 - i] = c;
        filter[3] -= 2 * c;
    }
    filter
}

// =====================================================================
// r475 — §7.17.6 sources: whole planes, or one stripe in place.
// =====================================================================

/// r475 — where a §7.17 kernel reads its samples from: the §7.17.6
/// `get_source_sample` routing (`UpscaledCurrFrame` at the stripe's
/// 2-row skirts, `UpscaledCdefFrame` inside the stripe) and the
/// in-rectangle `UpscaledCdefFrame` sample the self-guided projection
/// centres on. [`PlaneSource`] is the pair of whole planes;
/// [`StripeSource`] is one stripe's CDEF rows plus the saved pre-CDEF
/// boundary rows — the in-place driver's view.
pub(crate) trait LrSource {
    /// §7.17.6 `get_source_sample( plane, x, y )`.
    fn source(&self, x: i32, y: i32, geom: &LrBlockGeometry) -> i32;
    /// `UpscaledCdefFrame[ plane ][ y ][ x ]` for a position inside
    /// the plane (`0` outside — the kernels never write there).
    fn cdef_at(&self, x: i32, y: i32) -> i32;
    /// `out[ k ] = source( x0 + k, y )`.
    fn source_row(&self, x0: i32, y: i32, geom: &LrBlockGeometry, out: &mut [i32]) {
        for (k, o) in out.iter_mut().enumerate() {
            *o = self.source(x0 + k as i32, y, geom);
        }
    }
}

/// The §7.17 frame driver's view: `UpscaledCurrFrame` /
/// `UpscaledCdefFrame` as whole planes.
pub(crate) struct PlaneSource<'a, 'b> {
    pub curr: &'b [PlaneBuffer<'a>],
    pub cdef: &'b [PlaneBuffer<'a>],
    pub plane: u8,
}

impl LrSource for PlaneSource<'_, '_> {
    #[inline]
    fn source(&self, x: i32, y: i32, geom: &LrBlockGeometry) -> i32 {
        get_source_sample(self.curr, self.cdef, self.plane, x, y, geom)
    }

    #[inline]
    fn cdef_at(&self, x: i32, y: i32) -> i32 {
        let Some(buf) = self.cdef.get(self.plane as usize) else {
            return 0;
        };
        if y < 0 || x < 0 || y >= buf.rows as i32 || x >= buf.cols as i32 {
            return 0;
        }
        buf.samples[(y as usize) * (buf.cols as usize) + x as usize]
    }
}

/// r475 — one stripe's §7.17.6 source for the in-place driver: the
/// stripe's `UpscaledCdefFrame` rows `cdef_row0 ..` (a copy taken
/// before the stripe is overwritten), and the pre-CDEF rows at the
/// stripe's two skirts — `above` holds plane rows `StripeStartY - 2`
/// and `StripeStartY - 1`, `below` rows `StripeEndY + 1` and
/// `StripeEndY + 2`, each captured at its §7.17.6-clamped plane row
/// (`None` where the stripe touches the picture edge, which the
/// clamps make unreachable).
pub(crate) struct StripeSource<'a> {
    pub cdef: &'a [i32],
    pub cdef_row0: i32,
    pub cols: usize,
    pub above: Option<&'a [i32]>,
    pub below: Option<&'a [i32]>,
    pub plane_end_x: i32,
    pub plane_end_y: i32,
}

impl StripeSource<'_> {
    /// The source row for a clamped `y` (§7.17.6 routing) and the
    /// column range `[0, plane_end_x]` it covers.
    #[inline]
    fn row(&self, y: i32, geom: &LrBlockGeometry) -> &[i32] {
        let y = y.clamp(0, self.plane_end_y.max(0));
        if y < geom.stripe_start_y {
            let k = (y.max(geom.stripe_start_y - 2) - (geom.stripe_start_y - 2)) as usize;
            if let Some(a) = self.above {
                return &a[k * self.cols..(k + 1) * self.cols];
            }
        } else if y > geom.stripe_end_y {
            let k = (y.min(geom.stripe_end_y + 2) - (geom.stripe_end_y + 1)) as usize;
            if let Some(b) = self.below {
                return &b[k * self.cols..(k + 1) * self.cols];
            }
        }
        let yy = (y - self.cdef_row0).max(0) as usize;
        let start = (yy * self.cols).min(self.cdef.len().saturating_sub(self.cols));
        &self.cdef[start..start + self.cols]
    }
}

impl LrSource for StripeSource<'_> {
    #[inline]
    fn source(&self, x: i32, y: i32, geom: &LrBlockGeometry) -> i32 {
        let x = x.clamp(0, geom.plane_end_x.max(0)) as usize;
        self.row(y, geom)[x.min(self.cols - 1)]
    }

    #[inline]
    fn cdef_at(&self, x: i32, y: i32) -> i32 {
        if x < 0 || x > self.plane_end_x || y < 0 || y > self.plane_end_y {
            return 0;
        }
        let yy = y - self.cdef_row0;
        if yy < 0 {
            return 0;
        }
        self.cdef
            .get((yy as usize) * self.cols + x as usize)
            .copied()
            .unwrap_or(0)
    }

    fn source_row(&self, x0: i32, y: i32, geom: &LrBlockGeometry, out: &mut [i32]) {
        let row = self.row(y, geom);
        let end_x = geom.plane_end_x.max(0).min(self.cols as i32 - 1);
        for (k, o) in out.iter_mut().enumerate() {
            let x = (x0 + k as i32).clamp(0, end_x) as usize;
            *o = row[x];
        }
    }
}

/// r475 — the pre-CDEF rows §7.17.6 reads across every stripe
/// boundary of one plane: for boundary `s >= 1` at
/// `B = ( 64 s - 8 ) >> subY` (while `B <= PlaneEndY`), the four
/// plane rows `B - 2, B - 1, B, B + 1`, each clamped to
/// `[0, PlaneEndY]` as the fetch clamps them. Captured from the
/// deblocked frame before CDEF runs.
#[derive(Debug, Clone, Default)]
pub(crate) struct StripeBoundaryRows {
    pub cols: usize,
    /// `rows[ s - 1 ]` — four rows of `cols` samples for boundary `s`.
    pub rows: Vec<Vec<i32>>,
}

impl StripeBoundaryRows {
    /// Capture the boundary rows of `plane` (`sub_y` its subsampling).
    pub(crate) fn capture(plane: &PlaneBuffer<'_>, plane_end_y: i32, sub_y: u8) -> Self {
        let cols = plane.cols as usize;
        let mut rows = Vec::new();
        let mut s = 1i32;
        loop {
            let b = (64 * s - 8) >> sub_y;
            if b > plane_end_y {
                break;
            }
            let mut four = Vec::with_capacity(4 * cols);
            for k in -2..2 {
                let y = (b + k).clamp(0, plane_end_y.max(0)) as usize;
                let y = y.min(plane.rows.saturating_sub(1) as usize);
                four.extend_from_slice(&plane.samples[y * cols..(y + 1) * cols]);
            }
            rows.push(four);
            s += 1;
        }
        Self { cols, rows }
    }

    /// The two rows above stripe `s` (`s >= 1`) / below stripe `s`.
    fn above(&self, stripe: usize) -> Option<&[i32]> {
        if stripe == 0 {
            return None;
        }
        self.rows.get(stripe - 1).map(|r| &r[..2 * self.cols])
    }

    fn below(&self, stripe: usize) -> Option<&[i32]> {
        self.rows.get(stripe).map(|r| &r[2 * self.cols..])
    }
}

/// r475 — [`loop_restoration_frame`] in place: `planes` holds
/// `UpscaledCdefFrame` and receives `LrFrame`, one stripe at a time.
/// §7.17.6 reads CDEF samples only inside the unit's own stripe
/// (`StripeStartY ..= StripeEndY`) and pre-CDEF samples only at the
/// stripe's 2-row skirts, and a unit writes only its own rectangle, so
/// a stripe restores from a copy of its CDEF rows plus the captured
/// `boundary` rows — the frame-sized `UpscaledCurrFrame` and
/// `UpscaledCdefFrame` copies (2 × 73 MiB at 12 MP 4:2:0) are gone.
/// Every sample is computed from the same inputs as the frame driver.
pub(crate) fn loop_restoration_frame_in_place(
    ctx: &LoopRestorationFrameContext<'_>,
    planes: &mut [PlaneBuffer<'_>],
    boundary: &[StripeBoundaryRows],
) {
    if !ctx.lr_params.uses_lr {
        return;
    }
    let num_planes = ctx.num_planes.min(planes.len() as u8);
    let mut stripe_copy: Vec<i32> = Vec::new();
    for plane in 0..num_planes {
        let Some(bnd) = boundary.get(plane as usize) else {
            continue;
        };
        let pb = &mut planes[plane as usize];
        for stripe in plane_stripes(ctx, plane) {
            restore_stripe_in_place(ctx, plane, &stripe, bnd, pb, 0, &mut stripe_copy);
        }
    }
}

/// r475 — one stripe's restoration rectangles (`stripe_unit_rects`
/// grouped by `stripe_start_y`): the stripe ordinal, its plane rows
/// `[ y0, y1 )` and the rectangles.
#[derive(Debug, Clone)]
pub(crate) struct PlaneStripe {
    pub stripe: usize,
    pub y0: usize,
    pub y1: usize,
    pub rects: Vec<LrBlockGeometry>,
}

/// r475 — the stripes of `plane` that carry restoration (empty when
/// the plane's `FrameRestorationType` is `RESTORE_NONE`).
pub(crate) fn plane_stripes(ctx: &LoopRestorationFrameContext<'_>, plane: u8) -> Vec<PlaneStripe> {
    let mut out = Vec::new();
    if ctx.lr_params.frame_restoration_type[plane as usize] == FrameRestorationType::None {
        return out;
    }
    let rects = stripe_unit_rects(ctx, plane);
    let mut i = 0usize;
    while i < rects.len() {
        let stripe_start = rects[i].stripe_start_y;
        let mut j = i;
        while j < rects.len() && rects[j].stripe_start_y == stripe_start {
            j += 1;
        }
        let g = &rects[i];
        out.push(PlaneStripe {
            stripe: stripe_index(stripe_start, ctx, plane),
            y0: g.y as usize,
            y1: (g.y + g.h) as usize,
            rects: rects[i..j].to_vec(),
        });
        i = j;
    }
    out
}

/// r475 — restore one stripe in place: `out` holds the stripe's
/// CDEF rows (the whole plane, or the stripe's own row window with
/// `out_row0` its first plane row); the stripe is copied into
/// `stripe_copy` first and every unit reads the copy plus the
/// captured pre-CDEF `boundary` rows.
pub(crate) fn restore_stripe_in_place(
    ctx: &LoopRestorationFrameContext<'_>,
    plane: u8,
    stripe: &PlaneStripe,
    boundary: &StripeBoundaryRows,
    out: &mut PlaneBuffer<'_>,
    out_row0: usize,
    stripe_copy: &mut Vec<i32>,
) {
    let cols = out.cols as usize;
    let (y0, y1) = (stripe.y0, stripe.y1);
    let (w0, w1) = (y0 - out_row0, y1 - out_row0);
    stripe_copy.clear();
    stripe_copy.extend_from_slice(&out.samples[w0 * cols..w1 * cols]);
    let Some(g) = stripe.rects.first() else {
        return;
    };
    let src = StripeSource {
        cdef: stripe_copy,
        cdef_row0: y0 as i32,
        cols,
        above: boundary.above(stripe.stripe),
        below: boundary.below(stripe.stripe),
        plane_end_x: g.plane_end_x,
        plane_end_y: g.plane_end_y,
    };
    for geom in &stripe.rects {
        let r_type = (ctx.lr_type)(plane, geom.unit_row, geom.unit_col);
        match r_type {
            FrameRestorationType::Wiener => {
                wiener_filter_src(ctx, &src, out, out_row0 as i32, plane, geom);
            }
            FrameRestorationType::SgrProj => {
                self_guided_filter_src(ctx, &src, out, out_row0 as i32, plane, geom);
            }
            FrameRestorationType::None | FrameRestorationType::Switchable => {}
        }
    }
}

/// The stripe ordinal `s` of a rectangle's `StripeStartY = ( -8 + 64
/// s ) >> subY`.
fn stripe_index(stripe_start_y: i32, ctx: &LoopRestorationFrameContext<'_>, plane: u8) -> usize {
    let (_, sub_y) = subsampling_for_plane(plane, ctx.subsampling_x, ctx.subsampling_y);
    let mut s = 0i32;
    loop {
        let start = (-8 + s * 64) >> sub_y;
        if start >= stripe_start_y || s > 1 << 20 {
            return s.max(0) as usize;
        }
        s += 1;
    }
}

// =====================================================================
// §7.17.6 Get source sample process — av1-spec p.335 lines 18516-18550.
// =====================================================================

/// §7.17.6 source-sample fetch — av1-spec p.335 lines 18516-18550.
///
/// Snaps `(x, y)` to `[0, PlaneEndX] × [0, PlaneEndY]` and routes to
/// either `UpscaledCurrFrame[plane]` (when `y` lands above
/// `StripeStartY` or below `StripeEndY`, with the cropped
/// `StripeStartY - 2` / `StripeEndY + 2` neighbour-line reach) or
/// `UpscaledCdefFrame[plane]` (when `y` lands inside the stripe).
pub fn get_source_sample(
    curr_planes: &[PlaneBuffer<'_>],
    cdef_planes: &[PlaneBuffer<'_>],
    plane: u8,
    x: i32,
    y: i32,
    geom: &LrBlockGeometry,
) -> i32 {
    // av1-spec p.335 lines 18538-18541: clamp x / y to plane extent.
    let x = x.clamp(0, geom.plane_end_x.max(0));
    let mut y = y.clamp(0, geom.plane_end_y.max(0));
    // av1-spec p.335 lines 18542-18549: stripe-aware fetch routing.
    let p = plane as usize;
    if y < geom.stripe_start_y {
        y = y.max(geom.stripe_start_y - 2);
        return sample_at(curr_planes, p, x, y);
    } else if y > geom.stripe_end_y {
        y = y.min(geom.stripe_end_y + 2);
        return sample_at(curr_planes, p, x, y);
    }
    sample_at(cdef_planes, p, x, y)
}

/// Bounds-checked sample fetch from a [`PlaneBuffer`] used by
/// [`get_source_sample`]. The §7.17.6 clamping already pegs `(x, y)`
/// inside `[0, PlaneEndX] × [0, PlaneEndY]` and the stripe-boundary
/// fetches into `UpscaledCurrFrame` stay within
/// `[StripeStartY - 2, StripeEndY + 2]` — but the planes the caller
/// passes in may be sized differently, so the fetch here clamps
/// defensively to the buffer's `[0, rows-1] × [0, cols-1]` extent.
fn sample_at(planes: &[PlaneBuffer<'_>], plane: usize, x: i32, y: i32) -> i32 {
    let Some(buf) = planes.get(plane) else {
        return 0;
    };
    if buf.rows == 0 || buf.cols == 0 {
        return 0;
    }
    let yc = y.clamp(0, buf.rows as i32 - 1) as usize;
    let xc = x.clamp(0, buf.cols as i32 - 1) as usize;
    buf.samples[yc * (buf.cols as usize) + xc]
}

// =====================================================================
// Helpers — §3 / §5.9.20 numeric primitives.
// =====================================================================

/// §5.9.20 `count_units_in_frame(unitSize, frameSize)` per av1-spec
/// p.107 line 6406 — `Max((frameSize + (unitSize >> 1)) / unitSize, 1)`.
pub fn count_units_in_frame(unit_size: u32, frame_size: u32) -> u32 {
    if unit_size == 0 {
        return 1;
    }
    ((frame_size + (unit_size >> 1)) / unit_size).max(1)
}

/// §3 `Round2(x, n)` — `(x + (1 << (n - 1))) >> n` for `n > 0`,
/// identity for `n == 0`.
#[inline]
fn round2(x: u32, n: u8) -> u32 {
    if n == 0 {
        x
    } else {
        (x + (1u32 << (n - 1))) >> n
    }
}

/// `Round2(x, n)` over `i32` — same shape as [`round2`] but signed.
#[inline]
fn round2_i32(x: i32, n: u32) -> i32 {
    if n == 0 {
        x
    } else {
        (x + (1i32 << (n - 1))) >> n
    }
}

/// `Round2(x, n)` over `i64` — needed by the §7.17.3 box filter where
/// `a * n - d * d` and `b2` intermediates exceed the `i32` range at
/// the 10/12-bit profiles.
#[inline]
fn round2_i64(x: i64, n: u32) -> i64 {
    if n == 0 {
        x
    } else {
        (x + (1i64 << (n - 1))) >> n
    }
}

/// §3 `Clip3(a, b, x)` — clamp `x` to `[a, b]`.
#[inline]
fn clip3(a: i32, b: i32, x: i32) -> i32 {
    if x < a {
        a
    } else if x > b {
        b
    } else {
        x
    }
}

/// `(subX, subY)` for `plane` per av1-spec p.328 lines 18138-18142.
#[inline]
fn subsampling_for_plane(plane: u8, sub_x: u8, sub_y: u8) -> (u8, u8) {
    if plane == 0 {
        (0, 0)
    } else {
        (sub_x, sub_y)
    }
}

/// §7.11.3.2 `InterRound0` for `isCompound = 0` per av1-spec p.259
/// lines 14423 / 14427 — `3` for `BitDepth ∈ {8, 10}`, `5` for `12`.
#[inline]
fn inter_round0(bit_depth: u8) -> u32 {
    if bit_depth == 12 {
        5
    } else {
        3
    }
}

/// §7.11.3.2 `InterRound1` for `isCompound = 0` per av1-spec p.259
/// lines 14425 / 14429 — `11` for `BitDepth ∈ {8, 10}`, `9` for `12`.
#[inline]
fn inter_round1(bit_depth: u8) -> u32 {
    if bit_depth == 12 {
        9
    } else {
        11
    }
}

// =====================================================================
// Tests.
// =====================================================================

#[cfg(test)]
mod tests {
    use super::*;
    use crate::uncompressed_header_tail::RESTORATION_TILESIZE_MAX;

    fn make_lr_params(uses_lr: bool, sizes: [u32; 3]) -> LrParams {
        LrParams {
            frame_restoration_type: [FrameRestorationType::None; 3],
            uses_lr,
            uses_chroma_lr: false,
            lr_unit_shift: 0,
            lr_uv_shift: 0,
            loop_restoration_size: sizes,
            short_circuited: false,
        }
    }

    fn make_plane(rows: u32, cols: u32, fill: i32) -> Vec<i32> {
        vec![fill; (rows * cols) as usize]
    }

    #[test]
    fn wiener_taps_tables_match_spec() {
        // av1-spec p.74 lines 3884-3886 + p.107 lines 6484-6489.
        assert_eq!(WIENER_TAPS_MID, [3, -7, 15]);
        assert_eq!(WIENER_TAPS_MIN, [-5, -23, -17]);
        assert_eq!(WIENER_TAPS_MAX, [10, 8, 46]);
        assert_eq!(WIENER_TAPS_K, [1, 2, 3]);
        assert_eq!(SGRPROJ_XQD_MID, [-32, 31]);
        assert_eq!(SGRPROJ_XQD_MIN, [-96, -32]);
        assert_eq!(SGRPROJ_XQD_MAX, [31, 95]);
    }

    #[test]
    fn sgr_params_table_matches_spec() {
        // av1-spec p.332 lines 18395-18400 — spot-check rows 0, 10, 14.
        assert_eq!(SGR_PARAMS[0], [2, 12, 1, 4]);
        assert_eq!(SGR_PARAMS[10], [0, 0, 1, 5]);
        assert_eq!(SGR_PARAMS[14], [2, 30, 0, 0]);
        assert_eq!(SGR_PARAMS.len(), 1 << SGRPROJ_PARAMS_BITS);
    }

    #[test]
    fn wiener_coefficients_unit_dc_gain() {
        // av1-spec p.334 line 18499: filter[3] = 128, filter[i] =
        // coeff[i], filter[6-i] = coeff[i], filter[3] -= 2 * Σ coeff.
        // The symmetric DC gain Σ filter == 128 across all inputs.
        for coeff in [
            [0, 0, 0],
            [3, -7, 15],    // Wiener_Taps_Mid
            [-5, -23, -17], // Wiener_Taps_Min
            [10, 8, 46],    // Wiener_Taps_Max
        ] {
            let f = wiener_coefficients(&coeff);
            assert_eq!(f.iter().sum::<i32>(), 128);
            // Symmetry: filter[i] == filter[6 - i].
            assert_eq!(f[0], f[6]);
            assert_eq!(f[1], f[5]);
            assert_eq!(f[2], f[4]);
            // Transmitted taps land in slots 0..=2.
            assert_eq!(f[0], coeff[0]);
            assert_eq!(f[1], coeff[1]);
            assert_eq!(f[2], coeff[2]);
        }
    }

    #[test]
    fn wiener_coefficients_zero_centre_is_128() {
        // All-zero transmitted taps ⇒ filter = [0, 0, 0, 128, 0, 0, 0].
        let f = wiener_coefficients(&[0, 0, 0]);
        assert_eq!(f, [0, 0, 0, 128, 0, 0, 0]);
    }

    #[test]
    fn count_units_in_frame_matches_spec() {
        // av1-spec p.107 line 6406:
        //   count_units_in_frame(unitSize, frameSize) =
        //     Max((frameSize + (unitSize >> 1)) / unitSize, 1)
        assert_eq!(count_units_in_frame(64, 256), 4);
        assert_eq!(count_units_in_frame(64, 100), 2); // (100 + 32) / 64 = 2
        assert_eq!(count_units_in_frame(64, 0), 1); // floor at 1
        assert_eq!(count_units_in_frame(256, 240), 1); // (240 + 128) / 256 = 1
    }

    #[test]
    fn loop_restoration_frame_no_op_when_uses_lr_false() {
        // UsesLr == 0 ⇒ the driver just copies UpscaledCdefFrame into
        // LrFrame and returns. Per-plane closures are never consulted.
        let mut curr = make_plane(16, 16, 50);
        let mut cdef = make_plane(16, 16, 200);
        let mut lr = make_plane(16, 16, 0);
        let curr_buf = PlaneBuffer {
            rows: 16,
            cols: 16,
            samples: &mut curr,
        };
        let cdef_buf = PlaneBuffer {
            rows: 16,
            cols: 16,
            samples: &mut cdef,
        };
        let mut lr_buf = PlaneBuffer {
            rows: 16,
            cols: 16,
            samples: &mut lr,
        };
        let params = make_lr_params(false, [64; 3]);
        let ctx = LoopRestorationFrameContext {
            mi_rows: 4,
            mi_cols: 4,
            num_planes: 1,
            bit_depth: 8,
            subsampling_x: 0,
            subsampling_y: 0,
            frame_height: 16,
            upscaled_width: 16,
            lr_params: &params,
            lr_type: &|_, _, _| FrameRestorationType::None,
            lr_wiener: &|_, _, _, _, _| 0,
            lr_sgr_set: &|_, _, _| 0,
            lr_sgr_xqd: &|_, _, _, _| 0,
        };
        loop_restoration_frame(
            &ctx,
            std::slice::from_ref(&curr_buf),
            std::slice::from_ref(&cdef_buf),
            std::slice::from_mut(&mut lr_buf),
        );
        // LrFrame holds a copy of UpscaledCdefFrame.
        assert!(lr_buf.samples.iter().all(|&v| v == 200));
    }

    #[test]
    fn loop_restoration_frame_restore_none_keeps_cdef_copy() {
        // Per-plane FrameRestorationType == RESTORE_NONE ⇒ LrFrame
        // holds UpscaledCdefFrame, even when UsesLr == 1 (e.g. another
        // plane has loop restoration enabled).
        let mut curr = make_plane(16, 16, 50);
        let mut cdef = make_plane(16, 16, 175);
        let mut lr = make_plane(16, 16, 0);
        let curr_buf = PlaneBuffer {
            rows: 16,
            cols: 16,
            samples: &mut curr,
        };
        let cdef_buf = PlaneBuffer {
            rows: 16,
            cols: 16,
            samples: &mut cdef,
        };
        let mut lr_buf = PlaneBuffer {
            rows: 16,
            cols: 16,
            samples: &mut lr,
        };
        let mut params = make_lr_params(true, [64; 3]);
        params.frame_restoration_type[0] = FrameRestorationType::None;
        let ctx = LoopRestorationFrameContext {
            mi_rows: 4,
            mi_cols: 4,
            num_planes: 1,
            bit_depth: 8,
            subsampling_x: 0,
            subsampling_y: 0,
            frame_height: 16,
            upscaled_width: 16,
            lr_params: &params,
            lr_type: &|_, _, _| FrameRestorationType::None,
            lr_wiener: &|_, _, _, _, _| 0,
            lr_sgr_set: &|_, _, _| 0,
            lr_sgr_xqd: &|_, _, _, _| 0,
        };
        loop_restoration_frame(
            &ctx,
            std::slice::from_ref(&curr_buf),
            std::slice::from_ref(&cdef_buf),
            std::slice::from_mut(&mut lr_buf),
        );
        assert!(lr_buf.samples.iter().all(|&v| v == 175));
    }

    #[test]
    fn wiener_filter_identity_filter_recovers_source() {
        // With coeff = [0, 0, 0], wiener_coefficients returns
        // [0, 0, 0, 128, 0, 0, 0] — i.e. the identity 7-tap filter
        // scaled by 128. Two-pass: horizontal then vertical, with
        // InterRound0=3 and InterRound1=11 for BitDepth=8. The combined
        // gain is 128 * 128 >> 14 = 1, so a uniform input recovers as
        // itself.
        let mut curr = make_plane(64, 64, 100);
        let mut cdef = make_plane(64, 64, 100);
        let mut lr = make_plane(64, 64, 0);
        let curr_buf = PlaneBuffer {
            rows: 64,
            cols: 64,
            samples: &mut curr,
        };
        let cdef_buf = PlaneBuffer {
            rows: 64,
            cols: 64,
            samples: &mut cdef,
        };
        let mut lr_buf = PlaneBuffer {
            rows: 64,
            cols: 64,
            samples: &mut lr,
        };
        let mut params = make_lr_params(true, [RESTORATION_TILESIZE_MAX; 3]);
        params.frame_restoration_type[0] = FrameRestorationType::Wiener;
        let ctx = LoopRestorationFrameContext {
            mi_rows: 16,
            mi_cols: 16,
            num_planes: 1,
            bit_depth: 8,
            subsampling_x: 0,
            subsampling_y: 0,
            frame_height: 64,
            upscaled_width: 64,
            lr_params: &params,
            lr_type: &|_, _, _| FrameRestorationType::Wiener,
            lr_wiener: &|_, _, _, _, _| 0,
            lr_sgr_set: &|_, _, _| 0,
            lr_sgr_xqd: &|_, _, _, _| 0,
        };
        loop_restoration_frame(
            &ctx,
            std::slice::from_ref(&curr_buf),
            std::slice::from_ref(&cdef_buf),
            std::slice::from_mut(&mut lr_buf),
        );
        // The Wiener identity convolution over a uniform plane yields
        // the same uniform plane back. (Rounding constants cancel: a
        // uniform input means every horizontal sum is 128 * src; the
        // vertical pass scales that by 128 / 2^11 ≈ src after the
        // intermediate `>> 3` rounding.)
        for &v in lr_buf.samples.iter() {
            assert_eq!(v, 100);
        }
    }

    #[test]
    fn get_source_sample_inside_stripe_reads_cdef() {
        // y ∈ [StripeStartY, StripeEndY] ⇒ UpscaledCdefFrame.
        let mut curr = make_plane(16, 16, 11);
        let mut cdef = make_plane(16, 16, 99);
        let curr_buf = PlaneBuffer {
            rows: 16,
            cols: 16,
            samples: &mut curr,
        };
        let cdef_buf = PlaneBuffer {
            rows: 16,
            cols: 16,
            samples: &mut cdef,
        };
        let geom = LrBlockGeometry {
            unit_row: 0,
            unit_col: 0,
            x: 0,
            y: 4,
            w: 4,
            h: 4,
            stripe_start_y: 0,
            stripe_end_y: 8,
            plane_end_x: 15,
            plane_end_y: 15,
        };
        let v = get_source_sample(
            std::slice::from_ref(&curr_buf),
            std::slice::from_ref(&cdef_buf),
            0,
            3,
            5,
            &geom,
        );
        assert_eq!(v, 99);
    }

    #[test]
    fn get_source_sample_above_stripe_reads_curr_clamped() {
        // y < StripeStartY ⇒ UpscaledCurrFrame, with y pegged at
        // max(StripeStartY - 2, y).
        let mut curr = make_plane(16, 16, 33);
        let mut cdef = make_plane(16, 16, 77);
        let curr_buf = PlaneBuffer {
            rows: 16,
            cols: 16,
            samples: &mut curr,
        };
        let cdef_buf = PlaneBuffer {
            rows: 16,
            cols: 16,
            samples: &mut cdef,
        };
        let geom = LrBlockGeometry {
            unit_row: 0,
            unit_col: 0,
            x: 0,
            y: 8,
            w: 4,
            h: 4,
            stripe_start_y: 6,
            stripe_end_y: 14,
            plane_end_x: 15,
            plane_end_y: 15,
        };
        // y = 2 < stripe_start_y = 6; clamp pegs to max(6 - 2, 2) = 4.
        let v = get_source_sample(
            std::slice::from_ref(&curr_buf),
            std::slice::from_ref(&cdef_buf),
            0,
            3,
            2,
            &geom,
        );
        assert_eq!(v, 33);
    }

    #[test]
    fn get_source_sample_below_stripe_reads_curr() {
        // y > StripeEndY ⇒ UpscaledCurrFrame.
        let mut curr = make_plane(32, 16, 42);
        let mut cdef = make_plane(32, 16, 88);
        let curr_buf = PlaneBuffer {
            rows: 32,
            cols: 16,
            samples: &mut curr,
        };
        let cdef_buf = PlaneBuffer {
            rows: 32,
            cols: 16,
            samples: &mut cdef,
        };
        let geom = LrBlockGeometry {
            unit_row: 0,
            unit_col: 0,
            x: 0,
            y: 0,
            w: 4,
            h: 4,
            stripe_start_y: 0,
            stripe_end_y: 8,
            plane_end_x: 15,
            plane_end_y: 31,
        };
        // y = 12 > stripe_end_y = 8; clamp pegs to min(8 + 2, 12) = 10.
        let v = get_source_sample(
            std::slice::from_ref(&curr_buf),
            std::slice::from_ref(&cdef_buf),
            0,
            3,
            12,
            &geom,
        );
        assert_eq!(v, 42);
    }

    #[test]
    fn derive_block_geometry_luma_top_left() {
        let params = make_lr_params(true, [RESTORATION_TILESIZE_MAX; 3]);
        let ctx = LoopRestorationFrameContext {
            mi_rows: 16,
            mi_cols: 16,
            num_planes: 1,
            bit_depth: 8,
            subsampling_x: 0,
            subsampling_y: 0,
            frame_height: 64,
            upscaled_width: 64,
            lr_params: &params,
            lr_type: &|_, _, _| FrameRestorationType::None,
            lr_wiener: &|_, _, _, _, _| 0,
            lr_sgr_set: &|_, _, _| 0,
            lr_sgr_xqd: &|_, _, _, _| 0,
        };
        let geom = derive_block_geometry(&ctx, 0, 0, 0);
        assert_eq!(geom.unit_row, 0);
        assert_eq!(geom.unit_col, 0);
        assert_eq!(geom.x, 0);
        assert_eq!(geom.y, 0);
        assert_eq!(geom.w, MI_SIZE as u32);
        assert_eq!(geom.h, MI_SIZE as u32);
        assert_eq!(geom.plane_end_x, 63);
        assert_eq!(geom.plane_end_y, 63);
        // stripeNum = (0 + 8) / 64 = 0 ⇒ stripe_start_y = -8.
        assert_eq!(geom.stripe_start_y, -8);
        // stripe_end_y = -8 + 64 - 1 = 55.
        assert_eq!(geom.stripe_end_y, 55);
    }

    #[test]
    fn self_guided_uniform_zero_weights_recovers_source() {
        // Per av1-spec p.330 lines 18260-18271: when w0 == w1 == 0 (so
        // w2 = 128) and both r0 / r1 are non-zero, the projection is
        // v = 128 * flt1. On a uniform plane the box-filter output
        // closely tracks `u = cdef << RST_BITS` (variance term zeroes
        // to a == 1, so the convolution is dominated by `b ≈ u`), and
        // `Round2(128 * flt1, 11)` recovers the original sample (modulo
        // sub-LSB rounding noise — accepted within ±1 here).
        let mut curr = make_plane(16, 16, 50);
        let mut cdef = make_plane(16, 16, 222);
        let mut lr = make_plane(16, 16, 0);
        let curr_buf = PlaneBuffer {
            rows: 16,
            cols: 16,
            samples: &mut curr,
        };
        let cdef_buf = PlaneBuffer {
            rows: 16,
            cols: 16,
            samples: &mut cdef,
        };
        let mut lr_buf = PlaneBuffer {
            rows: 16,
            cols: 16,
            samples: &mut lr,
        };
        let mut params = make_lr_params(true, [RESTORATION_TILESIZE_MAX; 3]);
        params.frame_restoration_type[0] = FrameRestorationType::SgrProj;
        let ctx = LoopRestorationFrameContext {
            mi_rows: 4,
            mi_cols: 4,
            num_planes: 1,
            bit_depth: 8,
            subsampling_x: 0,
            subsampling_y: 0,
            frame_height: 16,
            upscaled_width: 16,
            lr_params: &params,
            lr_type: &|_, _, _| FrameRestorationType::SgrProj,
            lr_wiener: &|_, _, _, _, _| 0,
            // set 0 ⇒ (r0, eps0, r1, eps1) = (2, 12, 1, 4).
            lr_sgr_set: &|_, _, _| 0,
            // w0 = w1 = 0 ⇒ w2 = 128.
            lr_sgr_xqd: &|_, _, _, _| 0,
        };
        loop_restoration_frame(
            &ctx,
            std::slice::from_ref(&curr_buf),
            std::slice::from_ref(&cdef_buf),
            std::slice::from_mut(&mut lr_buf),
        );
        // Spot-check the deep interior so we avoid the boundary-fetch
        // path where the box-sum sees clipped neighbours.
        for y in 4..12 {
            for x in 4..12 {
                let v = lr_buf.samples[y * 16 + x];
                assert!((v - 222).abs() <= 1, "lr[{y},{x}] = {v} (expected ≈ 222)");
            }
        }
    }

    #[test]
    fn self_guided_w2_zero_recovers_source_directly() {
        // w0 = 0, w1 = 128 ⇒ w2 = 0; projection reduces to
        // v = 128 * u = 128 * (cdef << 4). Round2(., 11) collapses to
        // (cdef << 4) >> 4 = cdef. Recovers the source exactly without
        // running through the box filter's nonlinearity.
        let mut curr = make_plane(16, 16, 0);
        let mut cdef = make_plane(16, 16, 137);
        let mut lr = make_plane(16, 16, 0);
        let curr_buf = PlaneBuffer {
            rows: 16,
            cols: 16,
            samples: &mut curr,
        };
        let cdef_buf = PlaneBuffer {
            rows: 16,
            cols: 16,
            samples: &mut cdef,
        };
        let mut lr_buf = PlaneBuffer {
            rows: 16,
            cols: 16,
            samples: &mut lr,
        };
        let mut params = make_lr_params(true, [RESTORATION_TILESIZE_MAX; 3]);
        params.frame_restoration_type[0] = FrameRestorationType::SgrProj;
        let ctx = LoopRestorationFrameContext {
            mi_rows: 4,
            mi_cols: 4,
            num_planes: 1,
            bit_depth: 8,
            subsampling_x: 0,
            subsampling_y: 0,
            frame_height: 16,
            upscaled_width: 16,
            lr_params: &params,
            lr_type: &|_, _, _| FrameRestorationType::SgrProj,
            lr_wiener: &|_, _, _, _, _| 0,
            lr_sgr_set: &|_, _, _| 0,
            lr_sgr_xqd: &|_, _, _, i| if i == 1 { 128 } else { 0 },
        };
        loop_restoration_frame(
            &ctx,
            std::slice::from_ref(&curr_buf),
            std::slice::from_ref(&cdef_buf),
            std::slice::from_mut(&mut lr_buf),
        );
        for &v in lr_buf.samples.iter() {
            assert_eq!(v, 137);
        }
    }

    #[test]
    fn self_guided_set_10_r0_zero_uses_u_for_flt0() {
        // SGR_PARAMS[10] = [0, 0, 1, 5] ⇒ r0 == 0. Per av1-spec p.330
        // line 18263-18264, the projection substitutes `u` for `flt0`.
        // With w0 = 128, w1 = 0 (⇒ w2 = 0), the projection becomes
        // v = 0 * u + 128 * u + 0 * flt1 = 128 * u, recovering `cdef`.
        let mut curr = make_plane(16, 16, 0);
        let mut cdef = make_plane(16, 16, 200);
        let mut lr = make_plane(16, 16, 0);
        let curr_buf = PlaneBuffer {
            rows: 16,
            cols: 16,
            samples: &mut curr,
        };
        let cdef_buf = PlaneBuffer {
            rows: 16,
            cols: 16,
            samples: &mut cdef,
        };
        let mut lr_buf = PlaneBuffer {
            rows: 16,
            cols: 16,
            samples: &mut lr,
        };
        let mut params = make_lr_params(true, [RESTORATION_TILESIZE_MAX; 3]);
        params.frame_restoration_type[0] = FrameRestorationType::SgrProj;
        let ctx = LoopRestorationFrameContext {
            mi_rows: 4,
            mi_cols: 4,
            num_planes: 1,
            bit_depth: 8,
            subsampling_x: 0,
            subsampling_y: 0,
            frame_height: 16,
            upscaled_width: 16,
            lr_params: &params,
            lr_type: &|_, _, _| FrameRestorationType::SgrProj,
            lr_wiener: &|_, _, _, _, _| 0,
            // SGR_PARAMS[10] = (r0=0, eps0=0, r1=1, eps1=5).
            lr_sgr_set: &|_, _, _| 10,
            lr_sgr_xqd: &|_, _, _, i| if i == 0 { 128 } else { 0 },
        };
        loop_restoration_frame(
            &ctx,
            std::slice::from_ref(&curr_buf),
            std::slice::from_ref(&cdef_buf),
            std::slice::from_mut(&mut lr_buf),
        );
        for &v in lr_buf.samples.iter() {
            assert_eq!(v, 200);
        }
    }

    #[test]
    fn self_guided_set_14_r1_zero_uses_u_for_flt1() {
        // SGR_PARAMS[14] = [2, 30, 0, 0] ⇒ r1 == 0. Symmetric to the
        // r0 == 0 case: the projection substitutes `u` for `flt1`.
        // With w0 = 0, w1 = 128 (⇒ w2 = 0), v = 128 * u recovers cdef.
        let mut curr = make_plane(16, 16, 0);
        let mut cdef = make_plane(16, 16, 88);
        let mut lr = make_plane(16, 16, 0);
        let curr_buf = PlaneBuffer {
            rows: 16,
            cols: 16,
            samples: &mut curr,
        };
        let cdef_buf = PlaneBuffer {
            rows: 16,
            cols: 16,
            samples: &mut cdef,
        };
        let mut lr_buf = PlaneBuffer {
            rows: 16,
            cols: 16,
            samples: &mut lr,
        };
        let mut params = make_lr_params(true, [RESTORATION_TILESIZE_MAX; 3]);
        params.frame_restoration_type[0] = FrameRestorationType::SgrProj;
        let ctx = LoopRestorationFrameContext {
            mi_rows: 4,
            mi_cols: 4,
            num_planes: 1,
            bit_depth: 8,
            subsampling_x: 0,
            subsampling_y: 0,
            frame_height: 16,
            upscaled_width: 16,
            lr_params: &params,
            lr_type: &|_, _, _| FrameRestorationType::SgrProj,
            lr_wiener: &|_, _, _, _, _| 0,
            // SGR_PARAMS[14] = (r0=2, eps0=30, r1=0, eps1=0).
            lr_sgr_set: &|_, _, _| 14,
            lr_sgr_xqd: &|_, _, _, i| if i == 1 { 128 } else { 0 },
        };
        loop_restoration_frame(
            &ctx,
            std::slice::from_ref(&curr_buf),
            std::slice::from_ref(&cdef_buf),
            std::slice::from_mut(&mut lr_buf),
        );
        for &v in lr_buf.samples.iter() {
            assert_eq!(v, 88);
        }
    }

    #[test]
    fn box_filter_r0_returns_zeroed_buffer() {
        // av1-spec p.331 line 18294: r == 0 ⇒ box_filter terminates
        // immediately. We surface that by returning a zero-filled
        // buffer of the requested (h * w) extent — the §7.17.2 caller
        // substitutes `u` for `flt` when r == 0, so the zeros are
        // never read on the projection path.
        let mut curr = make_plane(8, 8, 0);
        let mut cdef = make_plane(8, 8, 100);
        let curr_buf = PlaneBuffer {
            rows: 8,
            cols: 8,
            samples: &mut curr,
        };
        let cdef_buf = PlaneBuffer {
            rows: 8,
            cols: 8,
            samples: &mut cdef,
        };
        let geom = LrBlockGeometry {
            unit_row: 0,
            unit_col: 0,
            x: 0,
            y: 0,
            w: 4,
            h: 4,
            stripe_start_y: -8,
            stripe_end_y: 55,
            plane_end_x: 7,
            plane_end_y: 7,
        };
        let out = box_filter(
            std::slice::from_ref(&curr_buf),
            std::slice::from_ref(&cdef_buf),
            0,
            &geom,
            8,
            0,
            0,
            0,
        );
        assert_eq!(out.len(), 16);
        assert!(out.iter().all(|&v| v == 0));
    }

    #[test]
    fn self_guided_clips_to_sample_range() {
        // The projection result is `Clip1(s)` per av1-spec p.330 line
        // 18271 — for 8-bit decode that's `[0, 255]`. Set up a
        // pathological w1 = 95 (max) + w0 = -96 (min) ⇒ w2 = 129. The
        // projection should still emerge inside `[0, 255]` for any
        // uniform input.
        let mut curr = make_plane(16, 16, 0);
        let mut cdef = make_plane(16, 16, 128);
        let mut lr = make_plane(16, 16, 0);
        let curr_buf = PlaneBuffer {
            rows: 16,
            cols: 16,
            samples: &mut curr,
        };
        let cdef_buf = PlaneBuffer {
            rows: 16,
            cols: 16,
            samples: &mut cdef,
        };
        let mut lr_buf = PlaneBuffer {
            rows: 16,
            cols: 16,
            samples: &mut lr,
        };
        let mut params = make_lr_params(true, [RESTORATION_TILESIZE_MAX; 3]);
        params.frame_restoration_type[0] = FrameRestorationType::SgrProj;
        let ctx = LoopRestorationFrameContext {
            mi_rows: 4,
            mi_cols: 4,
            num_planes: 1,
            bit_depth: 8,
            subsampling_x: 0,
            subsampling_y: 0,
            frame_height: 16,
            upscaled_width: 16,
            lr_params: &params,
            lr_type: &|_, _, _| FrameRestorationType::SgrProj,
            lr_wiener: &|_, _, _, _, _| 0,
            lr_sgr_set: &|_, _, _| 0,
            lr_sgr_xqd: &|_, _, _, i| if i == 0 { -96 } else { 95 },
        };
        loop_restoration_frame(
            &ctx,
            std::slice::from_ref(&curr_buf),
            std::slice::from_ref(&cdef_buf),
            std::slice::from_mut(&mut lr_buf),
        );
        for &v in lr_buf.samples.iter() {
            assert!((0..=255).contains(&v), "lr sample {v} escapes [0, 255]");
        }
    }
}