rs2d/
lib.rs

1use std::vec;
2
3use reed_solomon_erasure::*;
4use solana_merkle_tree::MerkleTree;
5pub type Shard = Vec<u8>;
6pub mod macros;
7
8#[derive(Debug)]
9pub struct EncodedData {
10    pub row_commitments: Vec<Vec<u8>>,
11    pub col_commitments: Vec<Vec<u8>>,
12    pub final_commitment: Vec<u8>,
13}
14pub fn encode(data: Vec<Shard>, shard_len: usize) -> EncodedData {
15    let data_array_length = data.len();
16    let sq_root = (data_array_length as f64).sqrt();
17
18    let mut sq_root_rounded = sq_root.ceil() as usize; // 45
19
20    if sq_root_rounded % 2 != 0 {
21        sq_root_rounded += 1; // doing this so we can split data and coding chunks equally
22    }
23
24    let mut data_matrix = vec![vec![vec![0; shard_len]; sq_root_rounded]; sq_root_rounded * 2];
25    for i in 0..data_array_length {
26        // 2000
27        let row = i / sq_root_rounded;
28        let col = i % sq_root_rounded;
29        data_matrix[row][col] = data[i].clone();
30    }
31
32    for i in 0..(sq_root_rounded * 2) {
33        let mut parity_prefill = vec![vec![0u8; shard_len]; sq_root_rounded];
34        data_matrix[i].append(&mut parity_prefill);
35    }
36
37    let mut vertical_erasure_swap: Vec<Vec<Shard>> = vec![];
38    let data_matrix = data_matrix
39        .iter_mut()
40        .enumerate()
41        .map(|(i, row)| {
42            let r = ReedSolomon::new(sq_root_rounded, sq_root_rounded).unwrap(); // assuming n:k is 1:1
43                                                                                 // let mut row = row.clone(); // m.erasure_matrix.push(r.clone());
44
45            let mut x: Vec<&mut [u8]> = row.iter_mut().map(|f| f.as_mut_slice()).collect();
46
47            r.encode(&mut x).unwrap();
48
49            if i < sq_root_rounded {
50                let y = x.iter().map(|f| f.to_vec()).collect::<Vec<Shard>>();
51
52                vertical_erasure_swap.push(y[sq_root_rounded..y.len()].to_vec().clone());
53            }
54            return x;
55        })
56        .collect::<Vec<Vec<&mut [u8]>>>();
57
58    let data_matrix = owned_vec!(data_matrix);
59    let mut data_matrix = data_matrix
60        .into_iter()
61        .enumerate()
62        .map(|(i, mut row)| {
63            if i < sq_root_rounded {
64                return row;
65            } else {
66                for item_index in 0..row.len() {
67                    if item_index < sq_root_rounded {
68                        row[item_index] =
69                            vertical_erasure_swap[i - sq_root_rounded][item_index].clone();
70                    }
71                }
72                return row;
73            }
74        })
75        .collect::<Vec<Vec<Shard>>>();
76
77    let data_matrix = data_matrix
78        .iter_mut()
79        .enumerate()
80        .map(|(i, row)| {
81            if i >= sq_root_rounded {
82                let r = ReedSolomon::new(sq_root_rounded, sq_root_rounded).unwrap(); // assuming n:k is 1:1
83                                                                                     // let mut row = row.clone(); // m.erasure_matrix.push(r.clone());
84
85                let mut x: Vec<&mut [u8]> = row.iter_mut().map(|f| f.as_mut_slice()).collect();
86
87                r.encode(&mut x).unwrap();
88
89                return x;
90            } else {
91                return row
92                    .iter_mut()
93                    .map(|f| f.as_mut_slice())
94                    .collect::<Vec<&mut [u8]>>();
95            }
96        })
97        .collect::<Vec<Vec<&mut [u8]>>>();
98
99    let row_commitments = data_matrix
100        .iter()
101        .map(|row| {
102            MerkleTree::new(&row)
103                .get_root()
104                .unwrap()
105                .as_ref()
106                .to_owned()
107        })
108        .collect::<Vec<Vec<u8>>>();
109
110    let data_matrix = owned_vec!(data_matrix);
111    let mut col_leaves = vec![];
112
113    for col_index in 0..sq_root_rounded * 2 {
114        for row_index in 0..sq_root_rounded * 2 {
115            col_leaves.push(data_matrix[row_index][col_index].clone());
116        }
117    }
118    let col_commitments = data_matrix
119        .iter()
120        .map(|col| {
121            MerkleTree::new(&col)
122                .get_root()
123                .unwrap()
124                .as_ref()
125                .to_owned()
126        })
127        .collect::<Vec<Vec<u8>>>();
128
129    let final_commitment =
130        MerkleTree::new(&[row_commitments.clone(), col_commitments.clone()].concat())
131            .get_root()
132            .unwrap()
133            .as_ref()
134            .to_vec();
135    EncodedData {
136        row_commitments,
137        col_commitments,
138        final_commitment,
139    }
140}
141
142#[cfg(test)]
143mod tests {
144    use super::*;
145    use rand::Rng;
146    #[test]
147    fn test_encode() {
148        // make a 2d vec of 23 numbers from 1 to 23
149        let mut original = vec![];
150        let data_len = 1560;
151        let shard_len = 1280;
152        for _ in 0..data_len {
153            let mut elem: Shard = Vec::with_capacity(shard_len as usize);
154            for _ in 0..shard_len {
155                elem.push(0);
156            }
157            // let mut elem = elem.into_boxed_slice();
158            for i in 0..shard_len {
159                elem[i as usize] = rand::thread_rng().gen::<u8>()
160            }
161            original.push(elem.clone());
162        }
163        encode(original, shard_len);
164    }
165}