Skip to main content

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 parameters(&self) -> Vec<(DataType, String)> {
11        vec![(DataType::U64, "arg".to_string())]
12    }
13
14    fn return_values(&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>) -> Result<(), RustShadowError> {
71            let arg = pop_encodable::<Self::Args>(stack)?;
72            push_encodable(stack, &arg.leading_zeros());
73            Ok(())
74        }
75
76        fn pseudorandom_args(
77            &self,
78            seed: [u8; 32],
79            bench_case: Option<BenchmarkCase>,
80        ) -> Self::Args {
81            match bench_case {
82                Some(BenchmarkCase::CommonCase) => 1 << 31,
83                Some(BenchmarkCase::WorstCase) => 1 << 62,
84                None => StdRng::from_seed(seed).random(),
85            }
86        }
87
88        fn corner_case_args(&self) -> Vec<Self::Args> {
89            let small = 0..10;
90            let medium = (27..35).map(|i| 1 << i);
91            let large = (0..10).map(|i| u64::MAX - i);
92
93            small.chain(medium).chain(large).collect()
94        }
95    }
96
97    #[macro_rules_attr::apply(test)]
98    fn unit() {
99        ShadowedClosure::new(LeadingZeros).test();
100    }
101}
102
103#[cfg(test)]
104mod benches {
105    use super::*;
106    use crate::test_prelude::*;
107
108    #[macro_rules_attr::apply(test)]
109    fn benchmark() {
110        ShadowedClosure::new(LeadingZeros).bench();
111    }
112}