Skip to main content

rustfs_erasure_codec/core/
workspace.rs

1extern crate alloc;
2
3use alloc::vec;
4use alloc::vec::Vec;
5
6use crate::Field;
7
8use super::ReedSolomon;
9
10/// Reusable parity scratch space for repeated verify calls.
11///
12/// This helper keeps the parity buffer allocation outside of `verify` so
13/// repeated callers can naturally take the `verify_with_buffer` fast path
14/// without having to manage `Vec<Vec<_>>` details themselves.
15#[derive(PartialEq, Debug, Clone)]
16pub struct VerifyWorkspace<F: Field> {
17    parity: Vec<Vec<F::Elem>>,
18}
19
20impl<F: Field> VerifyWorkspace<F> {
21    /// Create a new workspace with parity buffers sized for the given codec and shard length.
22    pub fn new(codec: &ReedSolomon<F>, shard_len: usize) -> Self {
23        let mut parity = Vec::with_capacity(codec.parity_shard_count);
24        for _ in 0..codec.parity_shard_count {
25            parity.push(vec![F::zero(); shard_len]);
26        }
27        Self { parity }
28    }
29
30    /// Returns the number of parity shard buffers.
31    pub fn parity_shards(&self) -> usize {
32        self.parity.len()
33    }
34
35    /// Returns the current shard buffer length, or `None` if there are no parity shards.
36    pub fn shard_len(&self) -> Option<usize> {
37        self.parity.first().map(Vec::len)
38    }
39
40    /// Resize parity buffers to match the given codec and shard length.
41    pub fn resize(&mut self, codec: &ReedSolomon<F>, shard_len: usize) {
42        if self.parity.len() < codec.parity_shard_count {
43            self.parity
44                .reserve(codec.parity_shard_count - self.parity.len());
45            while self.parity.len() < codec.parity_shard_count {
46                self.parity.push(Vec::new());
47            }
48        } else if self.parity.len() > codec.parity_shard_count {
49            self.parity.truncate(codec.parity_shard_count);
50        }
51
52        for shard in &mut self.parity {
53            shard.resize(shard_len, F::zero());
54        }
55    }
56
57    pub(crate) fn prepare(&mut self, codec: &ReedSolomon<F>, shard_len: usize) {
58        if self.parity.len() != codec.parity_shard_count
59            || self.shard_len() != Some(shard_len)
60            || self.parity.iter().any(|shard| shard.len() != shard_len)
61        {
62            self.resize(codec, shard_len);
63        }
64    }
65
66    pub(crate) fn as_mut_shards(&mut self) -> &mut [Vec<F::Elem>] {
67        &mut self.parity
68    }
69}