tasm_lib/arithmetic/u32/
leading_zeros.rs

1use triton_vm::prelude::*;
2
3use crate::prelude::*;
4
5#[derive(Debug, Copy, Clone, Eq, PartialEq, Hash)]
6pub struct LeadingZeros;
7
8impl BasicSnippet for LeadingZeros {
9    fn inputs(&self) -> Vec<(DataType, String)> {
10        vec![(DataType::U32, "arg".to_string())]
11    }
12
13    fn outputs(&self) -> Vec<(DataType, String)> {
14        vec![(DataType::U32, "leading_zeros(arg)".to_string())]
15    }
16
17    fn entrypoint(&self) -> String {
18        "tasmlib_arithmetic_u32_leading_zeros".to_string()
19    }
20
21    fn code(&self, _: &mut Library) -> Vec<LabelledInstruction> {
22        let entrypoint = self.entrypoint();
23        let non_zero_label = format!("{entrypoint}_non_zero");
24
25        triton_asm! {
26            // BEFORE: _ value
27            // AFTER:  _ (leading zeros in value)
28            {entrypoint}:
29                dup 0
30                skiz
31                  call {non_zero_label}
32
33                push -1
34                mul
35                addi 32
36
37                return
38
39            {non_zero_label}:
40                log_2_floor
41                addi 1
42                return
43        }
44    }
45}
46
47#[cfg(test)]
48mod tests {
49    use super::*;
50    use crate::test_prelude::*;
51
52    impl Closure for LeadingZeros {
53        type Args = u32;
54
55        fn rust_shadow(&self, stack: &mut Vec<BFieldElement>) {
56            let arg = pop_encodable::<Self::Args>(stack);
57            push_encodable(stack, &arg.leading_zeros());
58        }
59
60        fn pseudorandom_args(
61            &self,
62            seed: [u8; 32],
63            bench_case: Option<BenchmarkCase>,
64        ) -> Self::Args {
65            match bench_case {
66                Some(BenchmarkCase::CommonCase) => 1 << 15,
67                Some(BenchmarkCase::WorstCase) => u32::MAX,
68                None => StdRng::from_seed(seed).random(),
69            }
70        }
71
72        fn corner_case_args(&self) -> Vec<Self::Args> {
73            vec![0, 1, 2, 3, 1 << 28, 1 << 29, 1 << 30, 1 << 31, u32::MAX]
74        }
75    }
76
77    #[test]
78    fn unit() {
79        ShadowedClosure::new(LeadingZeros).test();
80    }
81}
82
83#[cfg(test)]
84mod benches {
85    use super::*;
86    use crate::test_prelude::*;
87
88    #[test]
89    fn benchmark() {
90        ShadowedClosure::new(LeadingZeros).bench();
91    }
92}