Skip to main content

tasm_lib/arithmetic/u64/
pow2.rs

1use std::collections::HashMap;
2
3use triton_vm::prelude::*;
4
5use crate::prelude::*;
6use crate::traits::basic_snippet::Reviewer;
7use crate::traits::basic_snippet::SignOffFingerprint;
8
9/// Raise 2 to the power of the given exponent.
10///
11/// ### Behavior
12///
13/// ```text
14/// BEFORE: _ [arg: u32]
15/// AFTER:  _ [2^arg: u64]
16/// ```
17///
18/// ### Preconditions
19///
20/// - the input `arg` is less than 64
21///
22/// ### Postconditions
23///
24/// - the output is properly [`BFieldCodec`] encoded
25#[derive(Debug, Clone, Copy, Eq, PartialEq, Hash)]
26pub struct Pow2;
27
28impl Pow2 {
29    pub const INPUT_TOO_LARGE_ERROR_ID: i128 = 360;
30}
31
32impl BasicSnippet for Pow2 {
33    fn parameters(&self) -> Vec<(DataType, String)> {
34        vec![(DataType::U32, "arg".to_string())]
35    }
36
37    fn return_values(&self) -> Vec<(DataType, String)> {
38        vec![(DataType::U64, "(2^arg)".to_string())]
39    }
40
41    fn entrypoint(&self) -> String {
42        "tasmlib_arithmetic_u64_pow2".to_string()
43    }
44
45    fn code(&self, _: &mut Library) -> Vec<LabelledInstruction> {
46        triton_asm! {
47            {self.entrypoint()}:
48                push 64
49                dup 1
50                lt
51                assert error_id {Self::INPUT_TOO_LARGE_ERROR_ID}
52
53                push 2
54                pow
55                split
56                return
57        }
58    }
59
60    fn sign_offs(&self) -> HashMap<Reviewer, SignOffFingerprint> {
61        let mut sign_offs = HashMap::new();
62        sign_offs.insert(Reviewer("ferdinand"), 0xcb940285f823bded.into());
63        sign_offs
64    }
65}
66
67#[cfg(test)]
68mod tests {
69    use super::*;
70    use crate::test_prelude::*;
71
72    impl Pow2 {
73        fn assert_expected_behavior(&self, arg: u32) {
74            let initial_stack = self.set_up_test_stack(arg);
75
76            let mut expected_stack = initial_stack.clone();
77            self.rust_shadow(&mut expected_stack).unwrap();
78
79            test_rust_equivalence_given_complete_state(
80                &ShadowedClosure::new(Self),
81                &initial_stack,
82                &[],
83                &NonDeterminism::default(),
84                &None,
85                Some(&expected_stack),
86            );
87        }
88    }
89
90    impl Closure for Pow2 {
91        type Args = u32;
92
93        fn rust_shadow(&self, stack: &mut Vec<BFieldElement>) -> Result<(), RustShadowError> {
94            let arg = pop_encodable::<Self::Args>(stack)?;
95            let pow = 2_u64
96                .checked_pow(arg)
97                .ok_or(RustShadowError::ArithmeticOverflow)?;
98            push_encodable(stack, &pow);
99            Ok(())
100        }
101
102        fn pseudorandom_args(
103            &self,
104            seed: [u8; 32],
105            bench_case: Option<BenchmarkCase>,
106        ) -> Self::Args {
107            let Some(bench_case) = bench_case else {
108                return StdRng::from_seed(seed).random_range(0..64);
109            };
110
111            match bench_case {
112                BenchmarkCase::CommonCase => 31,
113                BenchmarkCase::WorstCase => 63,
114            }
115        }
116    }
117
118    #[macro_rules_attr::apply(test)]
119    fn rust_shadow() {
120        ShadowedClosure::new(Pow2).test();
121    }
122
123    #[macro_rules_attr::apply(test)]
124    fn unit_test() {
125        for arg in 0..64 {
126            Pow2.assert_expected_behavior(arg);
127        }
128    }
129
130    #[macro_rules_attr::apply(proptest)]
131    fn negative_property_test(#[strategy(64_u32..)] arg: u32) {
132        let stack = Pow2.set_up_test_stack(arg);
133
134        test_assertion_failure(
135            &ShadowedClosure::new(Pow2),
136            InitVmState::with_stack(stack),
137            &[Pow2::INPUT_TOO_LARGE_ERROR_ID],
138        );
139    }
140}
141
142#[cfg(test)]
143mod benches {
144    use super::*;
145    use crate::test_prelude::*;
146
147    #[macro_rules_attr::apply(test)]
148    fn benchmark() {
149        ShadowedClosure::new(Pow2).bench();
150    }
151}