apollo_propeller: add reed-solomon erasure coding (#11047)

Author: sirandreww-starkware

Cargo.lock

@@ -315,6 +315,7 @@ quote = "1.0.26"
 rand = "0.8.5"
 rand_chacha = "0.3.1"
 rand_distr = "0.4.3"
+reed-solomon-simd = "3.1.0"
 regex = "1.10.4"
 replace_with = "0.1.7"
 reqwest = "0.12"

Cargo.toml

@@ -10,6 +10,7 @@ description = "Implementation of the Propeller algorithm for the Starknet sequen
 testing = []
 
 [dependencies]
+reed-solomon-simd.workspace = true
 sha2.workspace = true
 
 [dev-dependencies]

crates/apollo_propeller/Cargo.toml

@@ -1,3 +1,7 @@
 pub mod merkle;
 #[cfg(test)]
 mod merkle_test;
+// TODO(AndrewL): Consider renaming this to `erasure_coding` or `error_correction_code`.
+pub mod reed_solomon;
+#[cfg(test)]
+mod reed_solomon_test;

crates/apollo_propeller/src/lib.rs

@@ -0,0 +1,100 @@
+use reed_solomon_simd::{ReedSolomonDecoder, ReedSolomonEncoder};
+
+// TODO(AndrewL): Consider combining this with `generate_coding_shards`.
+// TODO(AndrewL): Consider adding custom error type and using it here.
+pub fn split_data_into_shards(message: Vec<u8>, num_data_shards: usize) -> Option<Vec<Vec<u8>>> {
+    if !message.len().is_multiple_of(num_data_shards) {
+        return None;
+    }
+    let shard_size = message.len() / num_data_shards;
+    Some(message.chunks_exact(shard_size).map(|chunk| chunk.to_vec()).collect())
+}
+
+/// Generate coding shards using Reed-Solomon encoding.
+// TODO(AndrewL): Consider adding custom error type and using it here.
+pub fn generate_coding_shards(
+    data_shards: &[Vec<u8>],
+    num_coding_shards: usize,
+) -> Result<Vec<Vec<u8>>, String> {
+    if num_coding_shards == 0 {
+        // ReedSolomonEncoder does not support 0 coding shards
+        return Ok(Vec::new());
+    }
+
+    let num_data_shards = data_shards.len();
+    // TODO(AndrewL): Consider accepting a shard size as an argument.
+    let shard_size = data_shards.first().ok_or("No data shards".to_string())?.len();
+
+    let mut encoder = ReedSolomonEncoder::new(num_data_shards, num_coding_shards, shard_size)
+        .map_err(|e| format!("Failed to create Reed-Solomon encoder: {}", e))?;
+
+    for shard in data_shards.iter().take(num_data_shards) {
+        encoder
+            .add_original_shard(shard)
+            .map_err(|e| format!("Failed to add data shard: {}", e))?;
+    }
+
+    let result = encoder.encode().map_err(|e| format!("Failed to encode: {}", e))?;
+
+    let coding_shards = result.recovery_iter().map(|shard| shard.to_vec()).collect();
+
+    Ok(coding_shards)
+}
+
+/// Reconstruct the original message from available shards using Reed-Solomon error correction.
+// TODO(AndrewL): Consider adding custom error type and using it here.
+// TODO(AndrewL): Rename this to `reconstruct_data_shards`.
+pub fn reconstruct_message_from_shards(
+    // TODO(AndrewL): Change this to a HashMap<usize, Vec<u8>>.
+    shards: &[(usize, Vec<u8>)],
+    num_data_shards: usize,
+    num_coding_shards: usize,
+) -> Result<Vec<Vec<u8>>, String> {
+    if num_coding_shards == 0 {
+        return Ok(shards.iter().map(|(_, s)| s.to_vec()).collect());
+    }
+    // TODO(AndrewL): Consider accepting a shard size as an argument.
+    let shard_size = shards.first().ok_or("No shards".to_string())?.1.len();
+
+    let mut decoder = ReedSolomonDecoder::new(num_data_shards, num_coding_shards, shard_size)
+        .map_err(|e| format!("Failed to create Reed-Solomon decoder: {}", e))?;
+
+    for (index, shard_data) in shards {
+        if *index < num_data_shards {
+            decoder
+                .add_original_shard(*index, shard_data)
+                .map_err(|e| format!("Failed to add original shard: {}", e))?;
+        } else {
+            decoder
+                .add_recovery_shard(index - num_data_shards, shard_data)
+                .map_err(|e| format!("Failed to add coding shard: {}", e))?;
+        }
+    }
+
+    let result = decoder.decode().map_err(|e| format!("Failed to decode: {}", e))?;
+
+    let mut shard_map = std::collections::HashMap::new();
+    for (index, shard) in shards {
+        shard_map.insert(index, shard);
+    }
+
+    let mut data_shards = Vec::with_capacity(num_data_shards);
+    for index in 0..num_data_shards {
+        if let Some(shard_shard) = shard_map.get(&index) {
+            data_shards.push(shard_shard.to_vec());
+        } else if let Some(restored_data) = result.restored_original(index) {
+            data_shards.push(restored_data.to_vec());
+        } else {
+            return Err(format!(
+                "Missing data shard at index {} and no restored data available",
+                index
+            ));
+        }
+    }
+
+    Ok(data_shards)
+}
+
+pub fn combine_data_shards(data_shards: Vec<Vec<u8>>) -> Vec<u8> {
+    data_shards.iter().flatten().copied().collect()
+}

crates/apollo_propeller/src/reed_solomon.rs

@@ -0,0 +1,57 @@
+use rstest::rstest;
+
+use crate::reed_solomon::{generate_coding_shards, reconstruct_message_from_shards};
+
+#[test]
+fn test_empty_generate_coding_shards() {
+    let data_shards = vec![vec![0, 1], vec![2, 3], vec![4, 5]];
+    let num_coding_shards = 0;
+    let coding_shards = generate_coding_shards(&data_shards, num_coding_shards).unwrap();
+    assert!(coding_shards.is_empty());
+}
+
+#[rstest]
+#[case(3, 2, 1, 8)]
+#[case(4, 2, 2, 16)]
+#[case(5, 3, 3, 8)]
+#[case(4, 4, 4, 32)]
+fn test_reed_solomon_with_lost_shards(
+    #[case] num_data_shards: u8,
+    #[case] num_coding_shards: usize,
+    #[case] num_lost_shards: usize,
+    #[case] shard_size: usize,
+) {
+    let data_shards: Vec<Vec<u8>> = (0..num_data_shards).map(|i| vec![i; shard_size]).collect();
+    let num_data_shards: usize = num_data_shards.into();
+    let original_data = data_shards.clone();
+
+    let coding_shards = generate_coding_shards(&data_shards, num_coding_shards).unwrap();
+    assert_eq!(coding_shards.len(), num_coding_shards);
+
+    let all_shards: Vec<(usize, Vec<u8>)> = data_shards
+        .iter()
+        .enumerate()
+        .map(|(i, s)| (i, s.clone()))
+        .chain(coding_shards.iter().enumerate().map(|(i, s)| (num_data_shards + i, s.clone())))
+        .collect();
+
+    let available_shards: Vec<(usize, Vec<u8>)> = all_shards
+        .into_iter()
+        .enumerate()
+        .filter(|(i, _)| *i % 2 != 0 || *i >= num_lost_shards * 2)
+        .map(|(_, shard)| shard)
+        .collect();
+
+    assert!(
+        available_shards.len() >= num_data_shards,
+        "Not enough shards to reconstruct (have {}, need {})",
+        available_shards.len(),
+        num_data_shards
+    );
+
+    let reconstructed_data =
+        reconstruct_message_from_shards(&available_shards, num_data_shards, num_coding_shards)
+            .unwrap();
+
+    assert_eq!(reconstructed_data, original_data, "Reconstructed data doesn't match original");
+}