1use std::collections::BTreeMap;
2
3use anyhow::anyhow;
4use fp::{
5 prime::{Prime, ValidPrime},
6 vector::FpVector,
7};
8
9use crate::{
10 algebra::{AdemAlgebra, Algebra, MilnorAlgebra, adem_algebra::AdemBasisElement},
11 milnor_algebra::{MilnorBasisElement, PPartEntry},
12 steenrod_parser::*,
13};
14
15pub struct SteenrodEvaluator {
16 pub adem: AdemAlgebra,
17 pub milnor: MilnorAlgebra,
18}
19
20impl SteenrodEvaluator {
21 pub fn new(p: ValidPrime) -> Self {
22 Self {
23 adem: AdemAlgebra::new(p, false),
24 milnor: MilnorAlgebra::new(p, false),
25 }
26 }
27
28 pub fn milnor_to_adem(&self, result: &mut FpVector, coeff: u32, degree: i32, input: &FpVector) {
29 let p = self.prime();
30 for (i, v) in input.iter_nonzero() {
31 self.milnor_to_adem_on_basis(result, (coeff * v) % p, degree, i);
32 }
33 }
34
35 pub fn adem_to_milnor(&self, result: &mut FpVector, coeff: u32, degree: i32, input: &FpVector) {
36 let p = self.prime();
37 for (i, v) in input.iter_nonzero() {
38 self.adem_to_milnor_on_basis(result, (coeff * v) % p, degree, i);
39 }
40 }
41
42 pub fn evaluate_algebra_adem(&self, input: &str) -> anyhow::Result<(i32, FpVector)> {
43 self.evaluate_algebra_node(None, parse_algebra(input)?)
44 }
45
46 pub fn evaluate_algebra_milnor(&self, input: &str) -> anyhow::Result<(i32, FpVector)> {
47 let adem_result = self.evaluate_algebra_adem(input);
48 if let Ok((degree, adem_vector)) = adem_result {
49 let mut milnor_vector = FpVector::new(adem_vector.prime(), adem_vector.len());
50 self.adem_to_milnor(&mut milnor_vector, 1, degree, &adem_vector);
51 Ok((degree, milnor_vector))
52 } else {
53 adem_result
54 }
55 }
56
57 pub fn evaluate_module_adem(
60 &self,
61 items: &str,
62 ) -> anyhow::Result<BTreeMap<String, (i32, FpVector)>> {
63 let mut result: BTreeMap<String, (i32, FpVector)> = BTreeMap::new();
64 if items.is_empty() {
65 return Ok(result);
66 }
67 for (op, g) in parse_module(items)? {
68 if let Some((deg, vec)) = result.get_mut(&g) {
69 let (_, adem_v) = self.evaluate_algebra_node(Some(*deg), op)?;
70 vec.add(&adem_v, 1);
71 } else {
72 let (deg, adem_v) = self.evaluate_algebra_node(None, op)?;
73 result.insert(g, (deg, adem_v));
74 }
75 }
76 Ok(result)
77 }
78
79 fn prime(&self) -> ValidPrime {
80 self.adem.prime()
81 }
82
83 fn compute_basis(&self, degree: i32) {
84 self.adem.compute_basis(degree);
85 self.milnor.compute_basis(degree);
86 }
87
88 fn dimension(&self, degree: i32) -> usize {
89 self.adem.dimension(degree)
90 }
91
92 fn evaluate_algebra_node(
93 &self,
94 mut output_degree: Option<i32>,
95 tree: AlgebraNode,
96 ) -> anyhow::Result<(i32, FpVector)> {
97 let p = self.prime();
98 match tree {
99 AlgebraNode::Sum(left, right) => {
100 let (degree, mut output_left) = self.evaluate_algebra_node(output_degree, *left)?;
101 let (_, output_right) = self.evaluate_algebra_node(Some(degree), *right)?;
102 output_left += &output_right;
103 Ok((degree, output_left))
104 }
105 AlgebraNode::Product(left, right) => {
106 let (degree_left, output_left) = self.evaluate_algebra_node(None, *left)?;
107 if let Some(degree) = output_degree {
108 if degree < degree_left {
109 return Err(anyhow!("Mismatched degree"));
110 }
111 output_degree = Some(degree - degree_left);
112 }
113 let (degree_right, output_right) =
114 self.evaluate_algebra_node(output_degree, *right)?;
115 let degree = degree_left + degree_right;
116 self.compute_basis(degree);
117 let mut result = FpVector::new(p, self.adem.dimension(degree));
118 self.adem.multiply_element_by_element(
119 result.as_slice_mut(),
120 1,
121 degree_left,
122 output_left.as_slice(),
123 degree_right,
124 output_right.as_slice(),
125 );
126 Ok((degree, result))
127 }
128 AlgebraNode::BasisElt(basis_elt) => {
129 self.evaluate_basis_element(output_degree, basis_elt)
130 }
131 AlgebraNode::Scalar(x) => {
132 if let Some(degree) = output_degree
133 && degree != 0
134 {
135 return Err(anyhow!("Mismatched Degree"));
136 }
137 let mut result = FpVector::new(p, 1);
138 result.set_entry(0, x.rem_euclid(p.as_i32()) as u32);
139 Ok((0, result))
140 }
141 }
142 }
143
144 fn evaluate_basis_element(
145 &self,
146 output_degree: Option<i32>,
147 basis_elt: AlgebraBasisElt,
148 ) -> anyhow::Result<(i32, FpVector)> {
149 let p = self.prime();
150 let q = self.adem.q();
151 let (degree, result) = match basis_elt {
152 AlgebraBasisElt::AList(p_or_b_list) => self.evaluate_p_or_b_list(&p_or_b_list),
153 AlgebraBasisElt::PList(p_list) => {
154 let degree = std::iter::zip(crate::algebra::combinatorics::xi_degrees(p), &p_list)
155 .map(|(&a, &b)| a * b as i32)
156 .sum::<i32>()
157 * q;
158 let elt = MilnorBasisElement {
159 degree,
160 p_part: p_list,
161 q_part: 0,
162 };
163
164 self.compute_basis(degree);
165 let mut result = FpVector::new(p, self.dimension(degree));
166 self.milnor_to_adem_on_basis(
167 &mut result,
168 1,
169 degree,
170 self.milnor.basis_element_to_index(&elt),
171 );
172 (degree, result)
173 }
174 AlgebraBasisElt::P(x) => {
175 self.compute_basis(x as i32 * q);
176 let (degree, idx) = self.adem.beps_pn(0, x);
177 let mut result = FpVector::new(p, self.dimension(degree));
178 result.set_entry(idx, 1);
179 (degree, result)
180 }
181 AlgebraBasisElt::Q(x) => {
182 let tau_degrees = crate::algebra::combinatorics::tau_degrees(p);
183 let degree = tau_degrees[x as usize];
184 self.compute_basis(degree);
185 let mut result = FpVector::new(p, self.dimension(degree));
186 self.adem_q(&mut result, 1, x);
187 (degree, result)
188 }
189 };
190 if let Some(requested_degree) = output_degree
191 && degree != requested_degree
192 {
193 return Err(anyhow!("Mismatched degree"));
194 }
195 Ok((degree, result))
196 }
197
198 fn adem_to_milnor_on_basis(&self, result: &mut FpVector, coeff: u32, degree: i32, idx: usize) {
202 let elt = self.adem.basis_element_from_index(degree, idx);
203 let p = self.prime();
204 let dim = self.dimension(elt.degree);
205 if dim == 1 {
206 result.set_entry(0, coeff);
207 return;
208 }
209 let mut tmp_vector_a = FpVector::new(p, 1);
210 let mut tmp_vector_b = FpVector::new(p, 0);
211
212 tmp_vector_a.set_entry(0, 1);
213
214 let mut bocksteins = elt.bocksteins;
215 let mut total_degree = 0;
216
217 for &sqn in &elt.ps {
218 let (deg, idx) = self.milnor.beps_pn(bocksteins & 1, sqn as PPartEntry);
219 bocksteins >>= 1;
220 self.compute_basis(total_degree + deg);
221
222 tmp_vector_b.set_scratch_vector_size(self.dimension(total_degree + deg));
223 self.milnor.multiply_element_by_basis_element(
224 tmp_vector_b.as_slice_mut(),
225 1,
226 total_degree,
227 tmp_vector_a.as_slice(),
228 deg,
229 idx,
230 );
231 total_degree += deg;
232 std::mem::swap(&mut tmp_vector_a, &mut tmp_vector_b);
233 }
234 if bocksteins & 1 == 0 {
235 result.add(&tmp_vector_a, coeff);
236 } else {
237 self.milnor.multiply_element_by_basis_element(
238 result.as_slice_mut(),
239 coeff,
240 total_degree,
241 tmp_vector_a.as_slice(),
242 1,
243 0,
244 );
245 }
246 }
247
248 fn milnor_to_adem_on_basis(&self, result: &mut FpVector, coeff: u32, degree: i32, idx: usize) {
251 if self.milnor.generic() {
252 self.milnor_to_adem_on_basis_generic(result, coeff, degree, idx);
253 } else {
254 self.milnor_to_adem_on_basis_2(result, coeff, degree, idx);
255 }
256 }
257
258 fn milnor_to_adem_on_basis_2(
259 &self,
260 result: &mut FpVector,
261 coeff: u32,
262 degree: i32,
263 idx: usize,
264 ) {
265 let elt = self.milnor.basis_element_from_index(degree, idx);
266 let p = self.prime();
267 let dim = self.dimension(elt.degree);
268 if dim == 1 {
269 result.set_entry(0, coeff);
270 return;
271 }
272 let mut t: Vec<u32> = vec![0; elt.p_part.len()];
273 t[elt.p_part.len() - 1] = elt.p_part[elt.p_part.len() - 1];
274 for i in (0..elt.p_part.len() - 1).rev() {
275 t[i] = elt.p_part[i] + 2 * t[i + 1];
276 }
277 let t_idx = self.adem.basis_element_to_index(&AdemBasisElement {
278 degree,
279 bocksteins: 0,
280 ps: t,
281 p_or_sq: p != 2,
282 });
283 let mut tmp_vector_a = FpVector::new(p, dim);
284 self.adem_to_milnor_on_basis(&mut tmp_vector_a, 1, degree, t_idx);
285 assert!(tmp_vector_a.entry(idx) == 1);
286 tmp_vector_a.set_entry(idx, 0);
287 self.milnor_to_adem(result, coeff, degree, &tmp_vector_a);
288 result.add_basis_element(t_idx, coeff);
289 }
290
291 fn milnor_to_adem_on_basis_generic(
292 &self,
293 result: &mut FpVector,
294 coeff: u32,
295 degree: i32,
296 idx: usize,
297 ) {
298 let elt = self.milnor.basis_element_from_index(degree, idx);
299 let p = self.prime();
300 let dim = self.dimension(elt.degree);
301 if dim == 1 {
302 result.set_entry(0, coeff);
303 return;
304 }
305 let t_len = std::cmp::max(
306 elt.p_part.len(),
307 (31u32.saturating_sub(elt.q_part.leading_zeros())) as usize,
308 );
309 let mut t = vec![0; t_len];
310 let last_p_part = if t_len <= elt.p_part.len() {
311 elt.p_part[t_len - 1]
312 } else {
313 0
314 };
315 t[t_len - 1] = last_p_part + ((elt.q_part >> (t_len)) & 1);
316 for i in (0..t_len - 1).rev() {
317 let p_part = if i < elt.p_part.len() {
318 elt.p_part[i]
319 } else {
320 0
321 };
322 t[i] = p_part + ((elt.q_part >> (i + 1)) & 1) + p * t[i + 1];
323 }
324 let t_idx = self.adem.basis_element_to_index(&AdemBasisElement {
325 degree,
326 bocksteins: elt.q_part,
327 ps: t,
328 p_or_sq: p != 2,
329 });
330 let mut tmp_vector_a = FpVector::new(p, dim);
331 self.adem_to_milnor_on_basis(&mut tmp_vector_a, 1, degree, t_idx);
332 assert!(tmp_vector_a.entry(idx) == 1);
333 tmp_vector_a.set_entry(idx, 0);
334 tmp_vector_a.scale(p - 1);
335 self.milnor_to_adem(result, coeff, degree, &tmp_vector_a);
336 result.add_basis_element(t_idx, coeff);
337 }
338
339 fn adem_q(&self, result: &mut FpVector, coeff: u32, qi: u32) {
341 let p = self.prime();
342 let degree = crate::algebra::combinatorics::tau_degrees(p)[qi as usize];
343 let mbe = if self.adem.generic() {
344 MilnorBasisElement {
345 degree,
346 q_part: 1 << qi,
347 p_part: vec![],
348 }
349 } else {
350 let mut p_part = vec![0; qi as usize + 1];
351 p_part[qi as usize] = 1;
352 MilnorBasisElement {
353 degree,
354 q_part: 0,
355 p_part,
356 }
357 };
358 let idx = self.milnor.basis_element_to_index(&mbe);
359 self.milnor_to_adem_on_basis(result, coeff, degree, idx);
360 }
361
362 fn evaluate_p_or_b_list(&self, list: &[BocksteinOrSq]) -> (i32, FpVector) {
363 let p = self.prime();
364 let q = self.adem.q();
365
366 let mut total_degree = 0;
367
368 let mut tmp_vector_a = FpVector::new(p, 1);
369 let mut tmp_vector_b = FpVector::new(p, 0);
370
371 tmp_vector_a.set_entry(0, 1);
372
373 for item in list {
374 let cur_elt = item.to_adem_basis_elt(q);
375
376 self.compute_basis(total_degree + cur_elt.degree);
377 tmp_vector_b.set_scratch_vector_size(self.dimension(total_degree + cur_elt.degree));
378 self.adem.multiply_element_by_basis_element(
379 tmp_vector_b.as_slice_mut(),
380 1,
381 total_degree,
382 tmp_vector_a.as_slice(),
383 cur_elt.degree,
384 self.adem.basis_element_to_index(&cur_elt),
385 );
386 total_degree += cur_elt.degree;
387 std::mem::swap(&mut tmp_vector_a, &mut tmp_vector_b);
388 }
389 (total_degree, tmp_vector_a)
390 }
391}
392
393#[cfg(test)]
394mod tests {
395 use expect_test::{Expect, expect};
396 use rstest::rstest;
397
398 use super::*;
399
400 #[test]
401 fn test_evaluate_2() {
402 let ev = SteenrodEvaluator::new(ValidPrime::new(2));
403
404 let check = |input, adem_output: Expect, milnor_output: Expect| {
405 let (degree, result) = ev.evaluate_algebra_adem(input).unwrap();
406 adem_output.assert_eq(&ev.adem.element_to_string(degree, result.as_slice()));
407
408 let (degree, result) = ev.evaluate_algebra_milnor(input).unwrap();
409 milnor_output.assert_eq(&ev.milnor.element_to_string(degree, result.as_slice()));
410 };
411
412 check(
413 "Sq2 * Sq2",
414 expect![[r#"Sq3 Sq1"#]],
415 expect![[r#"P(1, 1)"#]],
416 );
417 check("A(2 2)", expect![[r#"Sq3 Sq1"#]], expect![[r#"P(1, 1)"#]]);
418 check(
419 "Sq2 * Sq2 * Sq2 + Sq3 * Sq3",
420 expect![[r#"0"#]],
421 expect![[r#"0"#]],
422 );
423 check(
424 "Sq2 * (Sq2 * Sq2 + Sq4)",
425 expect![[r#"Sq6"#]],
426 expect![[r#"P(6)"#]],
427 );
428 check(
429 "Sq7 + Q2",
430 expect![[r#"Sq5 Sq2 + Sq6 Sq1 + Sq4 Sq2 Sq1"#]],
431 expect![[r#"P(7) + P(0, 0, 1)"#]],
432 );
433 check(
434 "(Q2 + Sq7) * Q1",
435 expect![[r#"Sq6 Sq3 Sq1"#]],
436 expect![[r#"P(7, 1) + P(3, 0, 1) + P(0, 1, 1)"#]],
437 );
438 }
439
440 #[test]
441 fn test_evaluate_3() {
442 let ev = SteenrodEvaluator::new(ValidPrime::new(3));
443
444 let check = |input, adem_output: Expect, milnor_output: Expect| {
445 let (degree, result) = ev.evaluate_algebra_adem(input).unwrap();
446 adem_output.assert_eq(&ev.adem.element_to_string(degree, result.as_slice()));
447
448 let (degree, result) = ev.evaluate_algebra_milnor(input).unwrap();
449 milnor_output.assert_eq(&ev.milnor.element_to_string(degree, result.as_slice()));
450 };
451
452 check("P1 * P1", expect![[r#"2 * P2"#]], expect![[r#"2 * P(2)"#]]);
453 check("A(1 1)", expect![[r#"2 * P2"#]], expect![[r#"2 * P(2)"#]]);
454 check(
455 "A(1 b 1)",
456 expect![[r#"b P2 + P2 b"#]],
457 expect![[r#"2 * Q_0 P(2) + Q_1 P(1)"#]],
458 );
459 check(
460 "A(4 2)",
461 expect![[r#"2 * P5 P1"#]],
462 expect![[r#"2 * P(2, 1)"#]],
463 );
464 check("A(5 2)", expect![[r#"0"#]], expect![[r#"0"#]]);
465 check(
466 "A(6 2)",
467 expect![[r#"P6 P2"#]],
468 expect![[r#"P(8) + P(4, 1) + P(0, 2)"#]],
469 );
470 }
471
472 #[rstest(p, max_degree, case(2, 32), case(3, 60))]
473 #[trace]
474 fn test_cob_adem_to_milnor(p: u32, max_degree: i32) {
475 let p = ValidPrime::new(p);
476 let ev = SteenrodEvaluator::new(p);
477 ev.compute_basis(max_degree);
478
479 for degree in 0..max_degree {
480 println!("degree : {degree}");
481 let dim = ev.dimension(degree);
482 let mut milnor_result = FpVector::new(p, dim);
483 let mut adem_result = FpVector::new(p, dim);
484 for i in 0..dim {
485 ev.milnor_to_adem_on_basis(&mut adem_result, 1, degree, i);
487 ev.adem_to_milnor(&mut milnor_result, 1, degree, &adem_result);
488 assert!(
489 milnor_result.entry(i) == 1,
490 "{} ==> {} ==> {}",
491 ev.milnor.basis_element_to_string(degree, i),
492 ev.adem.element_to_string(degree, adem_result.as_slice()),
493 ev.milnor
494 .element_to_string(degree, milnor_result.as_slice())
495 );
496 milnor_result.set_entry(i, 0);
497 assert!(
498 milnor_result.is_zero(),
499 "{} ==> {} ==> {}",
500 ev.milnor.basis_element_to_string(degree, i),
501 ev.adem.element_to_string(degree, adem_result.as_slice()),
502 ev.milnor
503 .element_to_string(degree, milnor_result.as_slice())
504 );
505 adem_result.set_to_zero();
506 milnor_result.set_to_zero();
507 }
508 }
509 }
510}