1use crate::modular::{mod_add, mod_mul, mod_neg, mod_sub};
7use crate::params::Params;
8
9#[derive(Clone, Debug)]
11pub struct Poly {
12 pub coeffs: Vec<u64>,
14}
15
16impl Poly {
17 pub fn zero(n: usize) -> Self {
19 Self {
20 coeffs: vec![0u64; n],
21 }
22 }
23
24 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 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 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 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 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 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 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 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 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 let p = tiny();
151 let n = p.n;
152 let mut a_coeffs = vec![0u64; n];
154 a_coeffs[n - 1] = 1;
155 let a = Poly { coeffs: a_coeffs };
156 let b = Poly::from_coeffs(&[0, 1], n);
158 let c = a.mul_schoolbook(&b, &p);
159 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}