add recovery from systematic chunks

reed-solomon-novelpoly/Cargo.toml

@@ -35,6 +35,7 @@ itertools = "0.12"
 reed-solomon-tester = { path = "../reed-solomon-tester" }
 rand = { version = "0.8.3", features = ["alloc", "small_rng"] }
 assert_matches = "1.5.0"
+quickcheck = { version = "1.0.3", default-features = false }
 
 [features]
 default = []

reed-solomon-novelpoly/src/novel_poly_basis/mod.rs

@@ -156,7 +156,9 @@ impl ReedSolomon {
 		Ok(shards)
 	}
 
-	/// each shard contains one symbol of one run of erasure coding
+	/// Reconstruct from chunks.
+	///
+	/// The result may be padded with zeros. Truncate the output to the expected byte length.
 	pub fn reconstruct<S: Shard>(&self, received_shards: Vec<Option<S>>) -> Result<Vec<u8>> {
 		let gap = self.n.saturating_sub(received_shards.len());
 
@@ -231,6 +233,52 @@ impl ReedSolomon {
 
 		Ok(acc)
 	}
+
+	/// Reconstruct from the set of systematic chunks.
+	/// Systematic chunks are the first `k` chunks, which contain the initial data.
+	///
+	/// Provide a vector containing chunk data. If too few chunks are provided, recovery is not
+	/// possible.
+	/// The result may be padded with zeros. Truncate the output to the expected byte length.
+	pub fn reconstruct_from_systematic<S: Shard>(&self, chunks: Vec<S>) -> Result<Vec<u8>> {
+		let Some(first_shard) = chunks.iter().next() else {
+			return Err(Error::NeedMoreShards { have: 0, min: self.k, all: self.n });
+		};
+		let shard_len = AsRef::<[[u8; 2]]>::as_ref(first_shard).len();
+
+		if shard_len == 0 {
+			return Err(Error::InconsistentShardLengths { first: 0, other: 0 });
+		}
+
+		if let Some(length) = chunks.iter().find_map(|c| {
+			let length = AsRef::<[[u8; 2]]>::as_ref(c).len();
+			if length != shard_len {
+				Some(length)
+			} else {
+				None
+			}
+		}) {
+			return Err(Error::InconsistentShardLengths { first: shard_len, other: length });
+		}
+
+		if chunks.len() < self.k {
+			return Err(Error::NeedMoreShards { have: chunks.len(), min: self.k, all: self.n });
+		}
+
+		let mut systematic_bytes = Vec::with_capacity(shard_len * 2 * self.k);
+
+		for i in 0..shard_len {
+			for chunk in chunks.iter().take(self.k) {
+				// No need to check for index out of bounds because i goes up to shard_len and
+				// we return an error for non uniform chunks.
+				let chunk = AsRef::<[[u8; 2]]>::as_ref(chunk)[i];
+				systematic_bytes.push(chunk[0]);
+				systematic_bytes.push(chunk[1]);
+			}
+		}
+
+		Ok(systematic_bytes)
+	}
 }
 
 #[cfg(test)]

reed-solomon-novelpoly/src/novel_poly_basis/tests.rs

@@ -9,6 +9,7 @@ use rand::rngs::SmallRng;
 use rand::seq::index::IndexVec;
 use rand::thread_rng;
 use reed_solomon_tester::*;
+use quickcheck::{Arbitrary, Gen, QuickCheck};
 
 /// Generate a random index
 fn rand_gf_element() -> Additive {
@@ -463,3 +464,34 @@ fn shard_len_is_reasonable() {
 	// needs 3 bytes to fit, rounded up to next even number.
 	assert_eq!(rs.shard_len(19), 6);
 }
+
+#[derive(Clone, Debug)]
+struct ArbitraryData(Vec<u8>);
+
+impl Arbitrary for ArbitraryData {
+	fn arbitrary(g: &mut Gen) -> Self {
+		// Limit the len to 1 mib, otherwise the test will take forever
+		let len = u32::arbitrary(g).saturating_add(2) % (1024 * 1024);
+
+		let data: Vec<u8> = (0..len).map(|_| u8::arbitrary(g)).collect();
+
+		ArbitraryData(data)
+	}
+}
+
+#[test]
+fn round_trip_systematic_quickcheck() {
+	fn property(available_data: ArbitraryData, n_validators: u16) {
+		let n_validators = n_validators.saturating_add(2);
+		let rs = CodeParams::derive_parameters(n_validators as usize, (n_validators as usize - 1) / 3 + 1)
+			.unwrap()
+			.make_encoder();
+		let kpow2 = rs.k;
+		let chunks = rs.encode::<WrappedShard>(&available_data.0).unwrap();
+		let mut res = rs.reconstruct_from_systematic(chunks.into_iter().take(kpow2).collect()).unwrap();
+		res.truncate(available_data.0.len());
+		assert_eq!(res, available_data.0);
+	}
+
+	QuickCheck::new().quickcheck(property as fn(ArbitraryData, u16))
+}