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 inputs(&self) -> Vec<(DataType, String)> {
10        vec![
11            (DataType::U32, "exponent".to_owned()),
12            (DataType::Xfe, "base".to_owned()),
13        ]
14    }
15
16    fn outputs(&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>) {
156            let (exponent, base) = pop_encodable::<Self::Args>(stack);
157            let result = base.mod_pow_u32(exponent);
158            push_encodable(stack, &result);
159        }
160
161        fn pseudorandom_args(
162            &self,
163            seed: [u8; 32],
164            bench_case: Option<BenchmarkCase>,
165        ) -> Self::Args {
166            let mut rng = StdRng::from_seed(seed);
167            let exponent = match bench_case {
168                Some(BenchmarkCase::CommonCase) => 1 << 25,
169                Some(BenchmarkCase::WorstCase) => u32::MAX,
170                None => rng.random(),
171            };
172
173            (exponent, rng.random())
174        }
175
176        fn corner_case_args(&self) -> Vec<Self::Args> {
177            let an_xfe = xfe!([14; 3]);
178
179            (0..=5)
180                .chain([u32::MAX - 1, u32::MAX])
181                .map(|exp| (exp, an_xfe))
182                .collect()
183        }
184    }
185
186    #[test]
187    fn mod_pow_u32_xfe_pbt() {
188        ShadowedClosure::new(XfeModPowU32).test()
189    }
190
191    #[test]
192    fn verify_crash_if_exponent_not_u32() {
193        let bfe_14 = BFieldElement::new(14);
194        let xfe_14 = XFieldElement::new([bfe_14, bfe_14, bfe_14]);
195        let xfe_14: Vec<_> = xfe_14.coefficients.into_iter().rev().collect();
196        let code = XfeModPowU32.link_for_isolated_run();
197
198        for exponent in [
199            1 << 32,
200            1 << 33,
201            1 << 32,
202            1 << 63,
203            BFieldElement::MAX - 1,
204            BFieldElement::MAX,
205        ] {
206            let init_stack = [empty_stack(), bfe_vec![exponent], xfe_14.clone()].concat();
207            let tvm_result = execute_with_terminal_state(
208                Program::new(&code),
209                &[],
210                &init_stack,
211                &NonDeterminism::default(),
212                None,
213            );
214            assert!(matches!(
215                tvm_result.unwrap_err(),
216                InstructionError::OpStackError(OpStackError::FailedU32Conversion(_))
217            ));
218        }
219    }
220}
221
222#[cfg(test)]
223mod benches {
224    use super::*;
225    use crate::test_prelude::*;
226
227    #[test]
228    fn benchmark() {
229        ShadowedClosure::new(XfeModPowU32).bench();
230    }
231}
tasm_lib/arithmetic/xfe/mod_pow_u32.rs

tasm_lib/arithmetic/xfe/
mod_pow_u32.rs
