erasure.rs

use anyhow::Result;
use log::debug;
use reed_solomon_erasure::galois_8::ReedSolomon;
use thiserror::Error;

#[derive(Error, Debug)]
pub enum CoderError {
    #[error("invalid configuration of data and parity shards: {0} {1}")]
    InvalidConfiguration(usize, usize),
    #[error("failed to encode data")]
    FailedToEncodeData,
    #[error("invalid shard data: {0}")]
    InvalidShardData(String),
    #[error("failed to verify sharded data: {0}")]
    FailedToVerify(String),
}

#[derive(Debug, Clone)]
pub struct Coder {
    pub rs: ReedSolomon,
}

pub type ShardsData<T> = Vec<T>;

impl Coder {
    pub fn new(data_shards: usize, parity_shards: usize) -> Result<Self> {
        let rs = ReedSolomon::new(data_shards, parity_shards)
            .map_err(|_| CoderError::InvalidConfiguration(data_shards, parity_shards))?;
        Ok(Self { rs })
    }
}

pub fn generate_shards<T>(ec: &Coder, input_data: &[u8]) -> Result<ShardsData<T>>
where
    T: From<Vec<u8>> + Clone,
    T: AsRef<[u8]> + AsMut<[u8]>,
    std::vec::Vec<T>: std::iter::FromIterator<std::vec::Vec<u8>>,
{
    let shard_size = input_data.len() / ec.rs.data_shard_count();
    let total_shards = ec.rs.data_shard_count() + ec.rs.parity_shard_count();

    let mut shards: ShardsData<T> = (0..total_shards)
        .map(|i| {
            if i < ec.rs.data_shard_count() {
                let start = i * shard_size;
                input_data[start..start + shard_size].to_vec()
            } else {
                vec![0; shard_size]
            }
        })
        .collect();

    ec.rs
        .encode(&mut shards)
        .map_err(|_| CoderError::FailedToEncodeData)?;

    Ok(shards)
}

pub fn reconstruct<T>(ec: &Coder, shards: &[&[u8]]) -> Result<ShardsData<T>>
where
    T: From<Vec<u8>> + Clone,
    T: AsRef<[u8]> + AsMut<[u8]>,
    std::vec::Vec<T>: std::iter::FromIterator<std::vec::Vec<u8>>,
{
    let mut reconstructed_shards: Vec<Option<Vec<u8>>> = shards
        .iter()
        .map(|shard| {
            if shard.iter().all(|&byte| byte == 0) {
                None
            } else {
                Some(shard.to_vec())
            }
        })
        .collect();

    ec.rs
        .reconstruct(&mut reconstructed_shards)
        .map_err(|e| CoderError::InvalidShardData(e.to_string()))?;

    let result: Vec<T> = reconstructed_shards.into_iter().flatten().collect();

    debug!("reconstructed shards: {:?}", result.len());
    ec.rs
        .verify(&result)
        .map_err(|e| CoderError::FailedToVerify(e.to_string()))?;

    Ok(result)
}