Skip to main content

par2_rs/
repair_transform.rs

1//! The transform arm of PAR2 repair: syndromes by DFT, solve by the m×m inverse.
2//!
3//! # Why
4//!
5//! The dense repair executor folds every available source into every missing
6//! output: `m * n_available` region folds, plus `m * m` for the recovery
7//! columns. At a large set that product is the whole runtime — 4096 missing
8//! slices against 14000 available ones is 57 million folds per stripe of the
9//! slice.
10//!
11//! The same answer can be reached in two cheaper steps.
12//!
13//! **Syndromes.** For a selected recovery exponent `e`,
14//!
15//! ```text
16//! S_e = R_e ^ sum over present i of D_i * c_i^e
17//! ```
18//!
19//! and that sum is exactly one output row of the multiplicative GF(2^16) DFT
20//! over the present slices' PAR2 constants. [`DftPlan`] computes a contiguous
21//! range of such rows in `region_folds()` folds instead of `m * n_available`,
22//! which at the shape above is a ~35x reduction. `S_e` is then the recovery
23//! row `R_e` with the present contributions removed, i.e. the parity of the
24//! *missing* slices alone.
25//!
26//! **Solve.** Those syndromes are `A * X`, where `A[r][c] = c_missing_c ^ e_r`
27//! is the m×m Vandermonde block over the missing slots and `X` the missing
28//! slices. [`crate::matrix`] already inverts `A` while planning the repair —
29//! `RepairPlan::decode_matrix` *is* `A^-1` — so the solve is `m * m` folds of
30//! that inverse against the syndrome rows.
31//!
32//! Total: `region_folds + m^2` against `m * (n_available + m)`. GF(2^16)
33//! addition is XOR and multiplication is exact, so re-associating the sum this
34//! way is bit-identical to the dense product, not merely equivalent. The dense
35//! path computes `A^-1 * (pre * present + I * R)`; this computes
36//! `A^-1 * (R + pre * present)`. Same terms, same field, different order.
37//!
38//! (The transform itself is a Good–Thomas factorisation of the length-65535
39//! cyclic DFT; its derivation, cost model and pruning live in
40//! [`reedsolomon_rs::gf16_dft`]. The idea of running PAR2 repair through a
41//! syndrome transform rather than a dense matrix product is not new — other
42//! Reed-Solomon implementations do it — but the schedule, the memory contract
43//! and the divergence probe here are this crate's own.)
44//!
45//! # What it costs in memory, and why that decides everything
46//!
47//! The dense executor streams one source at a time: two transfer buffers and a
48//! staging area, independent of `n_available`. The transform cannot. A DFT row
49//! is a sum over *every* present slice, so the same byte range of every present
50//! slice has to be resident at once.
51//!
52//! So the arm works in **bands**. A band is a byte range of each slice; one
53//! pass stages that range of every present slice, produces every missing
54//! slice's bytes for that range, writes them, and moves on. Within a band the
55//! work splits into **stripes**, which are what the transform and the solve
56//! actually consume and what the rayon workers take one at a time. A band is
57//! always a whole number of stripes, and both are 64-byte aligned.
58//!
59//! Every arena is summed before a byte is allocated, and the whole sum is
60//! charged against the caller's existing [`crate::repair::RepairOptions`]
61//! `memory_limit`:
62//!
63//! ```text
64//! (n_present + m + 2) * band                     staging, output rows, probe
65//! + workers * (range_len * stripe + dft scratch) per-worker syndrome rows
66//! + one spare dft scratch                        the short tail stripe
67//! + plan tables
68//! ```
69//!
70//! If the limit cannot buy a band of at least [`MIN_BAND`] bytes, or the slice
71//! would need more than [`MAX_PASSES`] bands, the arm declines and the dense
72//! executor runs unchanged. The decode matrix is *not* charged here: it is
73//! built by `plan_repair` under the separate matrix-workspace budget
74//! (`MATRIX_WORKSPACE_BUDGET_FLOOR`), exactly as it is for the dense path, and
75//! charging it twice would only make the arm decline where the dense path
76//! happily proceeds.
77//!
78//! # Safety
79//!
80//! Per band, one syndrome row is also computed the dense way over the staged
81//! bytes and compared with the transform's, and the solved rows are re-encoded
82//! at that exponent and compared with the same syndrome, so neither half of the
83//! arm is trusted on its own word. A mismatch abandons the arm and
84//! the caller reruns the whole repair on the dense path. Nothing is written
85//! until the probe for that band has passed, and the dense rerun recomputes
86//! and rewrites every missing byte from sources the arm never touches, so a
87//! band already written is simply overwritten.
88
89use std::collections::HashMap;
90use std::fs::File;
91use std::path::PathBuf;
92use std::sync::atomic::{AtomicUsize, Ordering};
93
94use rayon::prelude::*;
95use reedsolomon_rs::fft::TransformError;
96use reedsolomon_rs::gf16_dft::{DftPlan, DftScratch};
97use reedsolomon_rs::vandermonde_solve::{ConsecutiveSolvePlan, SolveError};
98use tracing::{debug, info, warn};
99
100use crate::error::{Par2Error, Result};
101use crate::gf;
102use crate::gf_simd::FactorSrc;
103use crate::matrix;
104use crate::par2_set::Par2FileSet;
105use crate::repair::{
106    RepairOptions, RepairPlan, StreamSourceReader, build_write_targets, check_cancel,
107    read_stream_source_chunk,
108};
109use crate::types::{ProgressPhase, ProgressStage, ProgressUpdate};
110use crate::verify::FileAccess;
111
112/// Smallest band the arm will accept. Below this the strided read of every
113/// present slice degenerates into one syscall per few kilobytes per slice.
114///
115/// 2 KiB rather than a rounder 4 KiB because the staging arena is
116/// `n_present * band`: on a 16k-block set (16320 present slices, 64 KiB
117/// slices) a 64 MiB limit — the default when this floor was measured — buys a
118/// band of 3840 bytes, and a 4 KiB floor would have refused every repair on
119/// that set at that limit. At
120/// 2 KiB the same set repairs in 1.8 s (m=512), 3.0 s (m=2048) and 5.9 s
121/// (m=4096) against the dense path's 2.2 s, 7.9 s and 20.3 s, with a peak RSS
122/// of 95, 190 and 246 MB against dense's 86, 149 and 264 MB.
123pub(crate) const MIN_BAND: usize = 2 * 1024;
124
125/// Most bands the arm will split a slice into. A pass re-walks every present
126/// file's descriptor set; past a couple of hundred that bookkeeping, not the
127/// arithmetic, sets the runtime.
128pub(crate) const MAX_PASSES: usize = 256;
129
130/// Missing slices below which the arm never engages, whatever the fold ratio.
131///
132/// Two reasons, both measured. The transform's cost is dominated by its
133/// length-257 accumulation, `live_buckets * m` folds, and `live_buckets`
134/// saturates at 257 as soon as the set is large; at small `m` that fixed cost
135/// is most of the work and the ratio against `m * n` is near 1. And the arm's
136/// setup — plan build, arena allocation, per-band probe — is charged whole
137/// against a repair whose dense form may take milliseconds. Everyday repairs
138/// of a handful of blocks must not notice this module exists, and with this
139/// floor they never reach it: the gate is three integer comparisons.
140pub(crate) const MIN_MISSING: usize = 256;
141
142/// The transform must beat the dense fold count by this factor to engage.
143///
144/// A fold is not a constant-cost unit across the two arms: the dense executor
145/// folds long contiguous source rows with prepared factors and a tuned
146/// controller, while the transform's folds are short stripe passes over
147/// scratch rows, several of them dependent. Requiring a 4x paper margin is the
148/// cheapest way to keep the arm out of the region where its better fold count
149/// does not survive contact with the memory system: measured on an AVX2
150/// desktop, a 2.6x paper advantage ran 1.8x slower than the dense path, while
151/// every shape at 12x or better won.
152const FOLD_MARGIN: u64 = 4;
153
154/// Hard cap on how far the covering exponent range may exceed `m`.
155///
156/// The selected exponents are normally the smallest available and therefore
157/// contiguous; a deleted middle recovery volume opens a gap. Rows inside the
158/// range but outside the selection are computed and discarded, and they also
159/// occupy per-worker syndrome memory, so a badly scattered selection is capped
160/// here as well as by the fold gate.
161const MAX_RANGE_MULTIPLE: usize = 4;
162
163/// Rows at or above which the closed-form consecutive solve beats the m×m
164/// product enough to be worth its scratch.
165///
166/// Measured in `reedsolomon-rs`: the crossover against the explicit inverse is
167/// around 600 rows (2.1x at 1024, 17x at 8192), and its setup is ~500x cheaper
168/// than the Gauss-Jordan at 1024. Below this the product wins and needs no
169/// scratch at all.
170const CONSECUTIVE_SOLVE_MIN_ROWS: usize = 512;
171
172/// Stripe length the arm aims for before memory forces it smaller.
173const STRIPE_TARGET: usize = 16 * 1024;
174
175/// Smallest stripe the arm will shrink to. The DFT only needs an even length;
176/// this is where the per-stripe bookkeeping stops being worth it.
177const STRIPE_FLOOR: usize = 256;
178
179/// Alignment of bands and stripes.
180const ALIGN: usize = 64;
181
182/// Sources folded into one syndrome-combining destination per kernel call.
183const SOLVE_BATCH: usize = 16;
184
185/// In-process override of the arm, for tests and A/B runs.
186#[derive(Clone, Copy, Debug, PartialEq, Eq)]
187pub enum TransformArm {
188    /// Never take the transform arm.
189    Off,
190    /// Take it wherever it is admissible, ignoring the win gate.
191    On,
192}
193
194thread_local! {
195    static ARM_OVERRIDE: std::cell::Cell<Option<TransformArm>> =
196        const { std::cell::Cell::new(None) };
197}
198
199/// Force the repair transform arm on or off for repairs started on this thread.
200///
201/// `None` restores the automatic gate. This is an in-process, thread-local
202/// switch: it takes precedence over `RARPAR_PAR2_TRANSFORM`, and it is read
203/// once, on the thread that calls `execute_repair_with_options`.
204pub fn set_transform_arm_override(value: Option<TransformArm>) {
205    ARM_OVERRIDE.with(|cell| cell.set(value));
206}
207
208thread_local! {
209    static ARM_STATS: std::cell::Cell<TransformArmStats> =
210        const {
211            std::cell::Cell::new(TransformArmStats {
212                executed: 0,
213                diverged: 0,
214                consecutive_solves: 0,
215            })
216        };
217}
218
219/// What the transform arm did on this thread, since the process started.
220///
221/// Repairs are driven from the thread that calls
222/// [`crate::repair::execute_repair_with_options`], so these counters are that
223/// thread's own and never race. Exposed so a caller (or a test) can tell a
224/// transform-arm repair from a dense one without parsing logs.
225#[derive(Clone, Copy, Debug, Default, PartialEq, Eq)]
226pub struct TransformArmStats {
227    /// Repairs completed on the transform arm.
228    pub executed: u64,
229    /// Repairs the arm abandoned after its probe diverged.
230    pub diverged: u64,
231    /// Repairs the arm chose the closed-form consecutive solve for, rather
232    /// than the explicit inverse. Counted when the solver is picked, so it
233    /// covers diverged runs too.
234    pub consecutive_solves: u64,
235}
236
237/// This thread's [`TransformArmStats`].
238pub fn transform_arm_stats() -> TransformArmStats {
239    ARM_STATS.with(|cell| cell.get())
240}
241
242fn record_consecutive_solve() {
243    ARM_STATS.with(|cell| {
244        let mut stats = cell.get();
245        stats.consecutive_solves += 1;
246        cell.set(stats);
247    });
248}
249
250fn record_executed() {
251    ARM_STATS.with(|cell| {
252        let mut stats = cell.get();
253        stats.executed += 1;
254        cell.set(stats);
255    });
256}
257
258fn record_diverged() {
259    ARM_STATS.with(|cell| {
260        let mut stats = cell.get();
261        stats.diverged += 1;
262        cell.set(stats);
263    });
264}
265
266// Corrupt the arm's own syndrome probe so the divergence path can be tested.
267#[cfg(test)]
268thread_local! {
269    static PROBE_FAULT: std::cell::Cell<bool> = const { std::cell::Cell::new(false) };
270}
271
272/// Make the next band's probe disagree with the transform, once.
273#[cfg(test)]
274fn take_probe_fault() -> bool {
275    PROBE_FAULT.with(|cell| cell.replace(false))
276}
277
278#[cfg(test)]
279thread_local! {
280    static SOLVE_FAULT: std::cell::Cell<bool> = const { std::cell::Cell::new(false) };
281}
282
283/// Corrupt the next band's solved output, once.
284#[cfg(test)]
285fn take_solve_fault() -> bool {
286    SOLVE_FAULT.with(|cell| cell.replace(false))
287}
288
289/// The current thread's transform-arm override, if any.
290pub fn transform_arm_override() -> Option<TransformArm> {
291    ARM_OVERRIDE.with(|cell| cell.get())
292}
293
294fn arm_setting() -> Option<TransformArm> {
295    if let Some(value) = transform_arm_override() {
296        return Some(value);
297    }
298    match std::env::var("RARPAR_PAR2_TRANSFORM").ok().as_deref() {
299        Some("0") => Some(TransformArm::Off),
300        Some("1") => Some(TransformArm::On),
301        _ => None,
302    }
303}
304
305/// Whether this build can engage a GPU repair arm at all.
306///
307/// The transform arm is a CPU alternative and is only offered where the CPU
308/// dense path would have run. Ranking it against a live GPU session would mean
309/// re-deciding GPU policy, so it declines outright in any build that carries a
310/// GPU backend. Default builds (and every shipped `rarpar` artifact) carry
311/// none, so the arm is live there.
312const GPU_ARM_POSSIBLE: bool = cfg!(any(
313    all(
314        feature = "metal",
315        target_os = "macos",
316        target_arch = "aarch64"
317    ),
318    feature = "wgpu"
319));
320
321/// What `try_execute` did.
322#[derive(Debug)]
323pub(crate) enum TransformOutcome {
324    /// The repair is complete; every missing slice has been written.
325    Executed,
326    /// The arm did not run. Nothing was written; run the dense path.
327    Declined(&'static str),
328    /// The arm ran, diverged from its own probe and stopped. Bands already
329    /// written are stale but harmless — rerun the dense path over the whole
330    /// plan, which rewrites them from untouched sources.
331    Diverged,
332}
333
334/// The stripe-granular half of the repair: syndrome rows in, missing rows out.
335///
336/// This is the seam a Forney-style O(m)-folds solver drops into. An
337/// implementation sees one stripe of every syndrome row and must fill one
338/// stripe of every missing row, in `missing_global_indices` order; it may
339/// assume every region has the same, even length and that the destinations are
340/// pairwise disjoint and uninitialised.
341pub(crate) trait StripeSolver: Send + Sync {
342    /// Working bytes one worker needs at this stripe length.
343    fn scratch_bytes(&self, stripe_len: usize) -> usize;
344
345    /// Turn one stripe of syndromes into one stripe of every missing slice.
346    ///
347    /// `syndromes` is the transform's whole output buffer for the stripe —
348    /// `range_len` rows of `stripe_len` bytes — and `row_offsets[r]` is where
349    /// the `r`-th selected recovery exponent's row sits in it. The buffer may
350    /// be overwritten. `outputs[c]` receives missing slice `c`, in
351    /// `missing_global_indices` order.
352    fn solve_stripe(
353        &self,
354        syndromes: &mut [u8],
355        row_offsets: &[usize],
356        stripe_len: usize,
357        scratch: &mut [u8],
358        outputs: &mut [&mut [u8]],
359        cancelled: &dyn Fn() -> bool,
360    ) -> Result<()>;
361}
362
363/// Apply `A^-1` as an explicit m×m stripe-wise product.
364///
365/// `m^2` folds, no scratch, and no requirement on the exponents at all. This
366/// is the arm's fallback solver and the oracle its faster sibling is checked
367/// against.
368struct DenseInverseSolver<'a> {
369    /// `A^-1`, `m` rows by `m` columns — `RepairPlan::decode_matrix`.
370    inverse: &'a matrix::Matrix,
371}
372
373impl StripeSolver for DenseInverseSolver<'_> {
374    fn scratch_bytes(&self, _stripe_len: usize) -> usize {
375        0
376    }
377
378    fn solve_stripe(
379        &self,
380        syndromes: &mut [u8],
381        row_offsets: &[usize],
382        stripe_len: usize,
383        _scratch: &mut [u8],
384        outputs: &mut [&mut [u8]],
385        cancelled: &dyn Fn() -> bool,
386    ) -> Result<()> {
387        debug_assert_eq!(outputs.len(), self.inverse.rows);
388        debug_assert_eq!(row_offsets.len(), self.inverse.cols);
389        let syndromes: &[u8] = syndromes;
390        let mut batch: Vec<FactorSrc<'_>> = Vec::with_capacity(SOLVE_BATCH);
391        for (row, out) in outputs.iter_mut().enumerate() {
392            if row.is_multiple_of(64) && cancelled() {
393                return Err(Par2Error::Cancelled);
394            }
395            out.fill(0);
396            let factors = self.inverse.row(row);
397            batch.clear();
398            for (column, &factor) in factors.iter().enumerate() {
399                if factor == 0 {
400                    continue;
401                }
402                batch.push(FactorSrc {
403                    factor,
404                    src: &syndromes[row_offsets[column] * stripe_len..][..stripe_len],
405                });
406                if batch.len() == SOLVE_BATCH {
407                    crate::gf_simd::mul_acc_input_batch(out, &batch);
408                    batch.clear();
409                }
410            }
411            if !batch.is_empty() {
412                crate::gf_simd::mul_acc_input_batch(out, &batch);
413                batch.clear();
414            }
415        }
416        Ok(())
417    }
418}
419
420/// The closed-form solve for a consecutive exponent run.
421///
422/// [`ConsecutiveSolvePlan`] replaces the m×m product with a locator
423/// correlation and a Forney evaluation — `O(m)` folds per unknown instead of
424/// `m`, and a setup that costs a fraction of the Gauss-Jordan the inverse
425/// needs. It only exists for `e0, e0+1, ...`, and it buys that speed with a
426/// large per-stripe scratch, so both the exponent shape and the memory
427/// contract have to admit it before the arm picks it up.
428struct ConsecutiveSolver {
429    plan: ConsecutiveSolvePlan,
430}
431
432impl StripeSolver for ConsecutiveSolver {
433    fn scratch_bytes(&self, stripe_len: usize) -> usize {
434        self.plan.scratch_bytes(stripe_len)
435    }
436
437    fn solve_stripe(
438        &self,
439        syndromes: &mut [u8],
440        row_offsets: &[usize],
441        stripe_len: usize,
442        scratch: &mut [u8],
443        outputs: &mut [&mut [u8]],
444        cancelled: &dyn Fn() -> bool,
445    ) -> Result<()> {
446        let rows = self.plan.rows();
447        debug_assert_eq!(outputs.len(), rows);
448        // A consecutive selection lands on the first `rows` rows of the
449        // transform's range in order, which is exactly the contiguous buffer
450        // the solve wants; the arm only builds this solver in that case.
451        debug_assert!(row_offsets.iter().copied().eq(0..rows));
452        self.plan
453            .solve_stripe(
454                &mut syndromes[..rows * stripe_len],
455                stripe_len,
456                scratch,
457                cancelled,
458            )
459            .map_err(|error| match error {
460                SolveError::Cancelled => Par2Error::Cancelled,
461                other => Par2Error::ReedSolomonError {
462                    reason: format!("repair transform solve failed: {other}"),
463                },
464            })?;
465        let solved: &[u8] = syndromes;
466        for (row, out) in outputs.iter_mut().enumerate() {
467            out.copy_from_slice(&solved[row * stripe_len..][..stripe_len]);
468        }
469        Ok(())
470    }
471}
472
473/// The band/stripe geometry and the arenas it implies.
474#[derive(Debug, Clone, Copy)]
475pub(crate) struct BandGeometry {
476    pub(crate) band: usize,
477    pub(crate) stripe: usize,
478    pub(crate) passes: usize,
479    pub(crate) workers: usize,
480    pub(crate) arena_bytes: usize,
481}
482
483/// Solve the memory contract: the most workers, at the largest stripe and the
484/// largest band, that fit `budget` once every arena is summed.
485///
486/// Returns `None` when no admissible geometry exists — the caller then either
487/// tries a solver with a smaller scratch or takes the dense path. `range_len`
488/// is the covering exponent range, `rows` the missing count, `present` the
489/// available source count.
490#[allow(clippy::too_many_arguments)]
491pub(crate) fn plan_geometry(
492    slice_size: usize,
493    present: usize,
494    rows: usize,
495    range_len: usize,
496    max_workers: usize,
497    plan_bytes: usize,
498    per_worker_bytes: &dyn Fn(usize) -> usize,
499    budget: usize,
500) -> Option<BandGeometry> {
501    let free = budget.checked_sub(plan_bytes)?;
502    // Per band, every present slice, every output row, and the two probe rows
503    // hold one band's bytes.
504    let per_band_byte = present.checked_add(rows)?.checked_add(2)?;
505
506    let mut workers = max_workers.max(1);
507    loop {
508        let mut stripe = STRIPE_TARGET.min(slice_size.next_multiple_of(ALIGN));
509        stripe = (stripe / ALIGN).max(1) * ALIGN;
510        loop {
511            // One spare per-worker arena covers the single short tail stripe a
512            // slice length that is not a multiple of the stripe leaves behind.
513            let per_worker = range_len
514                .checked_mul(stripe)
515                .and_then(|syndrome| syndrome.checked_add(per_worker_bytes(stripe)));
516            if let Some(per_worker) = per_worker
517                && let Some(worker_total) = per_worker
518                    .checked_mul(workers)
519                    .and_then(|total| total.checked_add(per_worker))
520                && let Some(band_budget) = free.checked_sub(worker_total)
521            {
522                let band = ((band_budget / per_band_byte) / stripe) * stripe;
523                let band = band.min(slice_size.next_multiple_of(stripe));
524                if band >= stripe && band >= MIN_BAND.min(slice_size) {
525                    let passes = slice_size.div_ceil(band);
526                    if passes <= MAX_PASSES {
527                        return Some(BandGeometry {
528                            band,
529                            stripe,
530                            passes,
531                            workers,
532                            arena_bytes: plan_bytes + worker_total + band * per_band_byte,
533                        });
534                    }
535                }
536            }
537            if stripe <= STRIPE_FLOOR {
538                break;
539            }
540            stripe = (stripe / 2).next_multiple_of(ALIGN).max(STRIPE_FLOOR);
541        }
542        if workers == 1 {
543            return None;
544        }
545        workers = (workers / 2).max(1);
546    }
547}
548
549/// Pick the solver and the geometry together: the two are one decision,
550/// because a solver's per-stripe scratch is part of the memory contract.
551///
552/// The closed-form solve is tried first when the exponents are consecutive and
553/// the row count is past its crossover against the m×m product; if its scratch
554/// cannot be afforded, the explicit inverse — which needs none — is tried at
555/// the same budget before the arm gives up.
556fn choose_solver<'a>(
557    plan: &'a RepairPlan,
558    dft: &DftPlan,
559    range_len: usize,
560    budget: usize,
561    workers: usize,
562) -> Option<(Box<dyn StripeSolver + 'a>, BandGeometry, &'static str)> {
563    let rows = plan.missing_slices.len();
564    let present = plan.available_input_global_indices.len();
565    let slice_size = plan.slice_size as usize;
566    let geometry_for = |solver: &dyn StripeSolver| {
567        plan_geometry(
568            slice_size,
569            present,
570            rows,
571            range_len,
572            workers,
573            dft.plan_bytes(),
574            // A worker holds the transform's scratch and the solver's at once,
575            // plus one slice descriptor per present source and per output row.
576            &|stripe| {
577                dft.scratch_bytes(stripe)
578                    .saturating_add(solver.scratch_bytes(stripe))
579                    .saturating_add(
580                        present
581                            .saturating_add(rows)
582                            .saturating_mul(std::mem::size_of::<&[u8]>()),
583                    )
584            },
585            budget,
586        )
587    };
588
589    if rows >= CONSECUTIVE_SOLVE_MIN_ROWS && range_len == rows {
590        let missing_logs: Vec<u16> = plan
591            .missing_global_indices
592            .iter()
593            .map(|&global| gf::log(plan.constants[global]))
594            .collect();
595        match ConsecutiveSolvePlan::build_for_exponents(&missing_logs, &plan.recovery_exponents) {
596            Ok(consecutive) => {
597                let solver = ConsecutiveSolver { plan: consecutive };
598                if let Some(geometry) = geometry_for(&solver) {
599                    record_consecutive_solve();
600                    return Some((Box::new(solver), geometry, "consecutive"));
601                }
602                debug!(
603                    budget,
604                    "the consecutive solve does not fit the memory limit; trying the inverse"
605                );
606            }
607            Err(SolveError::NonConsecutive) => {}
608            Err(error) => {
609                debug!(%error, "the consecutive solve refused the selection");
610            }
611        }
612    }
613
614    let solver = DenseInverseSolver {
615        inverse: &plan.decode_matrix,
616    };
617    let geometry = geometry_for(&solver)?;
618    Some((Box::new(solver), geometry, "inverse"))
619}
620
621/// Try to run the repair on the transform arm.
622///
623/// `Ok(TransformOutcome::Declined)` and `Ok(TransformOutcome::Diverged)` both
624/// mean the caller must run the dense path; `Declined` additionally guarantees
625/// nothing was written. Errors (cancellation, I/O, write failures) are the
626/// caller's to propagate — they are not arm-specific and the dense path would
627/// hit them too.
628pub(crate) fn try_execute(
629    plan: &RepairPlan,
630    par2_set: &Par2FileSet,
631    file_access: &mut dyn FileAccess,
632    options: &RepairOptions,
633    budget: usize,
634) -> Result<TransformOutcome> {
635    let setting = arm_setting();
636    if setting == Some(TransformArm::Off) {
637        return Ok(TransformOutcome::Declined("forced off"));
638    }
639    if GPU_ARM_POSSIBLE {
640        return Ok(TransformOutcome::Declined("GPU arm is compiled in"));
641    }
642    // The arm's stripes run on rayon workers. Where workers cannot be spawned
643    // (plain single-threaded wasm) the dense controller already runs inline;
644    // decline before anything here touches the pool.
645    if !reedsolomon_rs::threading::parallel_enabled() {
646        return Ok(TransformOutcome::Declined(
647            "no worker threads on this target",
648        ));
649    }
650
651    let rows = plan.missing_slices.len();
652    let present = plan.available_input_global_indices.len();
653    let slice_size = plan.slice_size as usize;
654    if setting != Some(TransformArm::On) && rows < MIN_MISSING {
655        return Ok(TransformOutcome::Declined("below the missing-slice floor"));
656    }
657    if rows == 0 || present == 0 || slice_size == 0 {
658        return Ok(TransformOutcome::Declined("degenerate shape"));
659    }
660    if plan.decode_matrix.rows != rows || plan.decode_matrix.cols != rows {
661        return Ok(TransformOutcome::Declined("decode matrix is not m x m"));
662    }
663
664    let Some(&lowest) = plan.recovery_exponents.iter().min() else {
665        return Ok(TransformOutcome::Declined("no recovery exponents"));
666    };
667    let highest = *plan
668        .recovery_exponents
669        .iter()
670        .max()
671        .expect("a non-empty selection has a maximum");
672    let Some(range_len) = (highest as usize)
673        .checked_sub(lowest as usize)
674        .map(|d| d + 1)
675    else {
676        return Ok(TransformOutcome::Declined("exponent range underflow"));
677    };
678    if highest >= 65535 {
679        return Ok(TransformOutcome::Declined("exponent outside the transform"));
680    }
681    if range_len > rows.saturating_mul(MAX_RANGE_MULTIPLE) {
682        return Ok(TransformOutcome::Declined(
683            "exponent selection too scattered",
684        ));
685    }
686
687    let slots: Vec<u16> = plan
688        .available_input_global_indices
689        .iter()
690        .map(|&global| gf::log(plan.constants[global]))
691        .collect();
692    let dft = match DftPlan::build(&slots, lowest..highest + 1) {
693        Ok(dft) => dft,
694        Err(error) => {
695            debug!(%error, "repair transform plan refused the shape");
696            return Ok(TransformOutcome::Declined("transform plan refused"));
697        }
698    };
699
700    let transform_folds = dft
701        .region_folds()
702        .saturating_add((rows as u64).saturating_mul(rows as u64));
703    let dense_folds = (rows as u64).saturating_mul((present + rows) as u64);
704    if setting != Some(TransformArm::On)
705        && transform_folds.saturating_mul(FOLD_MARGIN) > dense_folds
706    {
707        debug!(
708            transform_folds,
709            dense_folds, "repair transform arm does not beat the dense path by enough"
710        );
711        return Ok(TransformOutcome::Declined("fold margin not met"));
712    }
713
714    let workers = rayon::current_num_threads().max(1);
715    let Some((solver, geometry, solver_name)) =
716        choose_solver(plan, &dft, range_len, budget, workers)
717    else {
718        info!(
719            budget,
720            present,
721            rows,
722            "repair transform arm does not fit the memory limit; taking the dense path"
723        );
724        return Ok(TransformOutcome::Declined("memory limit too small"));
725    };
726
727    info!(
728        missing_slices = rows,
729        present_slices = present,
730        solver = solver_name,
731        band_bytes = geometry.band,
732        stripe_bytes = geometry.stripe,
733        passes = geometry.passes,
734        workers = geometry.workers,
735        arena_bytes = geometry.arena_bytes,
736        transform_folds,
737        dense_folds,
738        exponent_range = range_len,
739        "repairing with the transform arm"
740    );
741
742    run(
743        plan,
744        par2_set,
745        file_access,
746        options,
747        &dft,
748        solver.as_ref(),
749        geometry,
750        lowest,
751    )
752}
753
754/// One worker's private buffers. Sized once from the geometry; the short tail
755/// stripe is the only case that reallocates, and only its DFT scratch.
756struct Worker {
757    syndromes: Vec<u8>,
758    scratch: DftScratch,
759    solve_scratch: Vec<u8>,
760}
761
762#[allow(clippy::too_many_arguments)]
763fn run(
764    plan: &RepairPlan,
765    par2_set: &Par2FileSet,
766    file_access: &mut dyn FileAccess,
767    options: &RepairOptions,
768    dft: &DftPlan,
769    solver: &dyn StripeSolver,
770    geometry: BandGeometry,
771    first_exponent: u32,
772) -> Result<TransformOutcome> {
773    let rows = plan.missing_slices.len();
774    let present = plan.available_input_global_indices.len();
775    let slice_size = plan.slice_size as usize;
776    let range_len = dft.output_count();
777    let band = geometry.band;
778    let stripe = geometry.stripe;
779
780    let write_targets = build_write_targets(plan, par2_set)?;
781    let mut recovery_files: HashMap<PathBuf, File> = HashMap::new();
782    let mut source_reader: Option<StreamSourceReader> = None;
783
784    // Row `r` of `syndrome_row` is the position of recovery exponent `r` inside
785    // the transform's contiguous output range.
786    let syndrome_row: Vec<usize> = plan
787        .recovery_exponents
788        .iter()
789        .map(|&exponent| (exponent - first_exponent) as usize)
790        .collect();
791
792    let mut staging = vec![0u8; present.checked_mul(band).expect("band arena fits")];
793    let mut output = vec![0u8; rows.checked_mul(band).expect("band arena fits")];
794    let mut probe_seen = vec![0u8; band];
795    let mut probe_expect = vec![0u8; band];
796
797    let total_bytes = slice_size as u64;
798    let mut band_index = 0usize;
799    let mut band_start = 0usize;
800    while band_start < slice_size {
801        check_cancel(options)?;
802        let band_len = band.min(slice_size - band_start);
803
804        for source in 0..present {
805            if source % 64 == 0 {
806                check_cancel(options)?;
807            }
808            read_stream_source_chunk(
809                plan,
810                par2_set,
811                file_access,
812                &mut recovery_files,
813                &mut source_reader,
814                present,
815                source,
816                band_start,
817                &mut staging[source * band..source * band + band_len],
818            )?;
819        }
820        for row in 0..rows {
821            if row % 64 == 0 {
822                check_cancel(options)?;
823            }
824            read_stream_source_chunk(
825                plan,
826                par2_set,
827                file_access,
828                &mut recovery_files,
829                &mut source_reader,
830                present,
831                present + row,
832                band_start,
833                &mut output[row * band..row * band + band_len],
834            )?;
835        }
836
837        // The probe exponent rotates so a systematic error in one row cannot
838        // hide behind a band boundary.
839        let probe = band_index % rows;
840        probe_expect[..band_len].copy_from_slice(&output[probe * band..probe * band + band_len]);
841
842        transform_band(
843            dft,
844            solver,
845            options,
846            &staging,
847            &mut output,
848            &mut probe_seen[..band_len],
849            &syndrome_row,
850            TransformBand {
851                present,
852                rows,
853                range_len,
854                band,
855                band_len,
856                stripe,
857                workers: geometry.workers,
858                probe,
859            },
860        )?;
861
862        #[cfg(test)]
863        if take_probe_fault() {
864            probe_seen[0] ^= 0xFF;
865        }
866
867        let factors: Vec<u16> = plan
868            .available_input_global_indices
869            .iter()
870            .map(|&global| gf::pow(plan.constants[global], plan.recovery_exponents[probe]))
871            .collect();
872        dense_row(
873            &staging,
874            band,
875            band_len,
876            &factors,
877            &mut probe_expect[..band_len],
878        );
879        if probe_expect[..band_len] != probe_seen[..band_len] {
880            warn!(
881                band = band_index,
882                probe_exponent = plan.recovery_exponents[probe],
883                "repair transform arm diverged from its dense probe; falling back"
884            );
885            record_diverged();
886            return Ok(TransformOutcome::Diverged);
887        }
888
889        // The probe above vouches for the transform only. The solve gets its
890        // own: re-encoding the repaired rows at the probe exponent must give
891        // the dense syndrome back, so XORing that re-encoding over a row equal
892        // to it has to leave zeros. `m` folds per band, and no solver shares
893        // any of it.
894        #[cfg(test)]
895        if take_solve_fault() {
896            output[0] ^= 0xFF;
897        }
898        let factors: Vec<u16> = plan
899            .missing_global_indices
900            .iter()
901            .map(|&global| gf::pow(plan.constants[global], plan.recovery_exponents[probe]))
902            .collect();
903        dense_row(
904            &output,
905            band,
906            band_len,
907            &factors,
908            &mut probe_seen[..band_len],
909        );
910        if probe_seen[..band_len].iter().any(|&byte| byte != 0) {
911            warn!(
912                band = band_index,
913                probe_exponent = plan.recovery_exponents[probe],
914                "repair transform arm's solve does not re-encode to its syndrome; falling back"
915            );
916            record_diverged();
917            return Ok(TransformOutcome::Diverged);
918        }
919
920        check_cancel(options)?;
921        for (row, target) in write_targets.iter().enumerate() {
922            let write_offset = target.offset + band_start as u64;
923            let remaining = target.file_end.saturating_sub(write_offset);
924            let write_len = remaining.min(band_len as u64) as usize;
925            if write_len == 0 {
926                continue;
927            }
928            file_access
929                .write_file_range(
930                    &target.file_id,
931                    write_offset,
932                    &output[row * band..row * band + write_len],
933                )
934                .map_err(|error| Par2Error::RepairWriteFailed {
935                    filename: target.filename.clone(),
936                    offset: write_offset,
937                    source: error,
938                })?;
939        }
940
941        band_start += band_len;
942        band_index += 1;
943        if let Some(ref progress) = options.progress {
944            progress(ProgressUpdate {
945                stage: ProgressStage::Repairing,
946                current: band_index.min(u32::MAX as usize) as u32,
947                total: geometry.passes.min(u32::MAX as usize) as u32,
948                bytes_processed: band_start as u64,
949                total_bytes: Some(total_bytes),
950                phase: ProgressPhase::Whole,
951            });
952        }
953    }
954
955    info!(missing_slices = rows, "transform-arm repair complete");
956    record_executed();
957    Ok(TransformOutcome::Executed)
958}
959
960/// Shape of one band, passed whole so the stripe loop keeps one argument list.
961#[derive(Clone, Copy)]
962struct TransformBand {
963    present: usize,
964    rows: usize,
965    range_len: usize,
966    band: usize,
967    band_len: usize,
968    stripe: usize,
969    workers: usize,
970    probe: usize,
971}
972
973/// Transform and solve one staged band, in place over `output`.
974///
975/// `output` arrives holding each selected recovery block's bytes for the band
976/// and leaves holding each missing slice's repaired bytes. `probe_seen`
977/// receives the transform's own syndrome row for the probe exponent.
978#[allow(clippy::too_many_arguments)]
979fn transform_band(
980    dft: &DftPlan,
981    solver: &dyn StripeSolver,
982    options: &RepairOptions,
983    staging: &[u8],
984    output: &mut [u8],
985    probe_seen: &mut [u8],
986    syndrome_row: &[usize],
987    shape: TransformBand,
988) -> Result<()> {
989    let stripes = shape.band_len.div_ceil(shape.stripe);
990    let output_base = output.as_mut_ptr() as usize;
991    let probe_base = probe_seen.as_mut_ptr() as usize;
992    let next = AtomicUsize::new(0);
993    let failure: std::sync::Mutex<Option<Par2Error>> = std::sync::Mutex::new(None);
994    let cancelled = || {
995        options
996            .cancel
997            .as_ref()
998            .is_some_and(|token| token.is_cancelled())
999    };
1000
1001    rayon::scope(|scope| {
1002        for _ in 0..shape.workers {
1003            let next = &next;
1004            let failure = &failure;
1005            let cancelled = &cancelled;
1006            scope.spawn(move |_| {
1007                let mut worker = Worker {
1008                    syndromes: vec![0u8; shape.range_len * shape.stripe],
1009                    scratch: DftScratch::new(dft, shape.stripe),
1010                    solve_scratch: vec![0u8; solver.scratch_bytes(shape.stripe)],
1011                };
1012                let mut sources: Vec<&[u8]> = Vec::with_capacity(shape.present);
1013                loop {
1014                    let stripe_index = next.fetch_add(1, Ordering::Relaxed);
1015                    if stripe_index >= stripes {
1016                        return;
1017                    }
1018                    if failure.lock().expect("stripe failure lock").is_some() {
1019                        return;
1020                    }
1021                    let offset = stripe_index * shape.stripe;
1022                    let len = shape.stripe.min(shape.band_len - offset);
1023                    if worker.scratch.stripe_len() != len {
1024                        worker.scratch = DftScratch::new(dft, len);
1025                        worker.solve_scratch = vec![0u8; solver.scratch_bytes(len)];
1026                    }
1027
1028                    let syndromes = &mut worker.syndromes[..shape.range_len * len];
1029                    syndromes.fill(0);
1030                    // SAFETY: stripes are disjoint byte ranges of every row, and
1031                    // exactly one worker owns `stripe_index` at a time.
1032                    for (row, &position) in syndrome_row.iter().enumerate() {
1033                        let src = unsafe {
1034                            std::slice::from_raw_parts(
1035                                (output_base as *const u8).add(row * shape.band + offset),
1036                                len,
1037                            )
1038                        };
1039                        syndromes[position * len..position * len + len].copy_from_slice(src);
1040                    }
1041
1042                    sources.clear();
1043                    for source in 0..shape.present {
1044                        sources.push(&staging[source * shape.band + offset..][..len]);
1045                    }
1046                    if let Err(error) =
1047                        dft.transform_stripe(&sources, syndromes, &mut worker.scratch, cancelled)
1048                    {
1049                        let mapped = match error {
1050                            TransformError::Cancelled => Par2Error::Cancelled,
1051                            other => Par2Error::ReedSolomonError {
1052                                reason: format!("repair transform stripe failed: {other}"),
1053                            },
1054                        };
1055                        *failure.lock().expect("stripe failure lock") = Some(mapped);
1056                        return;
1057                    }
1058
1059                    let probe_position = syndrome_row[shape.probe];
1060                    // SAFETY: one worker owns this stripe of the probe row.
1061                    let probe_dst = unsafe {
1062                        std::slice::from_raw_parts_mut((probe_base as *mut u8).add(offset), len)
1063                    };
1064                    probe_dst.copy_from_slice(
1065                        &syndromes[probe_position * len..probe_position * len + len],
1066                    );
1067
1068                    let mut output_refs: Vec<&mut [u8]> = Vec::with_capacity(shape.rows);
1069                    for row in 0..shape.rows {
1070                        // SAFETY: as above — disjoint stripe of a distinct row.
1071                        output_refs.push(unsafe {
1072                            std::slice::from_raw_parts_mut(
1073                                (output_base as *mut u8).add(row * shape.band + offset),
1074                                len,
1075                            )
1076                        });
1077                    }
1078                    if let Err(error) = solver.solve_stripe(
1079                        syndromes,
1080                        syndrome_row,
1081                        len,
1082                        &mut worker.solve_scratch,
1083                        &mut output_refs,
1084                        cancelled,
1085                    ) {
1086                        *failure.lock().expect("stripe failure lock") = Some(error);
1087                        return;
1088                    }
1089                }
1090            });
1091        }
1092    });
1093
1094    match failure.into_inner().expect("stripe failure lock") {
1095        Some(error) => Err(error),
1096        None => Ok(()),
1097    }
1098}
1099
1100/// XOR one dense syndrome row over the staged band into `dst`.
1101///
1102/// `dst` arrives holding `R_e` and leaves holding `S_e`. This is the probe's
1103/// oracle: the same sum the transform computes, with the shared code path
1104/// being only the multiply-accumulate kernel.
1105fn dense_row(staging: &[u8], band: usize, band_len: usize, factors: &[u16], dst: &mut [u8]) {
1106    debug_assert_eq!(dst.len(), band_len);
1107    let chunk = ALIGN * 64;
1108    dst.par_chunks_mut(chunk).enumerate().for_each(|(at, out)| {
1109        let offset = at * chunk;
1110        let mut batch: Vec<FactorSrc<'_>> = Vec::with_capacity(SOLVE_BATCH);
1111        for (source, &factor) in factors.iter().enumerate() {
1112            if factor == 0 {
1113                continue;
1114            }
1115            batch.push(FactorSrc {
1116                factor,
1117                src: &staging[source * band + offset..][..out.len()],
1118            });
1119            if batch.len() == SOLVE_BATCH {
1120                crate::gf_simd::mul_acc_input_batch(out, &batch);
1121                batch.clear();
1122            }
1123        }
1124        if !batch.is_empty() {
1125            crate::gf_simd::mul_acc_input_batch(out, &batch);
1126        }
1127    });
1128}
1129
1130#[cfg(test)]
1131mod tests {
1132    use super::*;
1133
1134    #[test]
1135    fn geometry_declines_a_budget_that_cannot_buy_a_band() {
1136        let geometry = plan_geometry(
1137            262_144,
1138            14_000,
1139            4_096,
1140            4_096,
1141            8,
1142            1 << 20,
1143            &|stripe| 1275 * stripe,
1144            8 * 1024 * 1024,
1145        );
1146        assert!(geometry.is_none(), "{geometry:?}");
1147    }
1148
1149    #[test]
1150    fn geometry_stays_inside_the_budget_it_is_given() {
1151        for budget in [
1152            64 * 1024 * 1024,
1153            256 * 1024 * 1024,
1154            1024 * 1024 * 1024,
1155            4096 * 1024 * 1024,
1156        ] {
1157            let Some(geometry) = plan_geometry(
1158                262_144,
1159                14_000,
1160                2_048,
1161                2_048,
1162                10,
1163                512 * 1024,
1164                &|stripe| 1275 * stripe,
1165                budget,
1166            ) else {
1167                continue;
1168            };
1169            assert!(geometry.arena_bytes <= budget, "{geometry:?} for {budget}");
1170            assert!(
1171                geometry.band.is_multiple_of(geometry.stripe),
1172                "{geometry:?}"
1173            );
1174            assert!(geometry.band.is_multiple_of(ALIGN), "{geometry:?}");
1175            assert!(geometry.stripe.is_multiple_of(ALIGN), "{geometry:?}");
1176            assert!(geometry.band >= MIN_BAND, "{geometry:?}");
1177            assert!(geometry.passes <= MAX_PASSES, "{geometry:?}");
1178        }
1179    }
1180
1181    #[test]
1182    fn geometry_declines_when_the_band_would_need_too_many_passes() {
1183        // A slice far larger than the budget can band: the pass cap, not the
1184        // band floor, is what refuses it.
1185        let geometry = plan_geometry(
1186            64 * 1024 * 1024,
1187            14_000,
1188            2_048,
1189            2_048,
1190            10,
1191            512 * 1024,
1192            &|stripe| 1275 * stripe,
1193            64 * 1024 * 1024,
1194        );
1195        assert!(geometry.is_none(), "{geometry:?}");
1196    }
1197
1198    /// Damage `slices` of a single-file set, repair it with the arm in the
1199    /// given state, and return the restored bytes.
1200    fn repair_once(
1201        arm: Option<TransformArm>,
1202        file_data: &[u8],
1203        slice_size: u64,
1204        recovery: usize,
1205        damaged: &[usize],
1206        fault: bool,
1207    ) -> (Vec<u8>, TransformArmStats) {
1208        let (set, file_id) =
1209            crate::repair::tests::setup_repairable_set(file_data, slice_size, recovery);
1210        let mut broken = file_data.to_vec();
1211        for &slice in damaged {
1212            let start = slice * slice_size as usize;
1213            let end = (start + slice_size as usize).min(broken.len());
1214            broken[start..end].fill(0xA5);
1215        }
1216        let mut access = crate::verify::MemoryFileAccess::new();
1217        access.add_file(file_id, broken);
1218
1219        let verification = crate::verify::verify_all(&set, &access);
1220        let plan = crate::repair::plan_repair(&set, &verification).expect("plan");
1221
1222        set_transform_arm_override(arm);
1223        #[cfg(test)]
1224        PROBE_FAULT.with(|cell| cell.set(fault));
1225        let before = transform_arm_stats();
1226        crate::repair::execute_repair(&plan, &set, &mut access).expect("repair");
1227        let after = transform_arm_stats();
1228        set_transform_arm_override(None);
1229        PROBE_FAULT.with(|cell| cell.set(false));
1230
1231        let restored = crate::verify::FileAccess::read_file(&access, &file_id).expect("read back");
1232        (
1233            restored,
1234            TransformArmStats {
1235                executed: after.executed - before.executed,
1236                diverged: after.diverged - before.diverged,
1237                consecutive_solves: after.consecutive_solves - before.consecutive_solves,
1238            },
1239        )
1240    }
1241
1242    fn noise(len: usize, seed: u64) -> Vec<u8> {
1243        let mut state = seed | 1;
1244        (0..len)
1245            .map(|_| {
1246                state ^= state << 13;
1247                state ^= state >> 7;
1248                state ^= state << 17;
1249                (state >> 25) as u8
1250            })
1251            .collect()
1252    }
1253
1254    #[test]
1255    fn the_arm_is_bit_identical_to_the_dense_path() {
1256        // Slice sizes that are and are not multiples of 64, a partial last
1257        // slice, and a spread of damage patterns.
1258        for (slice_size, slices, damaged) in [
1259            (64u64, 16usize, vec![0usize, 3, 9]),
1260            (68, 12, vec![1, 2, 3, 4]),
1261            (256, 9, vec![0, 8]),
1262            (4, 20, vec![5, 6, 7]),
1263        ] {
1264            let full = slice_size as usize * slices;
1265            // A short final slice: the file stops before its last slice ends.
1266            let data = noise(full - slice_size as usize / 2, 0x51 + slice_size);
1267            let recovery = damaged.len() + 2;
1268            let (dense, dense_stats) = repair_once(
1269                Some(TransformArm::Off),
1270                &data,
1271                slice_size,
1272                recovery,
1273                &damaged,
1274                false,
1275            );
1276            let (transform, transform_stats) = repair_once(
1277                Some(TransformArm::On),
1278                &data,
1279                slice_size,
1280                recovery,
1281                &damaged,
1282                false,
1283            );
1284            assert_eq!(dense_stats.executed, 0, "slice_size={slice_size}");
1285            assert_eq!(transform_stats.executed, 1, "slice_size={slice_size}");
1286            assert_eq!(dense, data, "dense: slice_size={slice_size}");
1287            assert_eq!(transform, data, "transform: slice_size={slice_size}");
1288        }
1289    }
1290
1291    #[test]
1292    fn a_non_contiguous_exponent_selection_still_matches() {
1293        // Delete a middle recovery volume: the selection keeps its count but
1294        // its covering range is now wider than the selection itself.
1295        let slice_size = 64u64;
1296        let data = noise(64 * 20, 0xBEEF);
1297        let damaged = [2usize, 5, 11];
1298        let (set, file_id) = crate::repair::tests::setup_repairable_set(&data, slice_size, 8);
1299        let mut set = set;
1300        set.recovery_slices.remove(&1);
1301        set.recovery_slices.remove(&2);
1302
1303        let mut restored = Vec::new();
1304        for arm in [TransformArm::Off, TransformArm::On] {
1305            let mut broken = data.clone();
1306            for &slice in &damaged {
1307                broken[slice * 64..slice * 64 + 64].fill(0x5A);
1308            }
1309            let mut access = crate::verify::MemoryFileAccess::new();
1310            access.add_file(file_id, broken);
1311            let verification = crate::verify::verify_all(&set, &access);
1312            let plan = crate::repair::plan_repair(&set, &verification).expect("plan");
1313            assert!(
1314                plan.recovery_exponents.iter().max().unwrap()
1315                    - plan.recovery_exponents.iter().min().unwrap()
1316                    > 2,
1317                "the gap must survive selection: {:?}",
1318                plan.recovery_exponents
1319            );
1320            set_transform_arm_override(Some(arm));
1321            let before = transform_arm_stats();
1322            crate::repair::execute_repair(&plan, &set, &mut access).expect("repair");
1323            let after = transform_arm_stats();
1324            set_transform_arm_override(None);
1325            assert_eq!(
1326                after.executed - before.executed,
1327                u64::from(arm == TransformArm::On),
1328                "{arm:?}"
1329            );
1330            restored.push(crate::verify::FileAccess::read_file(&access, &file_id).unwrap());
1331        }
1332        assert_eq!(restored[0], data);
1333        assert_eq!(restored[1], data);
1334    }
1335
1336    #[test]
1337    fn a_probe_mismatch_abandons_the_arm_and_the_dense_path_finishes_the_repair() {
1338        let slice_size = 64u64;
1339        let data = noise(64 * 16, 0xC0FFEE);
1340        let damaged = [1usize, 4, 7];
1341        let (restored, stats) =
1342            repair_once(Some(TransformArm::On), &data, slice_size, 6, &damaged, true);
1343        assert_eq!(stats.diverged, 1, "the injected fault must be caught");
1344        assert_eq!(
1345            stats.executed, 0,
1346            "a diverged arm must not claim the repair"
1347        );
1348        assert_eq!(
1349            restored, data,
1350            "the dense rerun must still restore the file"
1351        );
1352    }
1353
1354    #[test]
1355    fn a_wrong_solve_abandons_the_arm_and_the_dense_path_finishes_the_repair() {
1356        // The syndrome probe cannot see a solver that turns right syndromes
1357        // into wrong slices; the re-encode check has to.
1358        let slice_size = 64u64;
1359        let data = noise(64 * 16, 0xBADC0DE);
1360        let damaged = [0usize, 5, 11];
1361        SOLVE_FAULT.with(|cell| cell.set(true));
1362        let (restored, stats) = repair_once(
1363            Some(TransformArm::On),
1364            &data,
1365            slice_size,
1366            6,
1367            &damaged,
1368            false,
1369        );
1370        SOLVE_FAULT.with(|cell| cell.set(false));
1371        assert_eq!(stats.diverged, 1, "the corrupted solve must be caught");
1372        assert_eq!(stats.executed, 0);
1373        assert_eq!(restored, data);
1374    }
1375
1376    #[test]
1377    fn a_cancelled_token_stops_the_arm_mid_band() {
1378        // `transform_band` is the arm's inner loop: a token cancelled while a
1379        // band is in flight must surface as `Cancelled`, not as a wrong answer.
1380        let slots: Vec<u16> = (1..=32u16).map(gf::log).collect();
1381        let dft = DftPlan::build(&slots, 0..4).expect("plan");
1382        let solver = DenseInverseSolver {
1383            inverse: &matrix::Matrix::identity(4),
1384        };
1385        let cancel = crate::types::CancellationToken::new();
1386        cancel.cancel();
1387        let options = RepairOptions {
1388            cancel: Some(cancel),
1389            progress: None,
1390            memory_limit: None,
1391        };
1392        let staging = vec![0u8; slots.len() * 128];
1393        let mut output = vec![0u8; 4 * 128];
1394        let mut probe = vec![0u8; 128];
1395        let error = transform_band(
1396            &dft,
1397            &solver,
1398            &options,
1399            &staging,
1400            &mut output,
1401            &mut probe,
1402            &[0, 1, 2, 3],
1403            TransformBand {
1404                present: slots.len(),
1405                rows: 4,
1406                range_len: 4,
1407                band: 128,
1408                band_len: 128,
1409                stripe: 128,
1410                workers: 2,
1411                probe: 0,
1412            },
1413        )
1414        .expect_err("a cancelled token must stop the band");
1415        assert!(matches!(error, Par2Error::Cancelled), "{error:?}");
1416    }
1417
1418    #[test]
1419    fn a_tiny_memory_limit_still_repairs_bit_identically() {
1420        // Force many passes: the band shrinks to the floor and the slice is
1421        // walked in pieces.
1422        let slice_size = 4096u64;
1423        let data = noise(4096 * 12, 0x1234);
1424        let damaged = [0usize, 5, 9];
1425        let (set, file_id) = crate::repair::tests::setup_repairable_set(&data, slice_size, 6);
1426
1427        let mut restored = Vec::new();
1428        for (arm, limit) in [
1429            (TransformArm::Off, None),
1430            (TransformArm::On, Some(1024 * 1024)),
1431        ] {
1432            let mut broken = data.clone();
1433            for &slice in &damaged {
1434                broken[slice * 4096..slice * 4096 + 4096].fill(0x11);
1435            }
1436            let mut access = crate::verify::MemoryFileAccess::new();
1437            access.add_file(file_id, broken);
1438            let verification = crate::verify::verify_all(&set, &access);
1439            let plan = crate::repair::plan_repair(&set, &verification).expect("plan");
1440            set_transform_arm_override(Some(arm));
1441            let before = transform_arm_stats();
1442            crate::repair::execute_repair_with_options(
1443                &plan,
1444                &set,
1445                &mut access,
1446                &RepairOptions {
1447                    cancel: None,
1448                    progress: None,
1449                    memory_limit: limit,
1450                },
1451            )
1452            .expect("repair");
1453            let after = transform_arm_stats();
1454            set_transform_arm_override(None);
1455            assert_eq!(
1456                after.executed - before.executed,
1457                u64::from(arm == TransformArm::On)
1458            );
1459            restored.push(crate::verify::FileAccess::read_file(&access, &file_id).unwrap());
1460        }
1461        assert_eq!(restored[0], data);
1462        assert_eq!(restored[1], data);
1463    }
1464
1465    #[test]
1466    fn a_limit_too_small_for_a_band_falls_back_to_the_dense_path() {
1467        let slice_size = 4096u64;
1468        let data = noise(4096 * 12, 0x99);
1469        let (set, file_id) = crate::repair::tests::setup_repairable_set(&data, slice_size, 6);
1470        let mut broken = data.clone();
1471        broken[..4096].fill(0x22);
1472        let mut access = crate::verify::MemoryFileAccess::new();
1473        access.add_file(file_id, broken);
1474        let verification = crate::verify::verify_all(&set, &access);
1475        let plan = crate::repair::plan_repair(&set, &verification).expect("plan");
1476
1477        // Ask the arm directly: what the dense path makes of so small a limit
1478        // is its own business and differs by kernel tier.
1479        set_transform_arm_override(Some(TransformArm::On));
1480        let before = transform_arm_stats();
1481        let broken_before = crate::verify::FileAccess::read_file(&access, &file_id).unwrap();
1482        let outcome = try_execute(
1483            &plan,
1484            &set,
1485            &mut access,
1486            &RepairOptions {
1487                cancel: None,
1488                progress: None,
1489                memory_limit: Some(16 * 1024),
1490            },
1491            16 * 1024,
1492        )
1493        .expect("declining is not an error");
1494        let after = transform_arm_stats();
1495        set_transform_arm_override(None);
1496        assert!(
1497            matches!(outcome, TransformOutcome::Declined(_)),
1498            "{outcome:?}"
1499        );
1500        assert_eq!(after, before, "a declined arm must not count as a run");
1501        assert_eq!(
1502            crate::verify::FileAccess::read_file(&access, &file_id).unwrap(),
1503            broken_before,
1504            "a declined arm must not have written anything"
1505        );
1506
1507        set_transform_arm_override(Some(TransformArm::Off));
1508        crate::repair::execute_repair(&plan, &set, &mut access).expect("dense repair");
1509        set_transform_arm_override(None);
1510        assert_eq!(
1511            crate::verify::FileAccess::read_file(&access, &file_id).unwrap(),
1512            data
1513        );
1514    }
1515
1516    #[test]
1517    fn the_arm_override_is_thread_local() {
1518        set_transform_arm_override(Some(TransformArm::Off));
1519        assert_eq!(transform_arm_override(), Some(TransformArm::Off));
1520        let seen = std::thread::spawn(transform_arm_override).join().unwrap();
1521        assert_eq!(seen, None);
1522        set_transform_arm_override(None);
1523        assert_eq!(transform_arm_override(), None);
1524    }
1525}