tasm_lib/arithmetic/u64/
leading_zeros.rs

1use triton_vm::prelude::*;
2
3use crate::arithmetic::u32::leading_zeros::LeadingZeros as U32LeadingZeroes;
4use crate::prelude::*;
5
6#[derive(Debug, Copy, Clone, Eq, PartialEq, Hash)]
7pub struct LeadingZeros;
8
9impl BasicSnippet for LeadingZeros {
10    fn inputs(&self) -> Vec<(DataType, String)> {
11        vec![(DataType::U64, "arg".to_string())]
12    }
13
14    fn outputs(&self) -> Vec<(DataType, String)> {
15        vec![(DataType::U32, "leading_zeros(arg)".to_string())]
16    }
17
18    fn entrypoint(&self) -> String {
19        "tasmlib_arithmetic_u64_leading_zeros".to_string()
20    }
21
22    fn code(&self, library: &mut Library) -> Vec<LabelledInstruction> {
23        let leading_zeros_u32 = library.import(Box::new(U32LeadingZeroes));
24
25        let entrypoint = self.entrypoint();
26        let hi_is_zero_label = format!("{entrypoint}_hi_is_zero");
27
28        triton_asm!(
29            // BEFORE: _ value_hi value_lo
30            // AFTER:  _ (leading_zeros as u32)
31            {entrypoint}:
32                pick 1
33                call {leading_zeros_u32}
34                // _ value_lo leading_zeros_value_hi
35
36                dup 0
37                push 32
38                eq
39                skiz
40                    call {hi_is_zero_label}
41
42                // _ temp leading_zeros
43
44                pick 1
45                pop 1
46                return
47
48            {hi_is_zero_label}:
49                // _ value_lo 32
50
51                pick 1
52                call {leading_zeros_u32}
53                // _ 32 leading_zeros_value_lo
54
55                addi 32
56                // _ 32 leading_zeros
57                return
58        )
59    }
60}
61
62#[cfg(test)]
63mod tests {
64    use super::*;
65    use crate::test_prelude::*;
66
67    impl Closure for LeadingZeros {
68        type Args = u64;
69
70        fn rust_shadow(&self, stack: &mut Vec<BFieldElement>) {
71            let arg = pop_encodable::<Self::Args>(stack);
72            push_encodable(stack, &arg.leading_zeros());
73        }
74
75        fn pseudorandom_args(
76            &self,
77            seed: [u8; 32],
78            bench_case: Option<BenchmarkCase>,
79        ) -> Self::Args {
80            match bench_case {
81                Some(BenchmarkCase::CommonCase) => 1 << 31,
82                Some(BenchmarkCase::WorstCase) => 1 << 62,
83                None => StdRng::from_seed(seed).random(),
84            }
85        }
86
87        fn corner_case_args(&self) -> Vec<Self::Args> {
88            let small = 0..10;
89            let medium = (27..35).map(|i| 1 << i);
90            let large = (0..10).map(|i| u64::MAX - i);
91
92            small.chain(medium).chain(large).collect()
93        }
94    }
95
96    #[test]
97    fn unit() {
98        ShadowedClosure::new(LeadingZeros).test();
99    }
100}
101
102#[cfg(test)]
103mod benches {
104    use super::*;
105    use crate::test_prelude::*;
106
107    #[test]
108    fn benchmark() {
109        ShadowedClosure::new(LeadingZeros).bench();
110    }
111}