Skip to main content

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}