treams_core/smatrix/compose/
saved.rs

1//! Fixed saved state for Redheffer composition.
2//!
3//! treams-rs extension.
4
5use 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    /// Byte count for saved composition state with `n` modes per direction.
17    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}