Add column positions function.

kate/recovery/Cargo.toml

@@ -10,6 +10,7 @@ derive_more = "0.99.17"
 dusk-bytes = "0.1.6"
 dusk-plonk = { git = "https://github.com/maticnetwork/plonk.git", tag = "v0.12.0-polygon-2" }
 getrandom = { version = "0.2", features = ["js"] }
+num = "0.4.0"
 once_cell = { version = "1.9.0", default-features = false }
 rand = "0.8.4"
 rand_chacha = "0.3"

kate/recovery/src/com.rs

@@ -1,8 +1,13 @@
-use std::{collections::HashMap, convert::TryFrom};
-
 use codec::Decode;
 use dusk_bytes::Serializable;
 use dusk_plonk::{fft::EvaluationDomain, prelude::BlsScalar};
+use num::ToPrimitive;
+use rand::seq::SliceRandom;
+use std::{
+	collections::{BTreeSet, HashMap},
+	convert::TryFrom,
+	iter::FromIterator,
+};
 use thiserror::Error;
 
 use crate::{config, data, index, matrix};
@@ -25,6 +30,30 @@ pub enum ReconstructionError {
 	RowCountExceeded,
 }
 
+/// From given positions, constructs related columns positions, up to given factor.
+/// E.g. if factor is 0.66, 66% of matched columns will be returned.
+/// Positions in columns are random.
+/// Function panics if factor is above 1.0.
+pub fn columns_positions(
+	dimensions: &matrix::Dimensions,
+	positions: Vec<matrix::Position>,
+	factor: f64,
+) -> Vec<matrix::Position> {
+	assert!(factor <= 1.0);
+
+	let cells = (factor * dimensions.extended_rows() as f64)
+		.to_usize()
+		.expect("result is lesser than usize maximum");
+
+	let rng = &mut rand::thread_rng();
+
+	BTreeSet::from_iter(positions.iter().map(|position| position.col))
+		.into_iter()
+		.map(|col| dimensions.col_positions(col))
+		.flat_map(|col| col.choose_multiple(rng, cells).cloned().collect::<Vec<_>>())
+		.collect::<Vec<matrix::Position>>()
+}
+
 /// Creates hash map of columns, each being hash map of cells, from vector of cells.
 /// Intention is to be able to find duplicates and to group cells by column.
 fn map_cells(

kate/recovery/src/matrix.rs

@@ -98,6 +98,17 @@ impl Dimensions {
 			.collect::<Vec<u32>>()
 	}
 
+	/// Cell positions for given column in extended matrix.
+	/// Empty if column index is not valid.
+	pub fn col_positions(&self, col: u16) -> Vec<Position> {
+		if self.cols() <= col {
+			return vec![];
+		}
+		(0..self.extended_rows())
+			.map(|row| Position { col, row })
+			.collect::<Vec<_>>()
+	}
+
 	/// Cell positions for given rows in extended matrix.
 	/// Empty if row index is not valid.
 	fn extended_row_positions(&self, row: u32) -> Vec<Position> {