s3p_core/
rs.rs

1use crate::errors::*;
2use reed_solomon_erasure::galois_8::ReedSolomon;
3
4/// Encode `data_shards` into `data_shards + parity_shards` shards.
5pub fn rs_encode(data: &[u8], data_shards: usize, parity_shards: usize) -> Result<Vec<Vec<u8>>> {
6    let shard_size = data.len().div_ceil(data_shards);
7    let total_shards = data_shards + parity_shards;
8
9    // data shards (padded) + parity shards (zero-initialized)
10    let mut shards: Vec<Vec<u8>> = Vec::with_capacity(total_shards);
11    for i in 0..data_shards {
12        let start = i * shard_size;
13        let end = (start + shard_size).min(data.len());
14        let mut shard = vec![0u8; shard_size];
15        if start < data.len() {
16            shard[..(end - start)].copy_from_slice(&data[start..end]);
17        }
18        shards.push(shard);
19    }
20    for _ in 0..parity_shards {
21        shards.push(vec![0u8; shard_size]);
22    }
23
24    // encode
25    let r =
26        ReedSolomon::new(data_shards, parity_shards).map_err(|e| S3pError::Rs(format!("{e}")))?;
27    r.encode(&mut shards.iter_mut().map(|v| &mut v[..]).collect::<Vec<_>>())
28        .map_err(|e| S3pError::Rs(format!("{e}")))?;
29
30    Ok(shards)
31}
32
33/// Attempt to reconstruct data from a subset of shards (some can be None).
34pub fn rs_reconstruct(
35    shards_in: Vec<Option<Vec<u8>>>,
36    data_shards: usize,
37    parity_shards: usize,
38) -> Result<Vec<u8>> {
39    let total_shards = data_shards + parity_shards;
40    if shards_in.len() != total_shards {
41        return Err(S3pError::Invalid("wrong shard count".into()));
42    }
43
44    // detect shard size from the first present shard
45    let shard_size = shards_in
46        .iter()
47        .filter_map(|o| o.as_ref().map(|v| v.len()))
48        .next()
49        .ok_or_else(|| S3pError::Invalid("no shards".into()))?;
50
51    // materialize buffers and remember presence flags
52    let mut shards: Vec<Vec<u8>> = Vec::with_capacity(total_shards);
53    let mut present: Vec<bool> = Vec::with_capacity(total_shards);
54
55    for opt in shards_in {
56        match opt {
57            Some(mut v) => {
58                if v.len() < shard_size {
59                    v.resize(shard_size, 0);
60                } else if v.len() > shard_size {
61                    v.truncate(shard_size);
62                }
63                present.push(true);
64                shards.push(v);
65            }
66            None => {
67                present.push(false);
68                shards.push(vec![0u8; shard_size]); // buffer to be filled by reconstruct
69            }
70        }
71    }
72
73    // (&mut [u8], bool) where bool = was_present
74    let mut tuples: Vec<(&mut [u8], bool)> = shards
75        .iter_mut()
76        .zip(present.iter().copied())
77        .map(|(v, was_present)| (&mut v[..], was_present))
78        .collect();
79
80    let r =
81        ReedSolomon::new(data_shards, parity_shards).map_err(|e| S3pError::Rs(format!("{e}")))?;
82    r.reconstruct(&mut tuples)
83        .map_err(|e| S3pError::Rs(format!("{e}")))?;
84
85    // join only data shards (first data_shards)
86    let mut out = Vec::with_capacity(data_shards * shard_size);
87    for s in shards.iter().take(data_shards) {
88        out.extend_from_slice(s);
89    }
90    Ok(out)
91}