miden_core/utils/col_matrix.rs
1use alloc::vec::Vec;
2
3#[derive(Debug)]
4pub struct ColMatrix<E> {
5 columns: Vec<Vec<E>>,
6}
7
8impl<E: Clone + Copy> ColMatrix<E> {
9 // CONSTRUCTOR
10 // --------------------------------------------------------------------------------------------
11 /// Returns a new [ColMatrix] instantiated with the data from the specified columns.
12 ///
13 /// # Panics
14 /// Panics if:
15 /// * The provided vector of columns is empty.
16 /// * Not all of the columns have the same number of elements.
17 /// * Number of rows is smaller than or equal to 1.
18 /// * Number of rows is not a power of two.
19 pub fn new(columns: Vec<Vec<E>>) -> Self {
20 assert!(!columns.is_empty(), "a matrix must contain at least one column");
21 let num_rows = columns[0].len();
22 assert!(num_rows > 1, "number of rows in a matrix must be greater than one");
23 assert!(num_rows.is_power_of_two(), "number of rows in a matrix must be a power of 2");
24 for column in columns.iter().skip(1) {
25 assert_eq!(column.len(), num_rows, "all matrix columns must have the same length");
26 }
27
28 Self { columns }
29 }
30 // PUBLIC ACCESSORS
31 // --------------------------------------------------------------------------------------------
32
33 /// Returns the number of columns in this matrix.
34 pub fn num_cols(&self) -> usize {
35 self.columns.len()
36 }
37
38 /// Returns the number of rows in this matrix.
39 pub fn num_rows(&self) -> usize {
40 self.columns[0].len()
41 }
42
43 /// Returns the element located at the specified column and row indexes in this matrix.
44 ///
45 /// # Panics
46 /// Panics if either `col_idx` or `row_idx` are out of bounds for this matrix.
47 pub fn get(&self, col_idx: usize, row_idx: usize) -> E {
48 self.columns[col_idx][row_idx]
49 }
50
51 /// Returns a reference to the column at the specified index.
52 pub fn get_column(&self, col_idx: usize) -> &[E] {
53 &self.columns[col_idx]
54 }
55
56 /// Returns a reference to the column at the specified index.
57 pub fn get_column_mut(&mut self, col_idx: usize) -> &mut [E] {
58 &mut self.columns[col_idx]
59 }
60
61 /// Returns an iterator over all columns in this matrix.
62 pub fn columns(&self) -> impl Iterator<Item = &[E]> {
63 self.columns.iter().map(|col| col.as_slice())
64 }
65
66 /// Copies values of all columns at the specified row into the specified row slice.
67 ///
68 /// # Panics
69 /// Panics if `row_idx` is out of bounds for this matrix.
70 pub fn read_row_into(&self, row_idx: usize, row: &mut [E]) {
71 for (column, value) in self.columns.iter().zip(row.iter_mut()) {
72 *value = column[row_idx];
73 }
74 }
75
76 /// Updates a row in this matrix at the specified index to the provided data.
77 ///
78 /// # Panics
79 /// Panics if `row_idx` is out of bounds for this matrix.
80 pub fn update_row(&mut self, row_idx: usize, row: &[E]) {
81 for (column, &value) in self.columns.iter_mut().zip(row) {
82 column[row_idx] = value;
83 }
84 }
85
86 /// Merges a column to the end of the matrix provided its length matches the matrix.
87 ///
88 /// # Panics
89 /// Panics if the column has a different length to other columns in the matrix.
90 pub fn merge_column(&mut self, column: Vec<E>) {
91 if let Some(first_column) = self.columns.first() {
92 assert_eq!(first_column.len(), column.len());
93 }
94 self.columns.push(column);
95 }
96
97 /// Removes a column of the matrix given its index.
98 ///
99 /// # Panics
100 /// Panics if the column index is out of range.
101 pub fn remove_column(&mut self, index: usize) -> Vec<E> {
102 assert!(index < self.num_cols(), "column index out of range");
103 self.columns.remove(index)
104 }
105}