From a3cea9a09c23b9bbc41a3784fd0adbdf847ba6b4 Mon Sep 17 00:00:00 2001
From: Daniel Sanchez <sanchez.quiros.daniel@gmail.com>
Date: Tue, 8 Oct 2024 09:15:41 +0200
Subject: [PATCH] DA: Implement data reconstruction from original columns
 (#808)

* Implement data reconstruction from original columns

* Make reconstruction public
---
 nomos-da/kzgrs-backend/src/lib.rs            |  1 +
 nomos-da/kzgrs-backend/src/reconstruction.rs | 62 ++++++++++++++++++++
 2 files changed, 63 insertions(+)
 create mode 100644 nomos-da/kzgrs-backend/src/reconstruction.rs

diff --git a/nomos-da/kzgrs-backend/src/lib.rs b/nomos-da/kzgrs-backend/src/lib.rs
index e8192b89..aea77b0f 100644
--- a/nomos-da/kzgrs-backend/src/lib.rs
+++ b/nomos-da/kzgrs-backend/src/lib.rs
@@ -2,6 +2,7 @@ pub mod common;
 pub mod dispersal;
 pub mod encoder;
 pub mod global;
+pub mod reconstruction;
 pub mod verifier;
 
 pub use kzgrs::KzgRsError;
diff --git a/nomos-da/kzgrs-backend/src/reconstruction.rs b/nomos-da/kzgrs-backend/src/reconstruction.rs
new file mode 100644
index 00000000..315d76ba
--- /dev/null
+++ b/nomos-da/kzgrs-backend/src/reconstruction.rs
@@ -0,0 +1,62 @@
+use crate::common::blob::DaBlob;
+use crate::common::Chunk;
+use kzgrs::BYTES_PER_FIELD_ELEMENT;
+
+/// Reconstruct original data from a set of `DaBlob`
+/// Warning! This does not interpolate so it should not be used on blobs which doesn't represent
+/// the original set of data.
+pub fn reconstruct_without_missing_data(blobs: &[DaBlob]) -> Vec<u8> {
+    // pick positions from columns
+    let mut data: Vec<((usize, usize), Vec<u8>)> = blobs
+        .iter()
+        .flat_map(|blob| {
+            blob.column
+                .iter()
+                .map(Chunk::as_bytes)
+                .enumerate()
+                .map(|(row, data)| ((row, blob.column_idx as usize), data.to_vec()))
+        })
+        .collect();
+    data.sort_unstable_by_key(|(k, _)| *k);
+    data.into_iter()
+        .flat_map(|(_, mut data)| {
+            assert_eq!(data.len(), BYTES_PER_FIELD_ELEMENT);
+            // pop last byte as we fill it empty
+            data.pop();
+            data.into_iter()
+        })
+        .collect()
+}
+
+#[cfg(test)]
+mod test {
+    use crate::common::blob::DaBlob;
+    use crate::common::ColumnIndex;
+    use crate::encoder::test::rand_data;
+    use crate::encoder::{DaEncoder, DaEncoderParams, EncodedData};
+    use crate::reconstruction::reconstruct_without_missing_data;
+    use nomos_core::da::DaEncoder as _;
+
+    #[test]
+    fn test_reconstruct() {
+        let data: Vec<u8> = rand_data(32);
+        let params = DaEncoderParams::default_with(4);
+        let encoder = DaEncoder::new(params);
+        let encoded_data: EncodedData = encoder.encode(&data).unwrap();
+        let blobs: Vec<DaBlob> = encoded_data
+            .chunked_data
+            .columns()
+            .enumerate()
+            .map(|(idx, column)| DaBlob {
+                column,
+                column_idx: idx as ColumnIndex,
+                column_commitment: Default::default(),
+                aggregated_column_commitment: Default::default(),
+                aggregated_column_proof: Default::default(),
+                rows_commitments: vec![],
+                rows_proofs: vec![],
+            })
+            .collect();
+        assert_eq!(data, reconstruct_without_missing_data(&blobs));
+    }
+}