1use 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
112pub(crate) const MIN_BAND: usize = 2 * 1024;
124
125pub(crate) const MAX_PASSES: usize = 256;
129
130pub(crate) const MIN_MISSING: usize = 256;
141
142const FOLD_MARGIN: u64 = 4;
153
154const MAX_RANGE_MULTIPLE: usize = 4;
162
163const CONSECUTIVE_SOLVE_MIN_ROWS: usize = 512;
171
172const STRIPE_TARGET: usize = 16 * 1024;
174
175const STRIPE_FLOOR: usize = 256;
178
179const ALIGN: usize = 64;
181
182const SOLVE_BATCH: usize = 16;
184
185#[derive(Clone, Copy, Debug, PartialEq, Eq)]
187pub enum TransformArm {
188 Off,
190 On,
192}
193
194thread_local! {
195 static ARM_OVERRIDE: std::cell::Cell<Option<TransformArm>> =
196 const { std::cell::Cell::new(None) };
197}
198
199pub 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#[derive(Clone, Copy, Debug, Default, PartialEq, Eq)]
226pub struct TransformArmStats {
227 pub executed: u64,
229 pub diverged: u64,
231 pub consecutive_solves: u64,
235}
236
237pub 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#[cfg(test)]
268thread_local! {
269 static PROBE_FAULT: std::cell::Cell<bool> = const { std::cell::Cell::new(false) };
270}
271
272#[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#[cfg(test)]
285fn take_solve_fault() -> bool {
286 SOLVE_FAULT.with(|cell| cell.replace(false))
287}
288
289pub 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
305const 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#[derive(Debug)]
323pub(crate) enum TransformOutcome {
324 Executed,
326 Declined(&'static str),
328 Diverged,
332}
333
334pub(crate) trait StripeSolver: Send + Sync {
342 fn scratch_bytes(&self, stripe_len: usize) -> usize;
344
345 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
363struct DenseInverseSolver<'a> {
369 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
420struct 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 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#[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#[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 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 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
549fn 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 &|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
621pub(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 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
754struct 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 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 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 #[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#[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#[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 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 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 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
1100fn 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 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 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 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 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 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 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 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 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 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}