Skip to main content

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}