tasm_lib/arithmetic/u32/
next_power_of_two.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/// Returns the smallest power of two greater than or equal to argument `arg`.
10/// Behaves like the [`rustc` method in debug mode][rustc_pow] for all inputs
11/// of type `u32`.
12///
13/// ### Behavior
14///
15/// ```text
16/// BEFORE: _ arg
17/// AFTER:  _ u32::next_power_of_two(arg)
18/// ```
19///
20/// ### Pre-conditions
21///
22/// - `arg` is a valid u32
23/// - `arg` is smaller than or equal to 2^31
24///
25/// ### Post-conditions
26///
27/// - the output is the smallest power of two greater than or equal to `arg`
28///
29/// [rustc_pow]: u32::next_power_of_two
30#[derive(Debug, Clone, Copy)]
31pub struct NextPowerOfTwo;
32
33impl NextPowerOfTwo {
34    pub const INPUT_TOO_LARGE_ERROR_ID: i128 = 130;
35}
36
37impl BasicSnippet for NextPowerOfTwo {
38    fn parameters(&self) -> Vec<(DataType, String)> {
39        vec![(DataType::U32, "self".to_owned())]
40    }
41
42    fn return_values(&self) -> Vec<(DataType, String)> {
43        vec![(DataType::U32, "next_power_of_two".to_owned())]
44    }
45
46    fn entrypoint(&self) -> String {
47        "tasmlib_arithmetic_u32_next_power_of_two".to_string()
48    }
49
50    fn code(&self, _: &mut Library) -> Vec<LabelledInstruction> {
51        let entrypoint = self.entrypoint();
52        let zero_or_one_label = format!("{entrypoint}_zero_or_one_label");
53        let greater_than_one_label = format!("{entrypoint}_greater_than_one");
54        triton_asm!(
55            {entrypoint}:
56                // _ arg
57
58                push 1
59                push 2
60                dup 2
61                lt
62                // _ arg 1 (arg < 2)
63
64                skiz
65                    call {zero_or_one_label}
66                // if arg < 2:  _ 1   0
67                // if arg >= 2: _ arg 1
68
69                skiz
70                    call {greater_than_one_label}
71                // _
72
73                return
74
75            {zero_or_one_label}:
76                // _ arg 1
77
78                pop 2
79                push 1
80                push 0
81                // _ 1 0
82
83                return
84
85            {greater_than_one_label}:
86                // _ arg
87
88                addi -1
89                log_2_floor
90                addi 1
91                // _ log₂(⌊value - 1⌋ + 1)
92
93                push 2
94                pow
95                // _ 2^log₂(⌊value - 1⌋ + 1)
96
97                /* Assert result *not* 2^{32} */
98                dup 0
99                push {1u64 << 32}
100                eq
101                push 0
102                eq
103                assert error_id {Self::INPUT_TOO_LARGE_ERROR_ID}
104
105                return
106        )
107    }
108
109    fn sign_offs(&self) -> HashMap<Reviewer, SignOffFingerprint> {
110        let mut sign_offs = HashMap::new();
111        sign_offs.insert(Reviewer("ferdinand"), 0x6ee79243e0c65303.into());
112        sign_offs
113    }
114}
115
116#[cfg(test)]
117mod tests {
118    use super::*;
119    use crate::test_prelude::*;
120
121    impl Closure for NextPowerOfTwo {
122        type Args = u32;
123
124        fn rust_shadow(&self, stack: &mut Vec<BFieldElement>) -> Result<(), RustShadowError> {
125            let arg = pop_encodable::<Self::Args>(stack)?;
126            let pow = arg
127                .checked_next_power_of_two()
128                .ok_or(RustShadowError::ArithmeticOverflow)?;
129            push_encodable(stack, &pow);
130
131            Ok(())
132        }
133
134        fn pseudorandom_args(
135            &self,
136            seed: [u8; 32],
137            bench_case: Option<BenchmarkCase>,
138        ) -> Self::Args {
139            match bench_case {
140                Some(BenchmarkCase::CommonCase) => (1 << 27) - 1,
141                Some(BenchmarkCase::WorstCase) => (1 << 31) - 1,
142                None => StdRng::from_seed(seed).next_u32() / 2,
143            }
144        }
145
146        fn corner_case_args(&self) -> Vec<Self::Args> {
147            let small_inputs = 0..=66;
148            let big_valid_inputs = (0..=5).map(|i| (1 << 31) - i);
149
150            small_inputs.chain(big_valid_inputs).collect()
151        }
152    }
153
154    impl NextPowerOfTwo {
155        fn prepare_vm_stack(&self, arg: u32) -> Vec<BFieldElement> {
156            [self.init_stack_for_isolated_run(), bfe_vec![arg]].concat()
157        }
158    }
159
160    #[macro_rules_attr::apply(test)]
161    fn next_power_of_two_pbt() {
162        ShadowedClosure::new(NextPowerOfTwo).test()
163    }
164
165    #[macro_rules_attr::apply(test)]
166    fn npo2_overflow_negative_test() {
167        let greater_two_pow_31 = (1..=5).map(|i| (1 << 31) + i);
168        let smaller_equal_u32_max = (0..=2).map(|i| u32::MAX - i);
169
170        for arg in greater_two_pow_31.chain(smaller_equal_u32_max) {
171            test_assertion_failure(
172                &ShadowedClosure::new(NextPowerOfTwo),
173                InitVmState::with_stack(NextPowerOfTwo.prepare_vm_stack(arg)),
174                &[NextPowerOfTwo::INPUT_TOO_LARGE_ERROR_ID],
175            );
176        }
177    }
178}
179
180#[cfg(test)]
181mod benches {
182    use super::*;
183    use crate::test_prelude::*;
184
185    #[macro_rules_attr::apply(test)]
186    fn benchmark() {
187        ShadowedClosure::new(NextPowerOfTwo).bench()
188    }
189}
tasm_lib/arithmetic/u32/next_power_of_two.rs

tasm_lib/arithmetic/u32/
next_power_of_two.rs
