tasm_lib/arithmetic/xfe/
mod_pow_u32.rs

1use triton_vm::prelude::*;
2
3use crate::prelude::*;
4
5#[derive(Debug, Copy, Clone, Eq, PartialEq, Hash)]
6pub struct XfeModPowU32;
7
8impl BasicSnippet for XfeModPowU32 {
9    fn parameters(&self) -> Vec<(DataType, String)> {
10        vec![
11            (DataType::U32, "exponent".to_owned()),
12            (DataType::Xfe, "base".to_owned()),
13        ]
14    }
15
16    fn return_values(&self) -> Vec<(DataType, String)> {
17        vec![(DataType::Xfe, "result".to_owned())]
18    }
19
20    fn entrypoint(&self) -> String {
21        "tasmlib_arithmetic_xfe_mod_pow_u32".to_owned()
22    }
23
24    // This implementation is far from optimized, not very efficient. To make a snippet
25    // with a shorter execution trace, you could e.g. implement this algorithm for
26    // statically known exponents.
27    fn code(&self, _library: &mut Library) -> Vec<LabelledInstruction> {
28        // Implemented as:
29        // ```rust
30        // fn mod_pow_u32(exponent: u32, base: XFieldElement) -> Self {
31        //     let mut x = base;
32        //     let mut acc = Self::one();
33        //     let mut i = exponent;
34
35        //     while i > 0 {
36        //         if i & 1 == 1 {
37        //             acc *= x;
38        //         }
39
40        //         x *= x;
41        //         i >>= 1;
42        //     }
43
44        //     acc
45        //
46
47        let entrypoint = self.entrypoint();
48        let loop_code_label = format!("{entrypoint}_loop");
49        let acc_mul_x_label = format!("{entrypoint}_acc_mul_x");
50
51        triton_asm!(
52            {entrypoint}:
53                // _ exponent [base]
54
55                push 0
56                push 0
57                push 1
58                // _ exponent [base] [1]
59
60                // Rename
61                // _ i [x] [acc]
62
63                call {loop_code_label}
64                // _ 0 [x] [result]
65
66                swap 4
67                pop 1
68                swap 4
69                pop 1
70                swap 4
71                pop 2
72                // _ [result]
73
74                return
75
76            // Invariant: i [x] [acc]
77            {loop_code_label}:
78
79                // Return iff i == 0
80                dup 6
81                push 0
82                eq
83                skiz
84                    return
85                // _ i [x] [acc]
86
87                dup 6
88                push 1
89                and
90                // _ i [x] [acc] (i & 1)
91
92                skiz
93                    call {acc_mul_x_label}
94
95                // _ i [x] [acc']
96
97                dup 5 dup 5 dup 5
98                dup 2 dup 2 dup 2
99                // _ i [x] [acc'] [x] [x]
100
101                xx_mul
102                // _ i [x] [acc'] [x * x]
103
104                swap 6
105                pop 1
106                swap 6
107                pop 1
108                swap 6
109                pop 1
110                // _ i [x'] [acc']
111
112                swap 6
113                // _ acc'_0 [x'] acc'_2 acc'_1 i
114
115                push 2
116                swap 1
117                // _ acc'_0 [x'] acc'_2 acc'_1 2 i
118
119                div_mod
120                pop 1
121                // _ acc'_0 [x'] acc'_2 acc'_1 (i / 2)
122
123                swap 6
124                // _ i' [x'] acc'_2 acc'_1 acc'_0
125                // _ i' [x'] [acc']
126
127                recurse
128
129            {acc_mul_x_label}:
130                 // _ [x] [acc]
131
132                 dup 5
133                 dup 5
134                 dup 5
135                 xx_mul
136                 // _ [x] [acc * x]
137
138                 return
139        )
140    }
141}
142
143#[cfg(test)]
144pub mod tests {
145    use twenty_first::math::traits::ModPowU32;
146
147    use super::*;
148    use crate::empty_stack;
149    use crate::execute_with_terminal_state;
150    use crate::test_prelude::*;
151
152    impl Closure for XfeModPowU32 {
153        type Args = (u32, XFieldElement);
154
155        fn rust_shadow(&self, stack: &mut Vec<BFieldElement>) -> Result<(), RustShadowError> {
156            let (exponent, base) = pop_encodable::<Self::Args>(stack)?;
157            let result = base.mod_pow_u32(exponent);
158            push_encodable(stack, &result);
159            Ok(())
160        }
161
162        fn pseudorandom_args(
163            &self,
164            seed: [u8; 32],
165            bench_case: Option<BenchmarkCase>,
166        ) -> Self::Args {
167            let mut rng = StdRng::from_seed(seed);
168            let exponent = match bench_case {
169                Some(BenchmarkCase::CommonCase) => 1 << 25,
170                Some(BenchmarkCase::WorstCase) => u32::MAX,
171                None => rng.random(),
172            };
173
174            (exponent, rng.random())
175        }
176
177        fn corner_case_args(&self) -> Vec<Self::Args> {
178            let an_xfe = xfe!([14; 3]);
179
180            (0..=5)
181                .chain([u32::MAX - 1, u32::MAX])
182                .map(|exp| (exp, an_xfe))
183                .collect()
184        }
185    }
186
187    #[macro_rules_attr::apply(test)]
188    fn mod_pow_u32_xfe_pbt() {
189        ShadowedClosure::new(XfeModPowU32).test()
190    }
191
192    #[macro_rules_attr::apply(test)]
193    fn verify_crash_if_exponent_not_u32() {
194        let bfe_14 = BFieldElement::new(14);
195        let xfe_14 = XFieldElement::new([bfe_14, bfe_14, bfe_14]);
196        let xfe_14: Vec<_> = xfe_14.coefficients.into_iter().rev().collect();
197        let code = XfeModPowU32.link_for_isolated_run();
198
199        for exponent in [
200            1 << 32,
201            1 << 33,
202            1 << 32,
203            1 << 63,
204            BFieldElement::MAX - 1,
205            BFieldElement::MAX,
206        ] {
207            let init_stack = [empty_stack(), bfe_vec![exponent], xfe_14.clone()].concat();
208            let tvm_result = execute_with_terminal_state(
209                Program::new(&code),
210                &[],
211                &init_stack,
212                &NonDeterminism::default(),
213                None,
214            );
215            assert!(matches!(
216                tvm_result.unwrap_err(),
217                InstructionError::OpStackError(OpStackError::FailedU32Conversion(_))
218            ));
219        }
220    }
221}
222
223#[cfg(test)]
224mod benches {
225    use super::*;
226    use crate::test_prelude::*;
227
228    #[macro_rules_attr::apply(test)]
229    fn benchmark() {
230        ShadowedClosure::new(XfeModPowU32).bench();
231    }
232}
tasm_lib/arithmetic/xfe/mod_pow_u32.rs

tasm_lib/arithmetic/xfe/
mod_pow_u32.rs
