1use std::collections::HashMap;
41
42mod coords;
43mod decode;
44mod encode;
45mod error;
46mod repair;
47mod transforms;
48
49pub use error::ClayError;
50
51const MAX_RS_SHARDS: usize = 32768;
52
53use decode::{decode as decode_chunks, RsCodec};
54use encode::encode as encode_chunks;
55use repair::{minimum_to_repair as min_repair, repair as repair_chunk};
56
57#[derive(Clone, Debug)]
59pub struct ClayCode {
60 pub k: usize,
62 pub m: usize,
64 pub n: usize,
66 pub d: usize,
68 pub q: usize,
70 pub t: usize,
72 pub nu: usize,
74 pub sub_chunk_no: usize,
76 pub beta: usize,
78 original_count: usize,
80 rs: RsCodec,
82}
83
84impl ClayCode {
85 pub fn new(k: usize, m: usize, d: usize) -> Result<Self, ClayError> {
95 if k < 1 {
96 return Err(ClayError::InvalidParameters("k must be at least 1".into()));
97 }
98 if m < 1 {
99 return Err(ClayError::InvalidParameters("m must be at least 1".into()));
100 }
101 if d < k + 1 || d > k + m - 1 {
102 return Err(ClayError::InvalidParameters(format!(
103 "d must be in range [{}, {}], got {}",
104 k + 1,
105 k + m - 1,
106 d
107 )));
108 }
109
110 let q = d - k + 1;
111 let n = k + m;
112
113 let nu = if n % q == 0 { 0 } else { q - (n % q) };
115
116 let t = (n + nu) / q;
117
118 let sub_chunk_no = checked_pow(q, t).ok_or_else(|| {
120 ClayError::Overflow(format!("q^t = {}^{} overflows", q, t))
121 })?;
122
123 let beta = sub_chunk_no / q; let original_count = k + nu;
127 let recovery_count = m;
128 if original_count > MAX_RS_SHARDS || recovery_count > MAX_RS_SHARDS {
129 return Err(ClayError::InvalidParameters(
130 "Total nodes exceeds reed-solomon limit of 32768".into(),
131 ));
132 }
133
134 let rs = RsCodec::new(original_count, recovery_count)
135 .map_err(|e| ClayError::InvalidParameters(format!("RS init failed: {:?}", e)))?;
136
137 Ok(ClayCode {
138 k,
139 m,
140 n,
141 d,
142 q,
143 t,
144 nu,
145 sub_chunk_no,
146 beta,
147 original_count,
148 rs,
149 })
150 }
151
152 pub fn new_default(k: usize, m: usize) -> Result<Self, ClayError> {
154 Self::new(k, m, k + m - 1)
155 }
156
157 fn encode_params(&self) -> encode::EncodeParams {
159 encode::EncodeParams {
160 k: self.k,
161 m: self.m,
162 n: self.n,
163 q: self.q,
164 t: self.t,
165 nu: self.nu,
166 sub_chunk_no: self.sub_chunk_no,
167 original_count: self.original_count,
168 }
169 }
170
171 pub fn encode(&self, data: &[u8]) -> Vec<Vec<u8>> {
179 encode_chunks(&self.encode_params(), &self.rs, data)
180 }
181
182 pub fn decode(
191 &self,
192 available: &HashMap<usize, Vec<u8>>,
193 erasures: &[usize],
194 ) -> Result<Vec<u8>, ClayError> {
195 decode_chunks(&self.encode_params(), &self.rs, available, erasures)
196 }
197
198 pub fn minimum_to_repair(
210 &self,
211 lost_node: usize,
212 available: &[usize],
213 ) -> Result<Vec<(usize, Vec<usize>)>, ClayError> {
214 min_repair(&self.encode_params(), lost_node, available)
215 }
216
217 pub fn repair(
229 &self,
230 lost_node: usize,
231 helper_data: &HashMap<usize, Vec<u8>>,
232 chunk_size: usize,
233 ) -> Result<Vec<u8>, ClayError> {
234 repair_chunk(&self.encode_params(), &self.rs, lost_node, helper_data, chunk_size)
235 }
236
237 pub fn normalized_repair_bandwidth(&self) -> f64 {
242 (self.d as f64) / ((self.k as f64) * (self.d - self.k + 1) as f64)
243 }
244}
245
246fn checked_pow(base: usize, exp: usize) -> Option<usize> {
248 let mut result: usize = 1;
249 let mut b = base;
250 let mut e = exp;
251 while e > 0 {
252 if e & 1 == 1 {
253 result = result.checked_mul(b)?;
254 }
255 e >>= 1;
256 if e > 0 {
257 b = b.checked_mul(b)?;
258 }
259 }
260 Some(result)
261}
262
263#[cfg(test)]
264mod tests {
265 use super::*;
266
267 #[test]
268 fn test_basic_encode_decode() {
269 let clay = ClayCode::new(4, 2, 5).unwrap();
270 let data = b"Test data for Clay codes - not empty!";
271 let chunks = clay.encode(data);
272 assert_eq!(chunks.len(), 6); let mut available: HashMap<usize, Vec<u8>> = HashMap::new();
276 for (i, chunk) in chunks.iter().enumerate() {
277 available.insert(i, chunk.clone());
278 }
279 let decoded = clay.decode(&available, &[]).unwrap();
280
281 assert_eq!(&decoded[..data.len()], &data[..]);
283 }
284
285 #[test]
286 fn test_decode_with_erasures() {
287 let clay = ClayCode::new(4, 2, 5).unwrap();
288 let data = b"Test data for Clay codes - testing erasure recovery!";
289 let chunks = clay.encode(data);
290
291 let mut available: HashMap<usize, Vec<u8>> = HashMap::new();
293 for (i, chunk) in chunks.iter().enumerate() {
294 if i != 0 {
295 available.insert(i, chunk.clone());
296 }
297 }
298 let decoded = clay.decode(&available, &[0]).unwrap();
299 assert_eq!(&decoded[..data.len()], &data[..]);
300
301 let mut available: HashMap<usize, Vec<u8>> = HashMap::new();
303 for (i, chunk) in chunks.iter().enumerate() {
304 if i != 5 {
305 available.insert(i, chunk.clone());
306 }
307 }
308 let decoded = clay.decode(&available, &[5]).unwrap();
309 assert_eq!(&decoded[..data.len()], &data[..]);
310
311 let mut available: HashMap<usize, Vec<u8>> = HashMap::new();
313 for (i, chunk) in chunks.iter().enumerate() {
314 if i != 0 && i != 5 {
315 available.insert(i, chunk.clone());
316 }
317 }
318 let decoded = clay.decode(&available, &[0, 5]).unwrap();
319 assert_eq!(&decoded[..data.len()], &data[..]);
320 }
321
322 #[test]
323 fn test_parameters() {
324 let clay = ClayCode::new(4, 2, 5).unwrap();
326 assert_eq!(clay.q, 2);
327 assert_eq!(clay.t, 3);
328 assert_eq!(clay.sub_chunk_no, 8); assert_eq!(clay.beta, 4); let clay2 = ClayCode::new(10, 4, 13).unwrap();
333 assert_eq!(clay2.q, 4);
334 assert_eq!(clay2.t, 4);
335 assert_eq!(clay2.sub_chunk_no, 256); assert_eq!(clay2.beta, 64); }
338
339 #[test]
340 fn test_minimum_to_repair() {
341 let clay = ClayCode::new(4, 2, 5).unwrap();
342 let available: Vec<usize> = vec![1, 2, 3, 4, 5];
343 let helper_info = clay.minimum_to_repair(0, &available).unwrap();
344
345 assert_eq!(helper_info.len(), 5);
347
348 for (_, indices) in &helper_info {
350 assert_eq!(indices.len(), 4);
351 }
352 }
353
354 #[test]
355 fn test_repair_bandwidth_verification() {
356 let clay = ClayCode::new(4, 2, 5).unwrap();
358 let data = b"Test data for bandwidth verification of Clay codes repair!";
359 let chunks = clay.encode(data);
360 let chunk_size = chunks[0].len();
361
362 let available: Vec<usize> = vec![1, 2, 3, 4, 5];
364 let helper_info = clay.minimum_to_repair(0, &available).unwrap();
365
366 let sub_chunk_size = chunk_size / clay.sub_chunk_no;
368 let total_repair_subchunks: usize = helper_info
369 .iter()
370 .map(|(_, indices)| indices.len())
371 .sum();
372 let total_repair_bytes = total_repair_subchunks * sub_chunk_size;
373
374 let full_decode_bytes = clay.k * chunk_size;
375
376 let ratio = total_repair_bytes as f64 / full_decode_bytes as f64;
378 println!(
379 "Repair bandwidth: {} bytes, Full decode: {} bytes, Ratio: {:.3}",
380 total_repair_bytes, full_decode_bytes, ratio
381 );
382
383 assert!(
384 total_repair_bytes < full_decode_bytes * 7 / 10,
385 "Repair bandwidth {} should be < 70% of full decode {}",
386 total_repair_bytes,
387 full_decode_bytes
388 );
389 }
390
391 #[test]
392 fn test_repair_correctness() {
393 let clay = ClayCode::new(4, 2, 5).unwrap();
394 let data = b"Test data for repair correctness verification!!!!";
395 let chunks = clay.encode(data);
396 let chunk_size = chunks[0].len();
397 let sub_chunk_size = chunk_size / clay.sub_chunk_no;
398
399 for lost_node in 0..clay.n {
401 let available: Vec<usize> = (0..clay.n).filter(|&i| i != lost_node).collect();
402 let helper_info = clay.minimum_to_repair(lost_node, &available).unwrap();
403
404 let mut partial_data: HashMap<usize, Vec<u8>> = HashMap::new();
406 for (helper_idx, indices) in &helper_info {
407 let mut helper_partial = Vec::new();
408 for &sc_idx in indices {
409 let start_byte = sc_idx * sub_chunk_size;
410 let end_byte = (sc_idx + 1) * sub_chunk_size;
411 helper_partial.extend_from_slice(&chunks[*helper_idx][start_byte..end_byte]);
412 }
413 partial_data.insert(*helper_idx, helper_partial);
414 }
415
416 let recovered = clay.repair(lost_node, &partial_data, chunk_size).unwrap();
418
419 assert_eq!(
421 recovered, chunks[lost_node],
422 "Repair failed for node {}",
423 lost_node
424 );
425 }
426 }
427
428 #[test]
429 fn test_various_parameters() {
430 let params = vec![
432 (4, 2, 5), (9, 3, 11), (10, 4, 13), ];
436
437 for (k, m, d) in params {
438 let clay = ClayCode::new(k, m, d).unwrap();
439 let data_size = k * clay.sub_chunk_no * 2;
440 let data: Vec<u8> = (0..data_size).map(|i| (i % 256) as u8).collect();
441 let chunks = clay.encode(&data);
442
443 let mut available: HashMap<usize, Vec<u8>> = HashMap::new();
445 for (i, chunk) in chunks.iter().enumerate() {
446 if i != 0 {
447 available.insert(i, chunk.clone());
448 }
449 }
450 let decoded = clay.decode(&available, &[0]).unwrap();
451 assert_eq!(
452 &decoded[..data.len()],
453 &data[..],
454 "Failed for params ({}, {}, {})",
455 k,
456 m,
457 d
458 );
459 }
460 }
461
462 #[test]
463 fn test_repair_all_nodes_various_params() {
464 let params = vec![(4, 2, 5), (9, 3, 11)];
465
466 for (k, m, d) in params {
467 let clay = ClayCode::new(k, m, d).unwrap();
468 let data_size = k * clay.sub_chunk_no;
469 let data: Vec<u8> = (0..data_size).map(|i| ((i * 7 + 13) % 256) as u8).collect();
470 let chunks = clay.encode(&data);
471 let chunk_size = chunks[0].len();
472 let sub_chunk_size = chunk_size / clay.sub_chunk_no;
473
474 for lost_node in 0..clay.n {
475 let available: Vec<usize> = (0..clay.n).filter(|&i| i != lost_node).collect();
476 let helper_info = clay.minimum_to_repair(lost_node, &available).unwrap();
477
478 let mut partial_data: HashMap<usize, Vec<u8>> = HashMap::new();
479 for (helper_idx, indices) in &helper_info {
480 let mut helper_partial = Vec::new();
481 for &sc_idx in indices {
482 let start_byte = sc_idx * sub_chunk_size;
483 let end_byte = (sc_idx + 1) * sub_chunk_size;
484 helper_partial.extend_from_slice(&chunks[*helper_idx][start_byte..end_byte]);
485 }
486 partial_data.insert(*helper_idx, helper_partial);
487 }
488
489 let recovered = clay.repair(lost_node, &partial_data, chunk_size).unwrap();
490 assert_eq!(
491 recovered, chunks[lost_node],
492 "Repair failed for node {} with params ({}, {}, {})",
493 lost_node, k, m, d
494 );
495 }
496 }
497 }
498
499 #[test]
500 fn test_decode_max_erasures() {
501 let clay = ClayCode::new(4, 2, 5).unwrap();
502 let data: Vec<u8> = (0..256).map(|i| (i % 256) as u8).collect();
503 let chunks = clay.encode(&data);
504
505 let patterns = vec![vec![0, 5], vec![0, 1], vec![4, 5], vec![1, 3]];
507
508 for erasures in patterns {
509 let mut available: HashMap<usize, Vec<u8>> = HashMap::new();
510 for (i, chunk) in chunks.iter().enumerate() {
511 if !erasures.contains(&i) {
512 available.insert(i, chunk.clone());
513 }
514 }
515 let decoded = clay.decode(&available, &erasures).unwrap();
516 assert_eq!(
517 &decoded[..data.len()],
518 &data[..],
519 "Failed for erasures {:?}",
520 erasures
521 );
522 }
523 }
524
525 #[test]
526 fn test_normalized_repair_bandwidth() {
527 let test_cases = vec![
528 ((4, 2, 5), 0.625),
529 ((9, 3, 11), 0.407),
530 ((10, 4, 13), 0.325),
531 ];
532
533 for ((k, m, d), expected) in test_cases {
534 let clay = ClayCode::new(k, m, d).unwrap();
535 let actual = clay.normalized_repair_bandwidth();
536 assert!(
537 (actual - expected).abs() < 0.01,
538 "Expected {}, got {} for ({}, {}, {})",
539 expected,
540 actual,
541 k,
542 m,
543 d
544 );
545 }
546 }
547
548 #[test]
549 fn test_random_data() {
550 use rand::Rng;
551 let mut rng = rand::thread_rng();
552
553 let clay = ClayCode::new(4, 2, 5).unwrap();
554 let data_size = clay.k * clay.sub_chunk_no * 4;
555 let data: Vec<u8> = (0..data_size).map(|_| rng.gen()).collect();
556 let chunks = clay.encode(&data);
557
558 let mut available: HashMap<usize, Vec<u8>> = HashMap::new();
560 for (i, chunk) in chunks.iter().enumerate() {
561 available.insert(i, chunk.clone());
562 }
563 let decoded = clay.decode(&available, &[]).unwrap();
564 assert_eq!(&decoded[..data.len()], &data[..]);
565
566 let mut available: HashMap<usize, Vec<u8>> = HashMap::new();
568 for (i, chunk) in chunks.iter().enumerate() {
569 if i != 2 {
570 available.insert(i, chunk.clone());
571 }
572 }
573 let decoded = clay.decode(&available, &[2]).unwrap();
574 assert_eq!(&decoded[..data.len()], &data[..]);
575 }
576
577 #[test]
578 fn test_checked_pow_overflow() {
579 assert!(checked_pow(2, 63).is_some());
581 assert!(checked_pow(2, 64).is_none()); assert!(checked_pow(10, 20).is_none()); }
584
585 #[test]
586 fn test_invalid_parameters() {
587 assert!(ClayCode::new(0, 2, 1).is_err());
589
590 assert!(ClayCode::new(4, 0, 3).is_err());
592
593 assert!(ClayCode::new(4, 2, 4).is_err()); assert!(ClayCode::new(4, 2, 6).is_err()); }
597
598 #[test]
599 fn test_clone_and_debug() {
600 let clay = ClayCode::new(4, 2, 5).unwrap();
601 let clay2 = clay.clone();
602 assert_eq!(clay2.k, clay.k);
603 assert_eq!(clay2.m, clay.m);
604 assert_eq!(clay2.d, clay.d);
605 let debug_str = format!("{:?}", clay);
607 assert!(debug_str.contains("ClayCode"));
608 }
609
610 #[test]
611 fn test_new_default() {
612 let clay_default = ClayCode::new_default(4, 2).unwrap();
613 let clay_explicit = ClayCode::new(4, 2, 4 + 2 - 1).unwrap();
614 assert_eq!(clay_default.k, clay_explicit.k);
615 assert_eq!(clay_default.m, clay_explicit.m);
616 assert_eq!(clay_default.d, clay_explicit.d);
617 assert_eq!(clay_default.q, clay_explicit.q);
618 assert_eq!(clay_default.t, clay_explicit.t);
619 assert_eq!(clay_default.sub_chunk_no, clay_explicit.sub_chunk_no);
620 assert_eq!(clay_default.beta, clay_explicit.beta);
621
622 let clay_default2 = ClayCode::new_default(10, 4).unwrap();
624 let clay_explicit2 = ClayCode::new(10, 4, 13).unwrap();
625 assert_eq!(clay_default2.d, clay_explicit2.d);
626 assert_eq!(clay_default2.sub_chunk_no, clay_explicit2.sub_chunk_no);
627 }
628
629 #[test]
630 fn test_decode_empty_available_with_erasures() {
631 let clay = ClayCode::new(4, 2, 5).unwrap();
632 let available: HashMap<usize, Vec<u8>> = HashMap::new();
633 let result = clay.decode(&available, &[0]);
634 assert!(
635 matches!(result, Err(ClayError::InvalidParameters(_))),
636 "Expected InvalidParameters error when available is empty but erasures is non-empty, got {:?}",
637 result
638 );
639 }
640
641 #[test]
644 fn test_decode_too_many_erasures() {
645 let clay = ClayCode::new(4, 2, 5).unwrap();
646 let data: Vec<u8> = (0..128).map(|i| (i % 256) as u8).collect();
647 let chunks = clay.encode(&data);
648
649 let mut available: HashMap<usize, Vec<u8>> = HashMap::new();
651 for (i, chunk) in chunks.iter().enumerate() {
652 if i > 2 {
653 available.insert(i, chunk.clone());
654 }
655 }
656
657 let result = clay.decode(&available, &[0, 1, 2]);
658 assert!(
659 matches!(result, Err(ClayError::TooManyErasures { max: 2, actual: 3 })),
660 "Expected TooManyErasures error, got {:?}",
661 result
662 );
663 }
664
665 #[test]
666 fn test_decode_inconsistent_chunk_sizes() {
667 let clay = ClayCode::new(4, 2, 5).unwrap();
668 let data: Vec<u8> = (0..128).map(|i| (i % 256) as u8).collect();
669 let chunks = clay.encode(&data);
670
671 let mut available: HashMap<usize, Vec<u8>> = HashMap::new();
672 for (i, chunk) in chunks.iter().enumerate() {
673 if i != 0 {
674 if i == 5 {
675 let mut bad_chunk = chunk.clone();
677 bad_chunk.push(0); available.insert(i, bad_chunk);
679 } else {
680 available.insert(i, chunk.clone());
681 }
682 }
683 }
684
685 let result = clay.decode(&available, &[0]);
686 assert!(
688 matches!(result, Err(ClayError::InconsistentChunkSizes { .. }))
689 || matches!(result, Err(ClayError::InvalidChunkSize { .. })),
690 "Expected InconsistentChunkSizes or InvalidChunkSize error, got {:?}",
691 result
692 );
693 }
694
695 #[test]
696 fn test_decode_invalid_chunk_index() {
697 let clay = ClayCode::new(4, 2, 5).unwrap();
698 let data: Vec<u8> = (0..128).collect();
699 let chunks = clay.encode(&data);
700
701 let mut available: HashMap<usize, Vec<u8>> = HashMap::new();
702 for (i, chunk) in chunks.iter().enumerate() {
703 available.insert(i, chunk.clone());
704 }
705 available.insert(100, vec![0u8; chunks[0].len()]);
707
708 let result = clay.decode(&available, &[]);
709 assert!(
710 matches!(result, Err(ClayError::InvalidParameters(_))),
711 "Expected InvalidParameters error for out-of-range index, got {:?}",
712 result
713 );
714 }
715
716 #[test]
717 fn test_decode_invalid_erasure_index() {
718 let clay = ClayCode::new(4, 2, 5).unwrap();
719 let data: Vec<u8> = (0..128).collect();
720 let chunks = clay.encode(&data);
721
722 let mut available: HashMap<usize, Vec<u8>> = HashMap::new();
723 for (i, chunk) in chunks.iter().enumerate() {
724 if i != 0 {
725 available.insert(i, chunk.clone());
726 }
727 }
728
729 let result = clay.decode(&available, &[100]);
731 assert!(
732 matches!(result, Err(ClayError::InvalidParameters(_))),
733 "Expected InvalidParameters error for out-of-range erasure, got {:?}",
734 result
735 );
736 }
737
738 #[test]
739 fn test_decode_available_erasure_overlap() {
740 let clay = ClayCode::new(4, 2, 5).unwrap();
741 let data: Vec<u8> = (0..128).collect();
742 let chunks = clay.encode(&data);
743
744 let mut available: HashMap<usize, Vec<u8>> = HashMap::new();
746 for (i, chunk) in chunks.iter().enumerate() {
747 available.insert(i, chunk.clone());
748 }
749
750 let result = clay.decode(&available, &[0]);
751 assert!(
752 matches!(result, Err(ClayError::InvalidParameters(ref msg)) if msg.contains("both")),
753 "Expected InvalidParameters error for overlap, got {:?}",
754 result
755 );
756 }
757
758 #[test]
759 fn test_decode_wrong_available_count() {
760 let clay = ClayCode::new(4, 2, 5).unwrap();
761 let data: Vec<u8> = (0..128).collect();
762 let chunks = clay.encode(&data);
763
764 let mut available: HashMap<usize, Vec<u8>> = HashMap::new();
766 for (i, chunk) in chunks.iter().enumerate() {
767 if i > 1 {
768 available.insert(i, chunk.clone());
769 }
770 }
771
772 let result = clay.decode(&available, &[0]);
774 assert!(
775 matches!(result, Err(ClayError::InvalidParameters(ref msg)) if msg.contains("Expected")),
776 "Expected InvalidParameters error for wrong count, got {:?}",
777 result
778 );
779 }
780}