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; if sq_root_rounded % 2 != 0 {
21 sq_root_rounded += 1; }
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 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(); 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(); 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 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 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}