treams_core/smatrix/compose/
saved.rs1use super::AddResidual;
6use crate::{
7 Result,
8 saved::{Reader, SavedState, Writer},
9 smatrix::{
10 saved::{invalid, product, read_matrix, total, write_matrix},
11 solve::InternalSolve,
12 },
13};
14
15impl AddResidual {
16 pub fn state_size(n: usize) -> Result<usize> {
18 let columns = product(&[2, n])?;
19 total(&[
20 8,
21 product(&[96, n, n])?,
22 InternalSolve::state_size(n, columns)?,
23 ])
24 }
25}
26
27impl SavedState for AddResidual {
28 fn save_state(&self) -> Result<Vec<u8>> {
29 let n = self.down.nrows();
30 let mut writer = Writer::new(Self::state_size(n)?);
31 writer.usize(n);
32 for matrix in self.lower.iter().chain(&self.upper).chain([&self.down]) {
33 write_matrix(&mut writer, matrix);
34 }
35 writer.raw(&self.solve.save_state()?);
36 Ok(writer.finish())
37 }
38
39 fn from_state(bytes: &[u8]) -> Result<Self> {
40 let mut reader = Reader::new(bytes);
41 let n = reader.usize()?;
42 if bytes.len() != Self::state_size(n)? {
43 return Err(invalid());
44 }
45 let columns = product(&[2, n])?;
46 let lower = [
47 read_matrix(&mut reader, n, n)?,
48 read_matrix(&mut reader, n, n)?,
49 ];
50 let upper = [
51 read_matrix(&mut reader, n, n)?,
52 read_matrix(&mut reader, n, n)?,
53 ];
54 let down = read_matrix(&mut reader, n, columns)?;
55 let solve = InternalSolve::from_state(reader.raw(InternalSolve::state_size(n, columns)?)?)?;
56 if solve.value.shape() != (n, columns) {
57 return Err(invalid());
58 }
59 reader.finish()?;
60 Ok(Self {
61 lower,
62 upper,
63 solve,
64 down,
65 })
66 }
67}