treams_core/
saved.rs

1//! Numeric derivative state carried by array-only framework callbacks.
2//!
3//! These bytes belong to one invocation of the current native extension. They
4//! are not a persistent file format. Each residual defines a fixed layout from
5//! its input shapes, including canonical storage for optional numerical data.
6//! treams-rs extension.
7
8use nalgebra::DMatrix;
9
10use crate::{Complex, Error, Result, numerics};
11
12/// A residual that can cross an array-only callback boundary without rerunning
13/// its forward calculation.
14pub trait SavedState: Sized {
15    /// Copy the saved numerical state into its fixed byte layout.
16    fn save_state(&self) -> Result<Vec<u8>>;
17
18    /// Restore numerical state from bytes produced by this extension.
19    fn from_state(bytes: &[u8]) -> Result<Self>;
20}
21
22/// Write fixed-width numbers without alignment assumptions or unsafe casts.
23#[derive(Debug)]
24pub struct Writer {
25    bytes: Vec<u8>,
26}
27
28impl Writer {
29    /// Allocate storage for a layout whose size has already been checked.
30    #[must_use]
31    pub fn new(capacity: usize) -> Self {
32        Self {
33            bytes: Vec::with_capacity(capacity),
34        }
35    }
36
37    /// Append a length or index as an unsigned 64-bit integer.
38    pub fn usize(&mut self, value: usize) {
39        self.bytes.extend_from_slice(&(value as u64).to_le_bytes());
40    }
41
42    /// Append one real scalar.
43    pub fn f64(&mut self, value: f64) {
44        self.bytes.extend_from_slice(&value.to_le_bytes());
45    }
46
47    /// Append a complex scalar, real part followed by imaginary part.
48    pub fn complex(&mut self, value: Complex) {
49        self.f64(value.re);
50        self.f64(value.im);
51    }
52
53    /// Append a tag or boolean.
54    pub fn byte(&mut self, value: u8) {
55        self.bytes.push(value);
56    }
57
58    /// Append one signed mode label.
59    pub fn i32(&mut self, value: i32) {
60        self.raw(&value.to_le_bytes());
61    }
62
63    /// Append an already encoded section.
64    pub fn raw(&mut self, bytes: &[u8]) {
65        self.bytes.extend_from_slice(bytes);
66    }
67
68    /// Return the complete byte buffer.
69    #[must_use]
70    pub fn finish(self) -> Vec<u8> {
71        self.bytes
72    }
73}
74
75/// Read a private invocation state; malformed or truncated buffers are errors.
76#[derive(Debug)]
77pub struct Reader<'a> {
78    bytes: &'a [u8],
79}
80
81impl<'a> Reader<'a> {
82    /// Borrow an encoded invocation state.
83    #[must_use]
84    pub const fn new(bytes: &'a [u8]) -> Self {
85        Self { bytes }
86    }
87
88    fn take<const N: usize>(&mut self) -> Result<[u8; N]> {
89        let Some((head, tail)) = self.bytes.split_at_checked(N) else {
90            return Err(invalid());
91        };
92        self.bytes = tail;
93        head.try_into().map_err(|_| invalid())
94    }
95
96    /// Read one length or index.
97    pub fn usize(&mut self) -> Result<usize> {
98        usize::try_from(u64::from_le_bytes(self.take()?)).map_err(|_| invalid())
99    }
100
101    /// Read one real scalar.
102    pub fn f64(&mut self) -> Result<f64> {
103        Ok(f64::from_le_bytes(self.take()?))
104    }
105
106    /// Read one complex scalar.
107    pub fn complex(&mut self) -> Result<Complex> {
108        Ok(Complex::new(self.f64()?, self.f64()?))
109    }
110
111    /// Read one tag or boolean.
112    pub fn byte(&mut self) -> Result<u8> {
113        Ok(u8::from_le_bytes(self.take()?))
114    }
115
116    /// Read one signed mode label.
117    pub fn i32(&mut self) -> Result<i32> {
118        Ok(i32::from_le_bytes(self.take()?))
119    }
120
121    /// Read a section without allocating or copying it.
122    pub fn raw(&mut self, len: usize) -> Result<&'a [u8]> {
123        let Some((head, tail)) = self.bytes.split_at_checked(len) else {
124            return Err(invalid());
125        };
126        self.bytes = tail;
127        Ok(head)
128    }
129
130    /// The number of bytes not yet read.
131    #[must_use]
132    pub const fn remaining_len(&self) -> usize {
133        self.bytes.len()
134    }
135
136    /// Validate an encoded vector count against the remaining storage before allocating.
137    pub fn count(&mut self, stride: usize) -> Result<usize> {
138        let count = self.usize()?;
139        if stride == 0 || count > self.remaining_len() / stride {
140            return Err(invalid());
141        }
142        Ok(count)
143    }
144
145    /// Reject trailing data after the complete expected layout.
146    pub fn finish(self) -> Result<()> {
147        if self.bytes.is_empty() {
148            Ok(())
149        } else {
150            Err(invalid())
151        }
152    }
153}
154
155pub(crate) fn invalid() -> Error {
156    Error::InvalidInput("invalid native derivative state".into())
157}
158
159pub(crate) fn total(parts: &[usize]) -> Result<usize> {
160    parts.iter().try_fold(0_usize, |size, &part| {
161        size.checked_add(part).ok_or_else(invalid)
162    })
163}
164
165pub(crate) fn product(parts: &[usize]) -> Result<usize> {
166    parts.iter().try_fold(1_usize, |size, &part| {
167        size.checked_mul(part).ok_or_else(invalid)
168    })
169}
170
171pub(crate) fn read_complex(reader: &mut Reader<'_>) -> Result<Complex> {
172    let value = reader.complex()?;
173    if !numerics::finite(value) {
174        return Err(invalid());
175    }
176    Ok(value)
177}
178
179pub(crate) fn read_f64(reader: &mut Reader<'_>) -> Result<f64> {
180    let value = reader.f64()?;
181    if !value.is_finite() {
182        return Err(invalid());
183    }
184    Ok(value)
185}
186
187pub(crate) fn write_matrix(writer: &mut Writer, matrix: &DMatrix<Complex>) {
188    for &value in matrix.as_slice() {
189        writer.complex(value);
190    }
191}
192
193pub(crate) fn read_matrix(
194    reader: &mut Reader<'_>,
195    rows: usize,
196    columns: usize,
197) -> Result<DMatrix<Complex>> {
198    let mut matrix = numerics::zeros(rows, columns)?;
199    for value in matrix.as_mut_slice() {
200        *value = read_complex(reader)?;
201    }
202    Ok(matrix)
203}
204
205/// Byte layout of mode labels followed by Cartesian positions, both length prefixed.
206fn basis_bytes(modes: usize, positions: usize, mode_bytes: usize) -> Result<usize> {
207    modes
208        .checked_mul(mode_bytes)
209        .and_then(|n| positions.checked_mul(24).and_then(|p| n.checked_add(p)))
210        .and_then(|n| n.checked_add(16))
211        .ok_or_else(invalid)
212}
213
214pub(crate) fn sw_basis_size(modes: usize, positions: usize) -> Result<usize> {
215    basis_bytes(modes, positions, 17)
216}
217
218pub(crate) fn cw_basis_size(modes: usize, positions: usize) -> Result<usize> {
219    basis_bytes(modes, positions, 21)
220}
221
222fn read_basis_counts(reader: &mut Reader<'_>, mode_bytes: usize) -> Result<(usize, usize)> {
223    let modes = reader.count(mode_bytes)?;
224    reader.raw(modes * mode_bytes)?;
225    let positions = reader.count(24)?;
226    reader.raw(positions * 24)?;
227    Ok((modes, positions))
228}
229
230pub(crate) fn read_sw_basis_dimensions(reader: &mut Reader<'_>) -> Result<(usize, usize)> {
231    read_basis_counts(reader, 17)
232}
233
234pub(crate) fn read_cw_basis_dimensions(reader: &mut Reader<'_>) -> Result<(usize, usize)> {
235    read_basis_counts(reader, 21)
236}
237
238/// Inspect and skip the complete encoded basis without allocating its arrays.
239pub(crate) fn read_basis_dimensions(reader: &mut Reader<'_>) -> Result<(usize, usize, bool)> {
240    let cylindrical = match reader.byte()? {
241        0 => false,
242        1 => true,
243        _ => return Err(invalid()),
244    };
245    let (modes, positions) = read_basis_counts(reader, if cylindrical { 21 } else { 17 })?;
246    Ok((modes, positions, cylindrical))
247}
248
249pub(crate) fn write_positions(writer: &mut Writer, positions: &[[f64; 3]]) {
250    writer.usize(positions.len());
251    for point in positions {
252        for &coordinate in point {
253            writer.f64(coordinate);
254        }
255    }
256}
257
258pub(crate) fn read_positions(reader: &mut Reader<'_>) -> Result<Vec<[f64; 3]>> {
259    (0..reader.count(24)?)
260        .map(|_| Ok([reader.f64()?, reader.f64()?, reader.f64()?]))
261        .collect()
262}
263
264pub(crate) fn write_sw_basis(writer: &mut Writer, basis: &crate::sw::Basis) {
265    writer.usize(basis.modes.len());
266    for &(position, mode) in &basis.modes {
267        writer.usize(position);
268        writer.i32(mode.l);
269        writer.i32(mode.m);
270        writer.byte(mode.pol);
271    }
272    write_positions(writer, &basis.positions);
273}
274
275pub(crate) fn read_sw_basis(reader: &mut Reader<'_>) -> Result<crate::sw::Basis> {
276    let modes = (0..reader.count(17)?)
277        .map(|_| {
278            Ok((
279                reader.usize()?,
280                crate::sw::Mode {
281                    l: reader.i32()?,
282                    m: reader.i32()?,
283                    pol: reader.byte()?,
284                },
285            ))
286        })
287        .collect::<Result<Vec<_>>>()?;
288    let basis = crate::sw::Basis {
289        modes,
290        positions: read_positions(reader)?,
291    };
292    basis.validate()?;
293    Ok(basis)
294}
295
296pub(crate) fn write_cw_basis(writer: &mut Writer, basis: &crate::cw::Basis) {
297    writer.usize(basis.modes.len());
298    for &(position, mode) in &basis.modes {
299        writer.usize(position);
300        writer.f64(mode.kz);
301        writer.i32(mode.m);
302        writer.byte(mode.pol);
303    }
304    write_positions(writer, &basis.positions);
305}
306
307pub(crate) fn read_cw_basis(reader: &mut Reader<'_>) -> Result<crate::cw::Basis> {
308    let modes = (0..reader.count(21)?)
309        .map(|_| {
310            Ok((
311                reader.usize()?,
312                crate::cw::Mode {
313                    kz: reader.f64()?,
314                    m: reader.i32()?,
315                    pol: reader.byte()?,
316                },
317            ))
318        })
319        .collect::<Result<Vec<_>>>()?;
320    let basis = crate::cw::Basis {
321        modes,
322        positions: read_positions(reader)?,
323    };
324    basis.validate()?;
325    Ok(basis)
326}
327
328pub(crate) fn write_basis(writer: &mut Writer, basis: &crate::basis::MultipoleBasis) {
329    match basis {
330        crate::basis::MultipoleBasis::Spherical(basis) => {
331            writer.byte(0);
332            write_sw_basis(writer, basis);
333        }
334        crate::basis::MultipoleBasis::Cylindrical(basis) => {
335            writer.byte(1);
336            write_cw_basis(writer, basis);
337        }
338    }
339}
340
341pub(crate) fn read_basis(reader: &mut Reader<'_>) -> Result<crate::basis::MultipoleBasis> {
342    match reader.byte()? {
343        0 => read_sw_basis(reader).map(Into::into),
344        1 => read_cw_basis(reader).map(Into::into),
345        _ => Err(invalid()),
346    }
347}
348
349pub(crate) fn write_radial(writer: &mut Writer, radial: crate::special::Radial) {
350    writer.byte(u8::from(radial == crate::special::Radial::Singular));
351}
352
353pub(crate) fn read_radial(reader: &mut Reader<'_>) -> Result<crate::special::Radial> {
354    match reader.byte()? {
355        0 => Ok(crate::special::Radial::Regular),
356        1 => Ok(crate::special::Radial::Singular),
357        _ => Err(invalid()),
358    }
359}
treams_core/saved.rs

treams_core/
saved.rs
