1use std::collections::HashMap;
42
43pub mod coords;
44pub mod decode;
45pub mod encode;
46pub mod error;
47pub mod repair;
48pub mod transforms;
49
50pub use error::ClayError;
51
52use decode::decode as decode_chunks;
53use encode::encode as encode_chunks;
54use repair::{minimum_to_repair as min_repair, repair as repair_chunk};
55
56pub struct ClayCode {
58 pub k: usize,
60 pub m: usize,
62 pub n: usize,
64 pub d: usize,
66 pub q: usize,
68 pub t: usize,
70 pub nu: usize,
72 pub sub_chunk_no: usize,
74 pub beta: usize,
76 original_count: usize,
78 recovery_count: usize,
80}
81
82impl ClayCode {
83 pub fn new(k: usize, m: usize, d: usize) -> Result<Self, ClayError> {
93 if k < 1 {
94 return Err(ClayError::InvalidParameters("k must be at least 1".into()));
95 }
96 if m < 1 {
97 return Err(ClayError::InvalidParameters("m must be at least 1".into()));
98 }
99 if d < k + 1 || d > k + m - 1 {
100 return Err(ClayError::InvalidParameters(format!(
101 "d must be in range [{}, {}], got {}",
102 k + 1,
103 k + m - 1,
104 d
105 )));
106 }
107
108 let q = d - k + 1;
109 let n = k + m;
110
111 let nu = if n % q == 0 { 0 } else { q - (n % q) };
113
114 let t = (n + nu) / q;
115
116 let sub_chunk_no = checked_pow(q, t).ok_or_else(|| {
118 ClayError::Overflow(format!("q^t = {}^{} overflows", q, t))
119 })?;
120
121 let beta = sub_chunk_no / q; let original_count = k + nu;
125 let recovery_count = m;
126 if original_count > 32768 || recovery_count > 32768 {
127 return Err(ClayError::InvalidParameters(
128 "Total nodes exceeds reed-solomon limit of 32768".into(),
129 ));
130 }
131
132 Ok(ClayCode {
133 k,
134 m,
135 n,
136 d,
137 q,
138 t,
139 nu,
140 sub_chunk_no,
141 beta,
142 original_count,
143 recovery_count,
144 })
145 }
146
147 pub fn new_default(k: usize, m: usize) -> Result<Self, ClayError> {
149 Self::new(k, m, k + m - 1)
150 }
151
152 fn encode_params(&self) -> encode::EncodeParams {
154 encode::EncodeParams {
155 k: self.k,
156 m: self.m,
157 n: self.n,
158 q: self.q,
159 t: self.t,
160 nu: self.nu,
161 sub_chunk_no: self.sub_chunk_no,
162 original_count: self.original_count,
163 recovery_count: self.recovery_count,
164 }
165 }
166
167 pub fn encode(&self, data: &[u8]) -> Vec<Vec<u8>> {
175 encode_chunks(&self.encode_params(), data)
176 }
177
178 pub fn decode(
187 &self,
188 available: &HashMap<usize, Vec<u8>>,
189 erasures: &[usize],
190 ) -> Result<Vec<u8>, ClayError> {
191 decode_chunks(&self.encode_params(), available, erasures)
192 }
193
194 pub fn minimum_to_repair(
206 &self,
207 lost_node: usize,
208 available: &[usize],
209 ) -> Result<Vec<(usize, Vec<usize>)>, ClayError> {
210 min_repair(&self.encode_params(), lost_node, available)
211 }
212
213 pub fn repair(
225 &self,
226 lost_node: usize,
227 helper_data: &HashMap<usize, Vec<u8>>,
228 chunk_size: usize,
229 ) -> Result<Vec<u8>, ClayError> {
230 repair_chunk(&self.encode_params(), lost_node, helper_data, chunk_size)
231 }
232
233 pub fn normalized_repair_bandwidth(&self) -> f64 {
238 (self.d as f64) / ((self.k as f64) * (self.d - self.k + 1) as f64)
239 }
240}
241
242fn checked_pow(base: usize, exp: usize) -> Option<usize> {
244 let mut result: usize = 1;
245 let mut b = base;
246 let mut e = exp;
247 while e > 0 {
248 if e & 1 == 1 {
249 result = result.checked_mul(b)?;
250 }
251 e >>= 1;
252 if e > 0 {
253 b = b.checked_mul(b)?;
254 }
255 }
256 Some(result)
257}
258
259#[cfg(test)]
260mod tests {
261 use super::*;
262
263 #[test]
264 fn test_basic_encode_decode() {
265 let clay = ClayCode::new(4, 2, 5).unwrap();
266 let data = b"Test data for Clay codes - not empty!";
267 let chunks = clay.encode(data);
268 assert_eq!(chunks.len(), 6); let mut available: HashMap<usize, Vec<u8>> = HashMap::new();
272 for (i, chunk) in chunks.iter().enumerate() {
273 available.insert(i, chunk.clone());
274 }
275 let decoded = clay.decode(&available, &[]).unwrap();
276
277 assert_eq!(&decoded[..data.len()], &data[..]);
279 }
280
281 #[test]
282 fn test_decode_with_erasures() {
283 let clay = ClayCode::new(4, 2, 5).unwrap();
284 let data = b"Test data for Clay codes - testing erasure recovery!";
285 let chunks = clay.encode(data);
286
287 let mut available: HashMap<usize, Vec<u8>> = HashMap::new();
289 for (i, chunk) in chunks.iter().enumerate() {
290 if i != 0 {
291 available.insert(i, chunk.clone());
292 }
293 }
294 let decoded = clay.decode(&available, &[0]).unwrap();
295 assert_eq!(&decoded[..data.len()], &data[..]);
296
297 let mut available: HashMap<usize, Vec<u8>> = HashMap::new();
299 for (i, chunk) in chunks.iter().enumerate() {
300 if i != 5 {
301 available.insert(i, chunk.clone());
302 }
303 }
304 let decoded = clay.decode(&available, &[5]).unwrap();
305 assert_eq!(&decoded[..data.len()], &data[..]);
306
307 let mut available: HashMap<usize, Vec<u8>> = HashMap::new();
309 for (i, chunk) in chunks.iter().enumerate() {
310 if i != 0 && i != 5 {
311 available.insert(i, chunk.clone());
312 }
313 }
314 let decoded = clay.decode(&available, &[0, 5]).unwrap();
315 assert_eq!(&decoded[..data.len()], &data[..]);
316 }
317
318 #[test]
319 fn test_parameters() {
320 let clay = ClayCode::new(4, 2, 5).unwrap();
322 assert_eq!(clay.q, 2);
323 assert_eq!(clay.t, 3);
324 assert_eq!(clay.sub_chunk_no, 8); assert_eq!(clay.beta, 4); let clay2 = ClayCode::new(10, 4, 13).unwrap();
329 assert_eq!(clay2.q, 4);
330 assert_eq!(clay2.t, 4);
331 assert_eq!(clay2.sub_chunk_no, 256); assert_eq!(clay2.beta, 64); }
334
335 #[test]
336 fn test_minimum_to_repair() {
337 let clay = ClayCode::new(4, 2, 5).unwrap();
338 let available: Vec<usize> = vec![1, 2, 3, 4, 5];
339 let helper_info = clay.minimum_to_repair(0, &available).unwrap();
340
341 assert_eq!(helper_info.len(), 5);
343
344 for (_, indices) in &helper_info {
346 assert_eq!(indices.len(), 4);
347 }
348 }
349
350 #[test]
351 fn test_repair_bandwidth_verification() {
352 let clay = ClayCode::new(4, 2, 5).unwrap();
354 let data = b"Test data for bandwidth verification of Clay codes repair!";
355 let chunks = clay.encode(data);
356 let chunk_size = chunks[0].len();
357
358 let available: Vec<usize> = vec![1, 2, 3, 4, 5];
360 let helper_info = clay.minimum_to_repair(0, &available).unwrap();
361
362 let sub_chunk_size = chunk_size / clay.sub_chunk_no;
364 let total_repair_subchunks: usize = helper_info
365 .iter()
366 .map(|(_, indices)| indices.len())
367 .sum();
368 let total_repair_bytes = total_repair_subchunks * sub_chunk_size;
369
370 let full_decode_bytes = clay.k * chunk_size;
371
372 let ratio = total_repair_bytes as f64 / full_decode_bytes as f64;
374 println!(
375 "Repair bandwidth: {} bytes, Full decode: {} bytes, Ratio: {:.3}",
376 total_repair_bytes, full_decode_bytes, ratio
377 );
378
379 assert!(
380 total_repair_bytes < full_decode_bytes * 7 / 10,
381 "Repair bandwidth {} should be < 70% of full decode {}",
382 total_repair_bytes,
383 full_decode_bytes
384 );
385 }
386
387 #[test]
388 fn test_repair_correctness() {
389 let clay = ClayCode::new(4, 2, 5).unwrap();
390 let data = b"Test data for repair correctness verification!!!!";
391 let chunks = clay.encode(data);
392 let chunk_size = chunks[0].len();
393 let sub_chunk_size = chunk_size / clay.sub_chunk_no;
394
395 for lost_node in 0..clay.n {
397 let available: Vec<usize> = (0..clay.n).filter(|&i| i != lost_node).collect();
398 let helper_info = clay.minimum_to_repair(lost_node, &available).unwrap();
399
400 let mut partial_data: HashMap<usize, Vec<u8>> = HashMap::new();
402 for (helper_idx, indices) in &helper_info {
403 let mut helper_partial = Vec::new();
404 for &sc_idx in indices {
405 let start_byte = sc_idx * sub_chunk_size;
406 let end_byte = (sc_idx + 1) * sub_chunk_size;
407 helper_partial.extend_from_slice(&chunks[*helper_idx][start_byte..end_byte]);
408 }
409 partial_data.insert(*helper_idx, helper_partial);
410 }
411
412 let recovered = clay.repair(lost_node, &partial_data, chunk_size).unwrap();
414
415 assert_eq!(
417 recovered, chunks[lost_node],
418 "Repair failed for node {}",
419 lost_node
420 );
421 }
422 }
423
424 #[test]
425 fn test_various_parameters() {
426 let params = vec![
428 (4, 2, 5), (9, 3, 11), (10, 4, 13), ];
432
433 for (k, m, d) in params {
434 let clay = ClayCode::new(k, m, d).unwrap();
435 let data_size = k * clay.sub_chunk_no * 2;
436 let data: Vec<u8> = (0..data_size).map(|i| (i % 256) as u8).collect();
437 let chunks = clay.encode(&data);
438
439 let mut available: HashMap<usize, Vec<u8>> = HashMap::new();
441 for (i, chunk) in chunks.iter().enumerate() {
442 if i != 0 {
443 available.insert(i, chunk.clone());
444 }
445 }
446 let decoded = clay.decode(&available, &[0]).unwrap();
447 assert_eq!(
448 &decoded[..data.len()],
449 &data[..],
450 "Failed for params ({}, {}, {})",
451 k,
452 m,
453 d
454 );
455 }
456 }
457
458 #[test]
459 fn test_repair_all_nodes_various_params() {
460 let params = vec![(4, 2, 5), (9, 3, 11)];
461
462 for (k, m, d) in params {
463 let clay = ClayCode::new(k, m, d).unwrap();
464 let data_size = k * clay.sub_chunk_no;
465 let data: Vec<u8> = (0..data_size).map(|i| ((i * 7 + 13) % 256) as u8).collect();
466 let chunks = clay.encode(&data);
467 let chunk_size = chunks[0].len();
468 let sub_chunk_size = chunk_size / clay.sub_chunk_no;
469
470 for lost_node in 0..clay.n {
471 let available: Vec<usize> = (0..clay.n).filter(|&i| i != lost_node).collect();
472 let helper_info = clay.minimum_to_repair(lost_node, &available).unwrap();
473
474 let mut partial_data: HashMap<usize, Vec<u8>> = HashMap::new();
475 for (helper_idx, indices) in &helper_info {
476 let mut helper_partial = Vec::new();
477 for &sc_idx in indices {
478 let start_byte = sc_idx * sub_chunk_size;
479 let end_byte = (sc_idx + 1) * sub_chunk_size;
480 helper_partial.extend_from_slice(&chunks[*helper_idx][start_byte..end_byte]);
481 }
482 partial_data.insert(*helper_idx, helper_partial);
483 }
484
485 let recovered = clay.repair(lost_node, &partial_data, chunk_size).unwrap();
486 assert_eq!(
487 recovered, chunks[lost_node],
488 "Repair failed for node {} with params ({}, {}, {})",
489 lost_node, k, m, d
490 );
491 }
492 }
493 }
494
495 #[test]
496 fn test_decode_max_erasures() {
497 let clay = ClayCode::new(4, 2, 5).unwrap();
498 let data: Vec<u8> = (0..256).map(|i| (i % 256) as u8).collect();
499 let chunks = clay.encode(&data);
500
501 let patterns = vec![vec![0, 5], vec![0, 1], vec![4, 5], vec![1, 3]];
503
504 for erasures in patterns {
505 let mut available: HashMap<usize, Vec<u8>> = HashMap::new();
506 for (i, chunk) in chunks.iter().enumerate() {
507 if !erasures.contains(&i) {
508 available.insert(i, chunk.clone());
509 }
510 }
511 let decoded = clay.decode(&available, &erasures).unwrap();
512 assert_eq!(
513 &decoded[..data.len()],
514 &data[..],
515 "Failed for erasures {:?}",
516 erasures
517 );
518 }
519 }
520
521 #[test]
522 fn test_normalized_repair_bandwidth() {
523 let test_cases = vec![
524 ((4, 2, 5), 0.625),
525 ((9, 3, 11), 0.407),
526 ((10, 4, 13), 0.325),
527 ];
528
529 for ((k, m, d), expected) in test_cases {
530 let clay = ClayCode::new(k, m, d).unwrap();
531 let actual = clay.normalized_repair_bandwidth();
532 assert!(
533 (actual - expected).abs() < 0.01,
534 "Expected {}, got {} for ({}, {}, {})",
535 expected,
536 actual,
537 k,
538 m,
539 d
540 );
541 }
542 }
543
544 #[test]
545 fn test_random_data() {
546 use rand::Rng;
547 let mut rng = rand::thread_rng();
548
549 let clay = ClayCode::new(4, 2, 5).unwrap();
550 let data_size = clay.k * clay.sub_chunk_no * 4;
551 let data: Vec<u8> = (0..data_size).map(|_| rng.gen()).collect();
552 let chunks = clay.encode(&data);
553
554 let mut available: HashMap<usize, Vec<u8>> = HashMap::new();
556 for (i, chunk) in chunks.iter().enumerate() {
557 available.insert(i, chunk.clone());
558 }
559 let decoded = clay.decode(&available, &[]).unwrap();
560 assert_eq!(&decoded[..data.len()], &data[..]);
561
562 let mut available: HashMap<usize, Vec<u8>> = HashMap::new();
564 for (i, chunk) in chunks.iter().enumerate() {
565 if i != 2 {
566 available.insert(i, chunk.clone());
567 }
568 }
569 let decoded = clay.decode(&available, &[2]).unwrap();
570 assert_eq!(&decoded[..data.len()], &data[..]);
571 }
572
573 #[test]
574 fn test_checked_pow_overflow() {
575 assert!(checked_pow(2, 63).is_some());
577 assert!(checked_pow(2, 64).is_none()); assert!(checked_pow(10, 20).is_none()); }
580
581 #[test]
582 fn test_invalid_parameters() {
583 assert!(ClayCode::new(0, 2, 1).is_err());
585
586 assert!(ClayCode::new(4, 0, 3).is_err());
588
589 assert!(ClayCode::new(4, 2, 4).is_err()); assert!(ClayCode::new(4, 2, 6).is_err()); }
593
594 #[test]
597 fn test_decode_too_many_erasures() {
598 let clay = ClayCode::new(4, 2, 5).unwrap();
599 let data: Vec<u8> = (0..128).map(|i| (i % 256) as u8).collect();
600 let chunks = clay.encode(&data);
601
602 let mut available: HashMap<usize, Vec<u8>> = HashMap::new();
604 for (i, chunk) in chunks.iter().enumerate() {
605 if i > 2 {
606 available.insert(i, chunk.clone());
607 }
608 }
609
610 let result = clay.decode(&available, &[0, 1, 2]);
611 assert!(
612 matches!(result, Err(ClayError::TooManyErasures { max: 2, actual: 3 })),
613 "Expected TooManyErasures error, got {:?}",
614 result
615 );
616 }
617
618 #[test]
619 fn test_decode_inconsistent_chunk_sizes() {
620 let clay = ClayCode::new(4, 2, 5).unwrap();
621 let data: Vec<u8> = (0..128).map(|i| (i % 256) as u8).collect();
622 let chunks = clay.encode(&data);
623
624 let mut available: HashMap<usize, Vec<u8>> = HashMap::new();
625 for (i, chunk) in chunks.iter().enumerate() {
626 if i != 0 {
627 if i == 5 {
628 let mut bad_chunk = chunk.clone();
630 bad_chunk.push(0); available.insert(i, bad_chunk);
632 } else {
633 available.insert(i, chunk.clone());
634 }
635 }
636 }
637
638 let result = clay.decode(&available, &[0]);
639 assert!(
641 matches!(result, Err(ClayError::InconsistentChunkSizes { .. }))
642 || matches!(result, Err(ClayError::InvalidChunkSize { .. })),
643 "Expected InconsistentChunkSizes or InvalidChunkSize error, got {:?}",
644 result
645 );
646 }
647
648 #[test]
649 fn test_decode_invalid_chunk_index() {
650 let clay = ClayCode::new(4, 2, 5).unwrap();
651 let data: Vec<u8> = (0..128).collect();
652 let chunks = clay.encode(&data);
653
654 let mut available: HashMap<usize, Vec<u8>> = HashMap::new();
655 for (i, chunk) in chunks.iter().enumerate() {
656 available.insert(i, chunk.clone());
657 }
658 available.insert(100, vec![0u8; chunks[0].len()]);
660
661 let result = clay.decode(&available, &[]);
662 assert!(
663 matches!(result, Err(ClayError::InvalidParameters(_))),
664 "Expected InvalidParameters error for out-of-range index, got {:?}",
665 result
666 );
667 }
668
669 #[test]
670 fn test_decode_invalid_erasure_index() {
671 let clay = ClayCode::new(4, 2, 5).unwrap();
672 let data: Vec<u8> = (0..128).collect();
673 let chunks = clay.encode(&data);
674
675 let mut available: HashMap<usize, Vec<u8>> = HashMap::new();
676 for (i, chunk) in chunks.iter().enumerate() {
677 if i != 0 {
678 available.insert(i, chunk.clone());
679 }
680 }
681
682 let result = clay.decode(&available, &[100]);
684 assert!(
685 matches!(result, Err(ClayError::InvalidParameters(_))),
686 "Expected InvalidParameters error for out-of-range erasure, got {:?}",
687 result
688 );
689 }
690
691 #[test]
692 fn test_decode_available_erasure_overlap() {
693 let clay = ClayCode::new(4, 2, 5).unwrap();
694 let data: Vec<u8> = (0..128).collect();
695 let chunks = clay.encode(&data);
696
697 let mut available: HashMap<usize, Vec<u8>> = HashMap::new();
699 for (i, chunk) in chunks.iter().enumerate() {
700 available.insert(i, chunk.clone());
701 }
702
703 let result = clay.decode(&available, &[0]);
704 assert!(
705 matches!(result, Err(ClayError::InvalidParameters(ref msg)) if msg.contains("both")),
706 "Expected InvalidParameters error for overlap, got {:?}",
707 result
708 );
709 }
710
711 #[test]
712 fn test_decode_wrong_available_count() {
713 let clay = ClayCode::new(4, 2, 5).unwrap();
714 let data: Vec<u8> = (0..128).collect();
715 let chunks = clay.encode(&data);
716
717 let mut available: HashMap<usize, Vec<u8>> = HashMap::new();
719 for (i, chunk) in chunks.iter().enumerate() {
720 if i > 1 {
721 available.insert(i, chunk.clone());
722 }
723 }
724
725 let result = clay.decode(&available, &[0]);
727 assert!(
728 matches!(result, Err(ClayError::InvalidParameters(ref msg)) if msg.contains("Expected")),
729 "Expected InvalidParameters error for wrong count, got {:?}",
730 result
731 );
732 }
733}