diff --git a/kate/recovery/src/data.rs b/kate/recovery/src/data.rs index 3356bfd1..0109fe5e 100644 --- a/kate/recovery/src/data.rs +++ b/kate/recovery/src/data.rs @@ -1,6 +1,6 @@ use std::{collections::HashMap, convert::TryInto}; -use crate::matrix::Position; +use crate::matrix::{Dimensions, Position, RowIndex}; /// Position and data of a cell in extended matrix #[derive(Default, Debug, Clone)] @@ -36,7 +36,7 @@ impl Cell { /// Merges cells data per row. /// Cells are sorted before merge. -pub fn rows(cells: &[Cell]) -> Vec<(u32, Vec)> { +pub fn rows(dimensions: &Dimensions, cells: &[&Cell]) -> Vec<(RowIndex, Vec)> { let mut sorted_cells = cells.to_vec(); sorted_cells @@ -44,11 +44,12 @@ pub fn rows(cells: &[Cell]) -> Vec<(u32, Vec)> { let mut rows = HashMap::new(); for cell in sorted_cells { - rows.entry(cell.position.row) + rows.entry(RowIndex(cell.position.row)) .or_insert_with(Vec::default) .extend(cell.data()); } + rows.retain(|_, row| row.len() != dimensions.row_byte_size()); rows.into_iter().collect::>() } @@ -65,7 +66,11 @@ impl From for DataCell { mod tests { use std::convert::TryInto; - use crate::{data::rows, data::Cell, matrix::Position}; + use crate::{ + data::rows, + data::Cell, + matrix::{Dimensions, Position}, + }; fn cell(position: Position, content: [u8; 80]) -> Cell { Cell { position, content } @@ -77,19 +82,20 @@ mod tests { #[test] fn rows_ok() { + let dimensions = Dimensions::new(2, 1).unwrap(); fn content(data: [u8; 32]) -> [u8; 80] { [&[0u8; 48], &data[..]].concat().try_into().unwrap() } let cells = [ - cell(position(1, 1), content([3; 32])), - cell(position(1, 0), content([2; 32])), - cell(position(0, 0), content([0; 32])), - cell(position(0, 1), content([1; 32])), + &cell(position(1, 1), content([3; 32])), + &cell(position(1, 0), content([2; 32])), + &cell(position(0, 0), content([0; 32])), + &cell(position(0, 1), content([1; 32])), ]; - let mut rows = rows(&cells); - rows.sort_by_key(|&(key, _)| key); + let mut rows = rows(&dimensions, &cells); + rows.sort_by_key(|(key, _)| key.0); let expected = [ [[0u8; 32], [1u8; 32]].concat(), @@ -98,8 +104,10 @@ mod tests { for i in 0..1 { let (row_index, row) = &rows[i]; - assert_eq!(*row_index, 0); + assert_eq!(row_index.0, 0); assert_eq!(*row, expected[i]); } + + // TODO: Update row test with incomplete row case } } diff --git a/kate/recovery/src/matrix.rs b/kate/recovery/src/matrix.rs index 670b4617..5985ebd2 100644 --- a/kate/recovery/src/matrix.rs +++ b/kate/recovery/src/matrix.rs @@ -2,7 +2,7 @@ use std::ops::Range; use serde::{Deserialize, Serialize}; -use crate::config; +use crate::config::{self, CHUNK_SIZE}; const EXTENSION_FACTOR_U32: u32 = config::EXTENSION_FACTOR as u32; @@ -18,6 +18,11 @@ impl Position { pub fn reference(&self, block_number: u32) -> String { format!("{}:{}:{}", block_number, self.col, self.row) } + + /// Checks if position is from extended row + pub fn is_extended(&self) -> bool { + self.row % 2 == 1 + } } /// Matrix partition (column-wise) @@ -27,6 +32,17 @@ pub struct Partition { pub fraction: u8, } +/// Matrix row index +#[derive(Debug, Clone, Hash, PartialEq, Eq, PartialOrd, Ord)] +pub struct RowIndex(pub u32); + +impl RowIndex { + /// Refrence in format `block_number:row_number` + pub fn reference(&self, block_number: u32) -> String { + format!("{}:{}", block_number, self.0) + } +} + /// Dimensions of a non-extended matrix. /// Extended matrix (with factor of 2) is a matrix where data is in odd rows and even rows contains erasure codes. /// Matrix is represented as list of cells, concatenated column by column, which is optimized for erasure coding. @@ -78,6 +94,11 @@ impl Dimensions { self.extended_rows() as u64 * self.cols as u64 } + /// Row size in bytes + pub fn row_byte_size(&self) -> usize { + CHUNK_SIZE * self.cols as usize + } + /// Extended matrix rows count. pub fn extended_rows(&self) -> u32 { (self.rows as u32) * EXTENSION_FACTOR_U32