Skip to main content

plat_core/
poly.rs

1//! Polynomial arithmetic in the ring Z_q[X]/(X^N + 1).
2//!
3//! All operations are modular arithmetic modulo a prime q,
4//! in the negacyclic polynomial ring of dimension N.
5
6use crate::modular::{mod_add, mod_mul, mod_neg, mod_sub};
7use crate::params::Params;
8
9/// A polynomial in Z_q[X]/(X^N + 1), stored as coefficient vector.
10#[derive(Clone, Debug)]
11pub struct Poly {
12    /// Coefficients in [0, q). Length must equal params.n.
13    pub coeffs: Vec<u64>,
14}
15
16impl Poly {
17    /// Create the zero polynomial of dimension n.
18    pub fn zero(n: usize) -> Self {
19        Self {
20            coeffs: vec![0u64; n],
21        }
22    }
23
24    /// Create a polynomial from coefficients (clamped to n, zero-padded if short).
25    pub fn from_coeffs(coeffs: &[u64], n: usize) -> Self {
26        let mut c = vec![0u64; n];
27        let len = coeffs.len().min(n);
28        c[..len].copy_from_slice(&coeffs[..len]);
29        Self { coeffs: c }
30    }
31
32    /// Create a constant polynomial (value in slot 0).
33    pub fn constant(value: u64, n: usize) -> Self {
34        let mut c = vec![0u64; n];
35        c[0] = value;
36        Self { coeffs: c }
37    }
38
39    /// Polynomial addition: (a + b) mod q, coefficient-wise.
40    pub fn add(&self, other: &Poly, params: &Params) -> Poly {
41        debug_assert_eq!(self.coeffs.len(), other.coeffs.len());
42        let coeffs = self
43            .coeffs
44            .iter()
45            .zip(&other.coeffs)
46            .map(|(&a, &b)| mod_add(a, b, params.q))
47            .collect();
48        Poly { coeffs }
49    }
50
51    /// Polynomial subtraction: (a - b) mod q.
52    pub fn sub(&self, other: &Poly, params: &Params) -> Poly {
53        debug_assert_eq!(self.coeffs.len(), other.coeffs.len());
54        let coeffs = self
55            .coeffs
56            .iter()
57            .zip(&other.coeffs)
58            .map(|(&a, &b)| mod_sub(a, b, params.q))
59            .collect();
60        Poly { coeffs }
61    }
62
63    /// Polynomial negation: (-a) mod q.
64    pub fn neg(&self, params: &Params) -> Poly {
65        let coeffs = self.coeffs.iter().map(|&a| mod_neg(a, params.q)).collect();
66        Poly { coeffs }
67    }
68
69    /// Scalar multiplication: (a * scalar) mod q.
70    pub fn scalar_mul(&self, scalar: u64, params: &Params) -> Poly {
71        let coeffs = self
72            .coeffs
73            .iter()
74            .map(|&a| mod_mul(a, scalar, params.q))
75            .collect();
76        Poly { coeffs }
77    }
78
79    /// Schoolbook polynomial multiplication in Z_q[X]/(X^N + 1).
80    /// This is O(N^2) — used for correctness testing and small N.
81    pub fn mul_schoolbook(&self, other: &Poly, params: &Params) -> Poly {
82        let n = params.n;
83        let q = params.q;
84        let mut result = vec![0u64; n];
85
86        for i in 0..n {
87            if self.coeffs[i] == 0 {
88                continue;
89            }
90            for j in 0..n {
91                if other.coeffs[j] == 0 {
92                    continue;
93                }
94                let prod = mod_mul(self.coeffs[i], other.coeffs[j], q);
95                let idx = i + j;
96                if idx < n {
97                    result[idx] = mod_add(result[idx], prod, q);
98                } else {
99                    // X^N ≡ -X^0 in negacyclic ring
100                    let idx = idx - n;
101                    result[idx] = mod_sub(result[idx], prod, q);
102                }
103            }
104        }
105
106        Poly { coeffs: result }
107    }
108}
109
110#[cfg(test)]
111mod tests {
112    use super::*;
113
114    fn tiny() -> Params {
115        Params::test_tiny()
116    }
117
118    #[test]
119    fn test_add_sub() {
120        let p = tiny();
121        let a = Poly::from_coeffs(&[1, 2, 3], p.n);
122        let b = Poly::from_coeffs(&[4, 5, 6], p.n);
123        let sum = a.add(&b, &p);
124        assert_eq!(sum.coeffs[0], 5);
125        assert_eq!(sum.coeffs[1], 7);
126        assert_eq!(sum.coeffs[2], 9);
127
128        let diff = sum.sub(&b, &p);
129        assert_eq!(diff.coeffs[0], 1);
130        assert_eq!(diff.coeffs[1], 2);
131        assert_eq!(diff.coeffs[2], 3);
132    }
133
134    #[test]
135    fn test_negacyclic_mul() {
136        let p = tiny();
137        // (1 + X) * (1 + X) = 1 + 2X + X^2 in normal ring
138        // but X^N wraps, so for small coefficients this is fine
139        let a = Poly::from_coeffs(&[1, 1], p.n);
140        let b = Poly::from_coeffs(&[1, 1], p.n);
141        let c = a.mul_schoolbook(&b, &p);
142        assert_eq!(c.coeffs[0], 1);
143        assert_eq!(c.coeffs[1], 2);
144        assert_eq!(c.coeffs[2], 1);
145    }
146
147    #[test]
148    fn test_negacyclic_wraparound() {
149        // In Z_q[X]/(X^N + 1), X^{N-1} * X = X^N = -1
150        let p = tiny();
151        let n = p.n;
152        // a = X^{N-1}
153        let mut a_coeffs = vec![0u64; n];
154        a_coeffs[n - 1] = 1;
155        let a = Poly { coeffs: a_coeffs };
156        // b = X
157        let b = Poly::from_coeffs(&[0, 1], n);
158        let c = a.mul_schoolbook(&b, &p);
159        // X^N = -1 mod (X^N + 1), so constant term should be q-1 (-1 mod q)
160        assert_eq!(c.coeffs[0], p.q - 1);
161        for i in 1..n {
162            assert_eq!(c.coeffs[i], 0);
163        }
164    }
165
166    #[test]
167    fn test_scalar_mul() {
168        let p = tiny();
169        let a = Poly::from_coeffs(&[3, 5, 7], p.n);
170        let b = a.scalar_mul(10, &p);
171        assert_eq!(b.coeffs[0], 30);
172        assert_eq!(b.coeffs[1], 50);
173        assert_eq!(b.coeffs[2], 70);
174    }
175}