From e307b850d9d8f431dd91390c97ecddeb0cb2b23e Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Aleksandar=20Terenti=C4=87?= <aleksandar@ethernal.tech>
Date: Wed, 7 Dec 2022 13:16:09 +0100
Subject: [PATCH] Return hash map of columns in reconstruct_columns.

---
 kate/recovery/src/com.rs | 26 +++++++++++++++++++-------
 1 file changed, 19 insertions(+), 7 deletions(-)

diff --git a/kate/recovery/src/com.rs b/kate/recovery/src/com.rs
index 32858dee..a381954b 100644
--- a/kate/recovery/src/com.rs
+++ b/kate/recovery/src/com.rs
@@ -4,13 +4,16 @@ use dusk_plonk::{fft::EvaluationDomain, prelude::BlsScalar};
 use num::ToPrimitive;
 use rand::seq::SliceRandom;
 use std::{
-	collections::{BTreeSet, HashMap},
+	collections::{HashMap, HashSet},
 	convert::TryFrom,
 	iter::FromIterator,
 };
 use thiserror::Error;
 
-use crate::{config, data, index, matrix};
+use crate::{
+	config::{self, CHUNK_SIZE},
+	data, index, matrix,
+};
 
 #[derive(Debug, Error)]
 pub enum ReconstructionError {
@@ -47,7 +50,9 @@ pub fn columns_positions(
 
 	let rng = &mut rand::thread_rng();
 
-	BTreeSet::from_iter(positions.iter().map(|position| position.col))
+	let columns: HashSet<u16> = HashSet::from_iter(positions.iter().map(|position| position.col));
+
+	columns
 		.into_iter()
 		.map(|col| dimensions.col_positions(col))
 		.flat_map(|col| col.choose_multiple(rng, cells).cloned().collect::<Vec<_>>())
@@ -158,10 +163,17 @@ pub fn reconstruct_extrinsics(
 	unflatten_padded_data(ranges, data).map_err(ReconstructionError::DataDecodingError)
 }
 
+/// Reconstructs columns for given cells.
+///
+/// # Arguments
+///
+/// * `dimensions` - Extended matrix dimensions
+/// * `cells` - Cells from required columns, at least 50% cells per column
 pub fn reconstruct_columns(
 	dimensions: &matrix::Dimensions,
-	cells: Vec<data::DataCell>,
-) -> Result<Vec<(u16, Vec<u8>)>, ReconstructionError> {
+	cells: &[data::Cell],
+) -> Result<HashMap<u16, Vec<[u8; CHUNK_SIZE]>>, ReconstructionError> {
+	let cells: Vec<data::DataCell> = cells.iter().cloned().map(Into::into).collect::<Vec<_>>();
 	let columns = map_cells(dimensions, cells)?;
 
 	columns
@@ -176,8 +188,8 @@ pub fn reconstruct_columns(
 			let column = reconstruct_column(dimensions.extended_rows(), &cells)
 				.map_err(ReconstructionError::ColumnReconstructionError)?
 				.iter()
-				.flat_map(BlsScalar::to_bytes)
-				.collect::<Vec<_>>();
+				.map(BlsScalar::to_bytes)
+				.collect::<Vec<[u8; CHUNK_SIZE]>>();
 
 			Ok((col, column))
 		})