Skip to main content

hekate_math/fft/
reed_solomon.rs

1// SPDX-License-Identifier: Apache-2.0
2// This file is part of the hekate-math project.
3// Copyright (C) 2026 Andrei Kochergin <andrei@oumuamua.dev>
4// Copyright (C) 2026 Oumuamua Labs <info@oumuamua.dev>.
5//
6// Licensed under the Apache License, Version 2.0 (the "License");
7// you may not use this file except in compliance with the License.
8// You may obtain a copy of the License at
9//
10//     http://www.apache.org/licenses/LICENSE-2.0
11//
12// Unless required by applicable law or agreed to in writing, software
13// distributed under the License is distributed on an "AS IS" BASIS,
14// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
15// See the License for the specific language governing permissions and
16// limitations under the License.
17
18use super::{AdditiveFft, FftError};
19use crate::{BinaryFieldExtras, Flat, HardwareField, PackedFlat};
20
21/// Error returned by the Reed–Solomon encoder.
22#[derive(Debug, Clone, Copy, PartialEq, Eq)]
23#[non_exhaustive]
24pub enum RsError {
25    BadRate { log_k: u32, log_n: u32 },
26    FieldTooSmall { log_n: u32, max_log_n: u32 },
27    BadLength { expected: usize, got: usize },
28}
29
30impl core::fmt::Display for RsError {
31    fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
32        match self {
33            RsError::BadRate { log_k, log_n } => {
34                write!(
35                    f,
36                    "ReedSolomon rate: log_k {log_k} must satisfy 1 <= log_k < log_n {log_n}"
37                )
38            }
39            RsError::FieldTooSmall { log_n, max_log_n } => {
40                write!(
41                    f,
42                    "ReedSolomon log_n {log_n} exceeds the maximum {max_log_n} \
43                     = min(field degree, usize::BITS - 1)"
44                )
45            }
46            RsError::BadLength { expected, got } => {
47                write!(f, "ReedSolomon buffer length {got}, expected {expected}")
48            }
49        }
50    }
51}
52
53impl core::error::Error for RsError {}
54
55impl From<FftError> for RsError {
56    fn from(e: FftError) -> Self {
57        match e {
58            FftError::BadLength { expected, got } => RsError::BadLength { expected, got },
59        }
60    }
61}
62
63/// Systematic Reed–Solomon `[n, k, n-k+1]` encoder
64/// over the additive-FFT subspaces `W_k ⊂ W_n`:
65/// `encode(msg)[..k] == msg` (Lin–Chung–Han 2014).
66pub struct ReedSolomon<F> {
67    fft_k: AdditiveFft<F>,
68    fft_n: AdditiveFft<F>,
69    k: usize,
70    n: usize,
71}
72
73impl<F: BinaryFieldExtras + HardwareField> ReedSolomon<F> {
74    /// `k = 2^log_k` message and `n = 2^log_n` codeword symbols,
75    /// `1 <= log_k < log_n <= min(F::BITS, usize::BITS - 1)`.
76    /// The only allocation is the two twiddle schedules.
77    pub fn new(log_k: u32, log_n: u32) -> Result<Self, RsError> {
78        if log_k < 1 || log_k >= log_n {
79            return Err(RsError::BadRate { log_k, log_n });
80        }
81
82        let max_log_n = F::BITS.min(usize::BITS as usize - 1) as u32;
83
84        if log_n > max_log_n {
85            return Err(RsError::FieldTooSmall { log_n, max_log_n });
86        }
87
88        Ok(Self {
89            fft_k: AdditiveFft::new(log_k),
90            fft_n: AdditiveFft::new(log_n),
91            k: 1usize << log_k,
92            n: 1usize << log_n,
93        })
94    }
95
96    pub fn message_len(&self) -> usize {
97        self.k
98    }
99
100    pub fn codeword_len(&self) -> usize {
101        self.n
102    }
103
104    /// Encode one row: `msg.len() == k`,
105    /// `out.len() == n`, `out[..k] == msg` on return.
106    pub fn encode_scalar(&self, msg: &[Flat<F>], out: &mut [Flat<F>]) -> Result<(), RsError> {
107        self.check(msg.len(), out.len())?;
108
109        out[..self.k].copy_from_slice(msg);
110        out[self.k..].fill(Flat::from_raw(F::ZERO));
111
112        self.fft_k.inverse_scalar(&mut out[..self.k])?;
113        self.fft_n.forward_scalar(out)?;
114
115        Ok(())
116    }
117
118    /// Encode `F::WIDTH` rows in lockstep. Lengths are in
119    /// packed elements: `msg.len() == k`, `out.len() == n`.
120    pub fn encode(&self, msg: &[PackedFlat<F>], out: &mut [PackedFlat<F>]) -> Result<(), RsError> {
121        self.check(msg.len(), out.len())?;
122
123        out[..self.k].copy_from_slice(msg);
124        out[self.k..].fill(PackedFlat::default());
125
126        self.fft_k.inverse(&mut out[..self.k])?;
127        self.fft_n.forward(out)?;
128
129        Ok(())
130    }
131
132    fn check(&self, msg_len: usize, out_len: usize) -> Result<(), RsError> {
133        if msg_len != self.k {
134            return Err(RsError::BadLength {
135                expected: self.k,
136                got: msg_len,
137            });
138        }
139
140        if out_len != self.n {
141            return Err(RsError::BadLength {
142                expected: self.n,
143                got: out_len,
144            });
145        }
146
147        Ok(())
148    }
149}