1use std::{io, ops::Deref};
2
3use byteorder::{LittleEndian, ReadBytesExt, WriteBytesExt};
4use itertools::Itertools;
5use serde::{Deserialize, Serialize};
6
7use super::Matrix;
8use crate::{
9 prime::ValidPrime,
10 vector::{FpSlice, FpSliceMut, FpVector},
11};
12
13#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
23#[repr(transparent)]
24pub struct Subspace {
25 matrix: Matrix,
26}
27
28impl Deref for Subspace {
31 type Target = Matrix;
32
33 fn deref(&self) -> &Self::Target {
34 &self.matrix
35 }
36}
37
38impl Subspace {
39 pub fn new(p: ValidPrime, dim: usize) -> Self {
40 let mut matrix = Matrix::new(p, dim + 1, dim);
44 matrix.initialize_pivots();
45 Self::from_matrix(matrix)
46 }
47
48 pub fn from_matrix(mut matrix: Matrix) -> Self {
50 matrix.row_reduce();
51 Self { matrix }
52 }
53
54 pub fn update_then_row_reduce<T, F: FnOnce(&mut Matrix) -> T>(&mut self, f: F) -> T {
56 let ret = f(&mut self.matrix);
57 self.matrix.row_reduce();
58 ret
59 }
60
61 pub fn from_bytes(p: ValidPrime, data: &mut impl io::Read) -> io::Result<Self> {
62 let rows = data.read_u64::<LittleEndian>()? as usize;
63 let ambient_dimension = data.read_u64::<LittleEndian>()? as usize;
64
65 let mut matrix = Matrix::from_bytes(p, rows, ambient_dimension, data)?;
66
67 matrix.pivots = Matrix::read_pivot(matrix.columns(), data)?;
68
69 Ok(Self { matrix })
70 }
71
72 pub fn to_bytes(&self, buffer: &mut impl io::Write) -> io::Result<()> {
73 buffer.write_u64::<LittleEndian>(self.matrix.rows() as u64)?;
74 buffer.write_u64::<LittleEndian>(self.ambient_dimension() as u64)?;
75
76 self.matrix.to_bytes(buffer)?;
77 Matrix::write_pivot(self.pivots(), buffer)
78 }
79
80 pub fn entire_space(p: ValidPrime, dim: usize) -> Self {
81 let mut result = Self::new(p, dim);
82 for i in 0..dim {
83 result.matrix.row_mut(i).set_entry(i, 1);
84 result.matrix.pivots_mut()[i] = i as isize;
85 }
86 result
87 }
88
89 pub fn add_vector(&mut self, row: FpSlice) -> usize {
97 let last_row = self.matrix.rows() - 1;
98 self.matrix.row_mut(last_row).assign(row);
99 self.matrix.row_reduce()
100 }
101
102 pub fn add_vectors(&mut self, mut rows: impl for<'a> FnMut(FpSliceMut<'a>) -> Option<()>) {
108 let num_rows = self.matrix.rows();
109 'outer: loop {
110 let first_row = self.dimension();
111 if first_row == num_rows {
112 return;
113 }
114
115 for i in first_row..num_rows {
116 if rows(self.matrix.row_mut(i)).is_none() {
117 break 'outer;
118 }
119 }
120 self.matrix.row_reduce();
121 }
122 self.matrix.row_reduce();
123 }
124
125 pub fn add_basis_elements(&mut self, mut rows: impl std::iter::Iterator<Item = usize>) {
126 self.add_vectors(|mut row| {
127 row.set_entry(rows.next()?, 1);
128 Some(())
129 });
130 }
131
132 pub fn reduce(&self, mut vector: FpSliceMut) {
135 assert_eq!(vector.as_slice().len(), self.ambient_dimension());
136 if self.matrix.rows() == 0 {
137 return;
138 }
139 let p = self.prime();
140 let iter = self
141 .pivots()
142 .iter()
143 .enumerate()
144 .filter(|(_, x)| **x >= 0)
145 .map(|(col, _)| col)
146 .zip(self.iter());
147 for (col, row) in iter {
148 let c = vector.as_slice().entry(col);
149 if c != 0 {
150 vector.add(row, p - c);
151 }
152 }
153 }
154
155 pub fn contains(&self, vector: FpSlice) -> bool {
156 let mut vector: FpVector = vector.to_owned();
157 self.reduce(vector.as_slice_mut());
158 vector.is_zero()
159 }
160
161 pub fn contains_space(&self, other: &Self) -> bool {
162 other.iter().all(|row| self.contains(row))
163 }
164
165 pub fn dimension(&self) -> usize {
166 self.pivots()
167 .iter()
168 .rev()
169 .find(|&&i| i >= 0)
170 .map_or(0, |&i| i as usize + 1)
171 }
172
173 pub fn is_empty(&self) -> bool {
175 self.matrix.rows() == 0 || self.matrix.row(0).is_zero()
176 }
177
178 pub fn ambient_dimension(&self) -> usize {
179 self.matrix.columns()
180 }
181
182 pub fn basis(&self) -> impl Iterator<Item = FpSlice<'_>> {
184 self.matrix.iter().take(self.dimension())
185 }
186
187 pub fn set_to_zero(&mut self) {
189 self.matrix.set_to_zero();
190 for x in self.matrix.pivots_mut() {
191 *x = -1;
192 }
193 }
194
195 pub fn set_to_entire(&mut self) {
197 self.matrix.set_to_zero();
198 for i in 0..self.matrix.columns() {
199 self.matrix.row_mut(i).set_entry(i, 1);
200 self.matrix.pivots_mut()[i] = i as isize;
201 }
202 }
203
204 pub fn iter(&self) -> impl Iterator<Item = FpSlice<'_>> {
205 self.matrix.iter().take(self.dimension())
206 }
207
208 pub fn iter_all_vectors(&self) -> impl Iterator<Item = FpVector> + '_ {
231 crate::prime::iter::combinations(self.prime(), self.dimension()).map(|combination| {
232 let mut vector = FpVector::new(self.prime(), self.ambient_dimension());
233 for (&c, v) in combination.iter().zip(self.matrix.iter()) {
234 vector.as_slice_mut().add(v, c);
235 }
236 vector
237 })
238 }
239
240 pub fn sum(&self, other: &Self) -> Self {
241 assert_eq!(self.prime(), other.prime());
242 assert_eq!(self.ambient_dimension(), other.ambient_dimension());
243
244 let mut sum = self.matrix.clone();
245 for other_row in other.iter() {
246 let mut new_row = sum.add_row();
247 new_row.assign(other_row);
248 }
249
250 let mut ret = Self::from_matrix(sum);
251 ret.matrix.trim(0, self.matrix.columns() + 1, 0);
252 ret
253 }
254}
255
256impl std::fmt::Display for Subspace {
257 fn fmt(&self, f: &mut std::fmt::Formatter) -> std::fmt::Result {
277 let dim = self.dimension();
278 let mut rows = self.matrix.iter().take(dim);
279
280 let output = if f.alternate() {
281 rows.join(", ")
282 } else {
283 rows.fold(String::new(), |mut output, row| {
284 use std::fmt::Write;
285 let _ = writeln!(output, "{row}");
286 output
287 })
288 };
289
290 write!(f, "{output}")?;
291 Ok(())
292 }
293}
294
295#[cfg(feature = "proptest")]
296pub mod arbitrary {
297 use proptest::prelude::*;
298
299 use super::*;
300 pub use crate::matrix::matrix_inner::arbitrary::MAX_COLUMNS as MAX_DIM;
301 use crate::matrix::matrix_inner::arbitrary::MatrixArbParams;
302
303 #[derive(Debug, Clone)]
304 pub struct SubspaceArbParams {
305 pub p: Option<ValidPrime>,
306 pub dim: BoxedStrategy<usize>,
307 }
308
309 impl Default for SubspaceArbParams {
310 fn default() -> Self {
311 Self {
312 p: None,
313 dim: (0..=MAX_DIM).boxed(),
314 }
315 }
316 }
317
318 impl Arbitrary for Subspace {
319 type Parameters = SubspaceArbParams;
320 type Strategy = BoxedStrategy<Self>;
321
322 fn arbitrary_with(args: Self::Parameters) -> Self::Strategy {
323 let p = match args.p {
324 Some(p) => Just(p).boxed(),
325 None => any::<ValidPrime>().boxed(),
326 };
327
328 (p, args.dim)
329 .prop_flat_map(move |(p, dim)| {
330 Matrix::arbitrary_rref_with(MatrixArbParams {
331 p: Some(p),
332 rows: (0..=dim + 1).boxed(),
333 columns: Just(dim).boxed(),
334 })
335 })
336 .prop_map(Self::from_matrix)
337 .boxed()
338 }
339 }
340}