treams_core/smatrix/illuminate/
saved.rs

1//! Fixed saved state for external and internal illumination fields.
2//!
3//! treams-rs extension.
4
5use super::IlluminateResidual;
6use crate::{
7    Result,
8    saved::{Reader, SavedState, Writer},
9    smatrix::{
10        saved::{invalid, product, read_matrix, read_stored, total, write_matrix, write_stored},
11        solve::InternalSolve,
12    },
13};
14
15impl IlluminateResidual {
16    /// Byte count for saved illumination state with mode and illumination counts.
17    pub fn state_size(n: usize, columns: usize) -> Result<usize> {
18        total(&[
19            24,
20            product(&[128, n, n])?,
21            product(&[48, n, columns])?,
22            InternalSolve::state_size(n, columns)?,
23        ])
24    }
25}
26
27impl SavedState for IlluminateResidual {
28    fn save_state(&self) -> Result<Vec<u8>> {
29        let (n, columns) = self.shape();
30        let mut writer = Writer::new(Self::state_size(n, columns)?);
31        writer.usize(n);
32        writer.usize(columns);
33        for block in self.lower.iter().chain(&self.upper) {
34            write_stored(&mut writer, block);
35        }
36        for matrix in self.incoming.iter().chain([&self.down]) {
37            write_matrix(&mut writer, matrix);
38        }
39        writer.raw(&self.solve.save_state()?);
40        Ok(writer.finish())
41    }
42
43    fn from_state(bytes: &[u8]) -> Result<Self> {
44        let mut reader = Reader::new(bytes);
45        let (n, columns) = (reader.usize()?, reader.usize()?);
46        if bytes.len() != Self::state_size(n, columns)? {
47            return Err(invalid());
48        }
49        let mut block = || read_stored(&mut reader, n);
50        let lower = [block()?, block()?, block()?, block()?];
51        let upper = [block()?, block()?, block()?, block()?];
52        let mut matrix = || read_matrix(&mut reader, n, columns);
53        let incoming = [matrix()?, matrix()?];
54        let down = matrix()?;
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            incoming,
64            solve,
65            down,
66        })
67    }
68}