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;
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 recovery_count: usize,
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 Ok(ClayCode {
135 k,
136 m,
137 n,
138 d,
139 q,
140 t,
141 nu,
142 sub_chunk_no,
143 beta,
144 original_count,
145 recovery_count,
146 })
147 }
148
149 pub fn new_default(k: usize, m: usize) -> Result<Self, ClayError> {
151 Self::new(k, m, k + m - 1)
152 }
153
154 fn encode_params(&self) -> encode::EncodeParams {
156 encode::EncodeParams {
157 k: self.k,
158 m: self.m,
159 n: self.n,
160 q: self.q,
161 t: self.t,
162 nu: self.nu,
163 sub_chunk_no: self.sub_chunk_no,
164 original_count: self.original_count,
165 recovery_count: self.recovery_count,
166 }
167 }
168
169 pub fn encode(&self, data: &[u8]) -> Vec<Vec<u8>> {
177 encode_chunks(&self.encode_params(), data)
178 }
179
180 pub fn decode(
189 &self,
190 available: &HashMap<usize, Vec<u8>>,
191 erasures: &[usize],
192 ) -> Result<Vec<u8>, ClayError> {
193 decode_chunks(&self.encode_params(), available, erasures)
194 }
195
196 pub fn minimum_to_repair(
208 &self,
209 lost_node: usize,
210 available: &[usize],
211 ) -> Result<Vec<(usize, Vec<usize>)>, ClayError> {
212 min_repair(&self.encode_params(), lost_node, available)
213 }
214
215 pub fn repair(
227 &self,
228 lost_node: usize,
229 helper_data: &HashMap<usize, Vec<u8>>,
230 chunk_size: usize,
231 ) -> Result<Vec<u8>, ClayError> {
232 repair_chunk(&self.encode_params(), lost_node, helper_data, chunk_size)
233 }
234
235 pub fn normalized_repair_bandwidth(&self) -> f64 {
240 (self.d as f64) / ((self.k as f64) * (self.d - self.k + 1) as f64)
241 }
242}
243
244fn checked_pow(base: usize, exp: usize) -> Option<usize> {
246 let mut result: usize = 1;
247 let mut b = base;
248 let mut e = exp;
249 while e > 0 {
250 if e & 1 == 1 {
251 result = result.checked_mul(b)?;
252 }
253 e >>= 1;
254 if e > 0 {
255 b = b.checked_mul(b)?;
256 }
257 }
258 Some(result)
259}
260
261#[cfg(test)]
262mod tests {
263 use super::*;
264
265 #[test]
266 fn test_basic_encode_decode() {
267 let clay = ClayCode::new(4, 2, 5).unwrap();
268 let data = b"Test data for Clay codes - not empty!";
269 let chunks = clay.encode(data);
270 assert_eq!(chunks.len(), 6); let mut available: HashMap<usize, Vec<u8>> = HashMap::new();
274 for (i, chunk) in chunks.iter().enumerate() {
275 available.insert(i, chunk.clone());
276 }
277 let decoded = clay.decode(&available, &[]).unwrap();
278
279 assert_eq!(&decoded[..data.len()], &data[..]);
281 }
282
283 #[test]
284 fn test_decode_with_erasures() {
285 let clay = ClayCode::new(4, 2, 5).unwrap();
286 let data = b"Test data for Clay codes - testing erasure recovery!";
287 let chunks = clay.encode(data);
288
289 let mut available: HashMap<usize, Vec<u8>> = HashMap::new();
291 for (i, chunk) in chunks.iter().enumerate() {
292 if i != 0 {
293 available.insert(i, chunk.clone());
294 }
295 }
296 let decoded = clay.decode(&available, &[0]).unwrap();
297 assert_eq!(&decoded[..data.len()], &data[..]);
298
299 let mut available: HashMap<usize, Vec<u8>> = HashMap::new();
301 for (i, chunk) in chunks.iter().enumerate() {
302 if i != 5 {
303 available.insert(i, chunk.clone());
304 }
305 }
306 let decoded = clay.decode(&available, &[5]).unwrap();
307 assert_eq!(&decoded[..data.len()], &data[..]);
308
309 let mut available: HashMap<usize, Vec<u8>> = HashMap::new();
311 for (i, chunk) in chunks.iter().enumerate() {
312 if i != 0 && i != 5 {
313 available.insert(i, chunk.clone());
314 }
315 }
316 let decoded = clay.decode(&available, &[0, 5]).unwrap();
317 assert_eq!(&decoded[..data.len()], &data[..]);
318 }
319
320 #[test]
321 fn test_parameters() {
322 let clay = ClayCode::new(4, 2, 5).unwrap();
324 assert_eq!(clay.q, 2);
325 assert_eq!(clay.t, 3);
326 assert_eq!(clay.sub_chunk_no, 8); assert_eq!(clay.beta, 4); let clay2 = ClayCode::new(10, 4, 13).unwrap();
331 assert_eq!(clay2.q, 4);
332 assert_eq!(clay2.t, 4);
333 assert_eq!(clay2.sub_chunk_no, 256); assert_eq!(clay2.beta, 64); }
336
337 #[test]
338 fn test_minimum_to_repair() {
339 let clay = ClayCode::new(4, 2, 5).unwrap();
340 let available: Vec<usize> = vec![1, 2, 3, 4, 5];
341 let helper_info = clay.minimum_to_repair(0, &available).unwrap();
342
343 assert_eq!(helper_info.len(), 5);
345
346 for (_, indices) in &helper_info {
348 assert_eq!(indices.len(), 4);
349 }
350 }
351
352 #[test]
353 fn test_repair_bandwidth_verification() {
354 let clay = ClayCode::new(4, 2, 5).unwrap();
356 let data = b"Test data for bandwidth verification of Clay codes repair!";
357 let chunks = clay.encode(data);
358 let chunk_size = chunks[0].len();
359
360 let available: Vec<usize> = vec![1, 2, 3, 4, 5];
362 let helper_info = clay.minimum_to_repair(0, &available).unwrap();
363
364 let sub_chunk_size = chunk_size / clay.sub_chunk_no;
366 let total_repair_subchunks: usize = helper_info
367 .iter()
368 .map(|(_, indices)| indices.len())
369 .sum();
370 let total_repair_bytes = total_repair_subchunks * sub_chunk_size;
371
372 let full_decode_bytes = clay.k * chunk_size;
373
374 let ratio = total_repair_bytes as f64 / full_decode_bytes as f64;
376 println!(
377 "Repair bandwidth: {} bytes, Full decode: {} bytes, Ratio: {:.3}",
378 total_repair_bytes, full_decode_bytes, ratio
379 );
380
381 assert!(
382 total_repair_bytes < full_decode_bytes * 7 / 10,
383 "Repair bandwidth {} should be < 70% of full decode {}",
384 total_repair_bytes,
385 full_decode_bytes
386 );
387 }
388
389 #[test]
390 fn test_repair_correctness() {
391 let clay = ClayCode::new(4, 2, 5).unwrap();
392 let data = b"Test data for repair correctness verification!!!!";
393 let chunks = clay.encode(data);
394 let chunk_size = chunks[0].len();
395 let sub_chunk_size = chunk_size / clay.sub_chunk_no;
396
397 for lost_node in 0..clay.n {
399 let available: Vec<usize> = (0..clay.n).filter(|&i| i != lost_node).collect();
400 let helper_info = clay.minimum_to_repair(lost_node, &available).unwrap();
401
402 let mut partial_data: HashMap<usize, Vec<u8>> = HashMap::new();
404 for (helper_idx, indices) in &helper_info {
405 let mut helper_partial = Vec::new();
406 for &sc_idx in indices {
407 let start_byte = sc_idx * sub_chunk_size;
408 let end_byte = (sc_idx + 1) * sub_chunk_size;
409 helper_partial.extend_from_slice(&chunks[*helper_idx][start_byte..end_byte]);
410 }
411 partial_data.insert(*helper_idx, helper_partial);
412 }
413
414 let recovered = clay.repair(lost_node, &partial_data, chunk_size).unwrap();
416
417 assert_eq!(
419 recovered, chunks[lost_node],
420 "Repair failed for node {}",
421 lost_node
422 );
423 }
424 }
425
426 #[test]
427 fn test_various_parameters() {
428 let params = vec![
430 (4, 2, 5), (9, 3, 11), (10, 4, 13), ];
434
435 for (k, m, d) in params {
436 let clay = ClayCode::new(k, m, d).unwrap();
437 let data_size = k * clay.sub_chunk_no * 2;
438 let data: Vec<u8> = (0..data_size).map(|i| (i % 256) as u8).collect();
439 let chunks = clay.encode(&data);
440
441 let mut available: HashMap<usize, Vec<u8>> = HashMap::new();
443 for (i, chunk) in chunks.iter().enumerate() {
444 if i != 0 {
445 available.insert(i, chunk.clone());
446 }
447 }
448 let decoded = clay.decode(&available, &[0]).unwrap();
449 assert_eq!(
450 &decoded[..data.len()],
451 &data[..],
452 "Failed for params ({}, {}, {})",
453 k,
454 m,
455 d
456 );
457 }
458 }
459
460 #[test]
461 fn test_repair_all_nodes_various_params() {
462 let params = vec![(4, 2, 5), (9, 3, 11)];
463
464 for (k, m, d) in params {
465 let clay = ClayCode::new(k, m, d).unwrap();
466 let data_size = k * clay.sub_chunk_no;
467 let data: Vec<u8> = (0..data_size).map(|i| ((i * 7 + 13) % 256) as u8).collect();
468 let chunks = clay.encode(&data);
469 let chunk_size = chunks[0].len();
470 let sub_chunk_size = chunk_size / clay.sub_chunk_no;
471
472 for lost_node in 0..clay.n {
473 let available: Vec<usize> = (0..clay.n).filter(|&i| i != lost_node).collect();
474 let helper_info = clay.minimum_to_repair(lost_node, &available).unwrap();
475
476 let mut partial_data: HashMap<usize, Vec<u8>> = HashMap::new();
477 for (helper_idx, indices) in &helper_info {
478 let mut helper_partial = Vec::new();
479 for &sc_idx in indices {
480 let start_byte = sc_idx * sub_chunk_size;
481 let end_byte = (sc_idx + 1) * sub_chunk_size;
482 helper_partial.extend_from_slice(&chunks[*helper_idx][start_byte..end_byte]);
483 }
484 partial_data.insert(*helper_idx, helper_partial);
485 }
486
487 let recovered = clay.repair(lost_node, &partial_data, chunk_size).unwrap();
488 assert_eq!(
489 recovered, chunks[lost_node],
490 "Repair failed for node {} with params ({}, {}, {})",
491 lost_node, k, m, d
492 );
493 }
494 }
495 }
496
497 #[test]
498 fn test_decode_max_erasures() {
499 let clay = ClayCode::new(4, 2, 5).unwrap();
500 let data: Vec<u8> = (0..256).map(|i| (i % 256) as u8).collect();
501 let chunks = clay.encode(&data);
502
503 let patterns = vec![vec![0, 5], vec![0, 1], vec![4, 5], vec![1, 3]];
505
506 for erasures in patterns {
507 let mut available: HashMap<usize, Vec<u8>> = HashMap::new();
508 for (i, chunk) in chunks.iter().enumerate() {
509 if !erasures.contains(&i) {
510 available.insert(i, chunk.clone());
511 }
512 }
513 let decoded = clay.decode(&available, &erasures).unwrap();
514 assert_eq!(
515 &decoded[..data.len()],
516 &data[..],
517 "Failed for erasures {:?}",
518 erasures
519 );
520 }
521 }
522
523 #[test]
524 fn test_normalized_repair_bandwidth() {
525 let test_cases = vec![
526 ((4, 2, 5), 0.625),
527 ((9, 3, 11), 0.407),
528 ((10, 4, 13), 0.325),
529 ];
530
531 for ((k, m, d), expected) in test_cases {
532 let clay = ClayCode::new(k, m, d).unwrap();
533 let actual = clay.normalized_repair_bandwidth();
534 assert!(
535 (actual - expected).abs() < 0.01,
536 "Expected {}, got {} for ({}, {}, {})",
537 expected,
538 actual,
539 k,
540 m,
541 d
542 );
543 }
544 }
545
546 #[test]
547 fn test_random_data() {
548 use rand::Rng;
549 let mut rng = rand::thread_rng();
550
551 let clay = ClayCode::new(4, 2, 5).unwrap();
552 let data_size = clay.k * clay.sub_chunk_no * 4;
553 let data: Vec<u8> = (0..data_size).map(|_| rng.gen()).collect();
554 let chunks = clay.encode(&data);
555
556 let mut available: HashMap<usize, Vec<u8>> = HashMap::new();
558 for (i, chunk) in chunks.iter().enumerate() {
559 available.insert(i, chunk.clone());
560 }
561 let decoded = clay.decode(&available, &[]).unwrap();
562 assert_eq!(&decoded[..data.len()], &data[..]);
563
564 let mut available: HashMap<usize, Vec<u8>> = HashMap::new();
566 for (i, chunk) in chunks.iter().enumerate() {
567 if i != 2 {
568 available.insert(i, chunk.clone());
569 }
570 }
571 let decoded = clay.decode(&available, &[2]).unwrap();
572 assert_eq!(&decoded[..data.len()], &data[..]);
573 }
574
575 #[test]
576 fn test_checked_pow_overflow() {
577 assert!(checked_pow(2, 63).is_some());
579 assert!(checked_pow(2, 64).is_none()); assert!(checked_pow(10, 20).is_none()); }
582
583 #[test]
584 fn test_invalid_parameters() {
585 assert!(ClayCode::new(0, 2, 1).is_err());
587
588 assert!(ClayCode::new(4, 0, 3).is_err());
590
591 assert!(ClayCode::new(4, 2, 4).is_err()); assert!(ClayCode::new(4, 2, 6).is_err()); }
595
596 #[test]
597 fn test_clone_and_debug() {
598 let clay = ClayCode::new(4, 2, 5).unwrap();
599 let clay2 = clay.clone();
600 assert_eq!(clay2.k, clay.k);
601 assert_eq!(clay2.m, clay.m);
602 assert_eq!(clay2.d, clay.d);
603 let debug_str = format!("{:?}", clay);
605 assert!(debug_str.contains("ClayCode"));
606 }
607
608 #[test]
609 fn test_new_default() {
610 let clay_default = ClayCode::new_default(4, 2).unwrap();
611 let clay_explicit = ClayCode::new(4, 2, 4 + 2 - 1).unwrap();
612 assert_eq!(clay_default.k, clay_explicit.k);
613 assert_eq!(clay_default.m, clay_explicit.m);
614 assert_eq!(clay_default.d, clay_explicit.d);
615 assert_eq!(clay_default.q, clay_explicit.q);
616 assert_eq!(clay_default.t, clay_explicit.t);
617 assert_eq!(clay_default.sub_chunk_no, clay_explicit.sub_chunk_no);
618 assert_eq!(clay_default.beta, clay_explicit.beta);
619
620 let clay_default2 = ClayCode::new_default(10, 4).unwrap();
622 let clay_explicit2 = ClayCode::new(10, 4, 13).unwrap();
623 assert_eq!(clay_default2.d, clay_explicit2.d);
624 assert_eq!(clay_default2.sub_chunk_no, clay_explicit2.sub_chunk_no);
625 }
626
627 #[test]
628 fn test_decode_empty_available_with_erasures() {
629 let clay = ClayCode::new(4, 2, 5).unwrap();
630 let available: HashMap<usize, Vec<u8>> = HashMap::new();
631 let result = clay.decode(&available, &[0]);
632 assert!(
633 matches!(result, Err(ClayError::InvalidParameters(_))),
634 "Expected InvalidParameters error when available is empty but erasures is non-empty, got {:?}",
635 result
636 );
637 }
638
639 #[test]
642 fn test_decode_too_many_erasures() {
643 let clay = ClayCode::new(4, 2, 5).unwrap();
644 let data: Vec<u8> = (0..128).map(|i| (i % 256) as u8).collect();
645 let chunks = clay.encode(&data);
646
647 let mut available: HashMap<usize, Vec<u8>> = HashMap::new();
649 for (i, chunk) in chunks.iter().enumerate() {
650 if i > 2 {
651 available.insert(i, chunk.clone());
652 }
653 }
654
655 let result = clay.decode(&available, &[0, 1, 2]);
656 assert!(
657 matches!(result, Err(ClayError::TooManyErasures { max: 2, actual: 3 })),
658 "Expected TooManyErasures error, got {:?}",
659 result
660 );
661 }
662
663 #[test]
664 fn test_decode_inconsistent_chunk_sizes() {
665 let clay = ClayCode::new(4, 2, 5).unwrap();
666 let data: Vec<u8> = (0..128).map(|i| (i % 256) as u8).collect();
667 let chunks = clay.encode(&data);
668
669 let mut available: HashMap<usize, Vec<u8>> = HashMap::new();
670 for (i, chunk) in chunks.iter().enumerate() {
671 if i != 0 {
672 if i == 5 {
673 let mut bad_chunk = chunk.clone();
675 bad_chunk.push(0); available.insert(i, bad_chunk);
677 } else {
678 available.insert(i, chunk.clone());
679 }
680 }
681 }
682
683 let result = clay.decode(&available, &[0]);
684 assert!(
686 matches!(result, Err(ClayError::InconsistentChunkSizes { .. }))
687 || matches!(result, Err(ClayError::InvalidChunkSize { .. })),
688 "Expected InconsistentChunkSizes or InvalidChunkSize error, got {:?}",
689 result
690 );
691 }
692
693 #[test]
694 fn test_decode_invalid_chunk_index() {
695 let clay = ClayCode::new(4, 2, 5).unwrap();
696 let data: Vec<u8> = (0..128).collect();
697 let chunks = clay.encode(&data);
698
699 let mut available: HashMap<usize, Vec<u8>> = HashMap::new();
700 for (i, chunk) in chunks.iter().enumerate() {
701 available.insert(i, chunk.clone());
702 }
703 available.insert(100, vec![0u8; chunks[0].len()]);
705
706 let result = clay.decode(&available, &[]);
707 assert!(
708 matches!(result, Err(ClayError::InvalidParameters(_))),
709 "Expected InvalidParameters error for out-of-range index, got {:?}",
710 result
711 );
712 }
713
714 #[test]
715 fn test_decode_invalid_erasure_index() {
716 let clay = ClayCode::new(4, 2, 5).unwrap();
717 let data: Vec<u8> = (0..128).collect();
718 let chunks = clay.encode(&data);
719
720 let mut available: HashMap<usize, Vec<u8>> = HashMap::new();
721 for (i, chunk) in chunks.iter().enumerate() {
722 if i != 0 {
723 available.insert(i, chunk.clone());
724 }
725 }
726
727 let result = clay.decode(&available, &[100]);
729 assert!(
730 matches!(result, Err(ClayError::InvalidParameters(_))),
731 "Expected InvalidParameters error for out-of-range erasure, got {:?}",
732 result
733 );
734 }
735
736 #[test]
737 fn test_decode_available_erasure_overlap() {
738 let clay = ClayCode::new(4, 2, 5).unwrap();
739 let data: Vec<u8> = (0..128).collect();
740 let chunks = clay.encode(&data);
741
742 let mut available: HashMap<usize, Vec<u8>> = HashMap::new();
744 for (i, chunk) in chunks.iter().enumerate() {
745 available.insert(i, chunk.clone());
746 }
747
748 let result = clay.decode(&available, &[0]);
749 assert!(
750 matches!(result, Err(ClayError::InvalidParameters(ref msg)) if msg.contains("both")),
751 "Expected InvalidParameters error for overlap, got {:?}",
752 result
753 );
754 }
755
756 #[test]
757 fn test_decode_wrong_available_count() {
758 let clay = ClayCode::new(4, 2, 5).unwrap();
759 let data: Vec<u8> = (0..128).collect();
760 let chunks = clay.encode(&data);
761
762 let mut available: HashMap<usize, Vec<u8>> = HashMap::new();
764 for (i, chunk) in chunks.iter().enumerate() {
765 if i > 1 {
766 available.insert(i, chunk.clone());
767 }
768 }
769
770 let result = clay.decode(&available, &[0]);
772 assert!(
773 matches!(result, Err(ClayError::InvalidParameters(ref msg)) if msg.contains("Expected")),
774 "Expected InvalidParameters error for wrong count, got {:?}",
775 result
776 );
777 }
778}