minroot_cat/schedule.rs
1//! Exponent scanning as a catamorphism over [`Stream`].
2//!
3//! The fifth-root computation `x^((4p-3)/5)` processes the exponent
4//! bit-by-bit via square-and-multiply. Each bit produces a control
5//! signal telling the pipeline whether to square only or square-and-multiply.
6//!
7//! This is a **catamorphism** (fold) over the exponent bits, producing
8//! a [`Stream`] of [`RoundControl`] signals. Equivalently, it is an
9//! **anamorphism** (unfold) that generates the control stream from the
10//! exponent state.
11//!
12//! [`Stream`]: comp_cat_rs::effect::stream::Stream
13
14use std::sync::Arc;
15
16use comp_cat_rs::effect::io::Io;
17use comp_cat_rs::effect::stream::Stream;
18use minroot_core::field::Curve;
19
20/// Control signal for a single pipeline round.
21///
22/// Tells the multiply stage whether to multiply by the base or bypass.
23#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
24pub enum RoundControl {
25 /// Square only (exponent bit is 0): bypass the multiplier.
26 SquareOnly,
27 /// Square and multiply (exponent bit is 1): multiply by base.
28 SquareAndMultiply,
29}
30
31impl RoundControl {
32 /// Returns `true` if this round includes a multiplication.
33 #[must_use]
34 pub fn is_multiply(self) -> bool {
35 self == Self::SquareAndMultiply
36 }
37}
38
39/// State for the exponent scanner unfold.
40///
41/// Tracks the current bit position as we scan the exponent
42/// from MSB to LSB.
43#[derive(Debug, Clone)]
44struct ExponentState {
45 /// The exponent limbs (little-endian).
46 limbs: [u64; 4],
47 /// Current bit position (counts down from `num_bits - 1` to 0).
48 /// `None` means we have finished scanning.
49 position: Option<usize>,
50}
51
52/// Generates the control signal stream for a fifth-root computation.
53///
54/// Scans the exponent `(4p - 3) / 5` bit-by-bit from MSB to LSB,
55/// producing a [`Stream`] of [`RoundControl`] values. This is the
56/// anamorphism (unfold) that drives the pipeline FSM.
57///
58/// The stream produces exactly `curve.exponent_bits()` elements.
59#[must_use]
60pub fn exponent_schedule(curve: Curve) -> Stream<core::convert::Infallible, RoundControl> {
61 let limbs = curve.fifth_root_exponent();
62 let num_bits = curve.exponent_bits();
63 let init = ExponentState {
64 limbs,
65 position: num_bits.checked_sub(1),
66 };
67
68 Stream::unfold(
69 init,
70 Arc::new(|state: ExponentState| {
71 Io::pure(state.position.map(|pos| {
72 let limb_idx = pos / 64;
73 let bit_idx = pos % 64;
74 let bit = (state.limbs[limb_idx] >> bit_idx) & 1;
75 let control = if bit == 1 {
76 RoundControl::SquareAndMultiply
77 } else {
78 RoundControl::SquareOnly
79 };
80 let next_pos = pos.checked_sub(1);
81 let next_state = ExponentState {
82 limbs: state.limbs,
83 position: next_pos,
84 };
85 (control, next_state)
86 }))
87 }),
88 )
89}
90
91/// Counts the number of multiply rounds in the exponent schedule.
92///
93/// This equals the Hamming weight of the exponent, which determines
94/// the average pipeline utilization (multiplies are more expensive
95/// than square-only rounds in terms of power consumption).
96#[must_use]
97pub fn hamming_weight(curve: Curve) -> Io<core::convert::Infallible, usize> {
98 exponent_schedule(curve).fold(
99 0usize,
100 Arc::new(|count, ctrl| {
101 if ctrl.is_multiply() {
102 count + 1
103 } else {
104 count
105 }
106 }),
107 )
108}
109
110#[cfg(test)]
111mod tests {
112 use super::*;
113
114 /// Unwraps an infallible result via `into_ok`.
115 fn infallible_ok<T>(r: Result<T, core::convert::Infallible>) -> T {
116 r.unwrap_or_else(|e| match e {})
117 }
118
119 #[test]
120 fn schedule_length_matches_exponent_bits() {
121 let count = infallible_ok(
122 exponent_schedule(Curve::Pallas)
123 .fold(0usize, Arc::new(|n, _| n + 1))
124 .run(),
125 );
126 assert_eq!(count, 254);
127 }
128
129 #[test]
130 fn first_bit_is_square_and_multiply_for_pallas() {
131 // MSB of Pallas exponent: 0x33... = 0b0011_0011...
132 // Bit 253 (MSB of 254-bit number) is 1.
133 let first = infallible_ok(
134 exponent_schedule(Curve::Pallas)
135 .take(1)
136 .collect()
137 .run(),
138 );
139 assert_eq!(first[0], RoundControl::SquareAndMultiply);
140 }
141
142 #[test]
143 fn hamming_weight_is_positive() {
144 let w = infallible_ok(hamming_weight(Curve::Pallas).run());
145 assert!(w > 0);
146 }
147}