tasm_lib/arithmetic/u64/
leading_zeros.rs

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
use triton_vm::prelude::*;

use crate::arithmetic::u32::leading_zeros::LeadingZeros as U32LeadingZeroes;
use crate::prelude::*;

#[derive(Debug, Copy, Clone, Eq, PartialEq, Hash)]
pub struct LeadingZeros;

impl BasicSnippet for LeadingZeros {
    fn inputs(&self) -> Vec<(DataType, String)> {
        vec![(DataType::U64, "arg".to_string())]
    }

    fn outputs(&self) -> Vec<(DataType, String)> {
        vec![(DataType::U32, "leading_zeros(arg)".to_string())]
    }

    fn entrypoint(&self) -> String {
        "tasmlib_arithmetic_u64_leading_zeros".to_string()
    }

    fn code(&self, library: &mut Library) -> Vec<LabelledInstruction> {
        let leading_zeros_u32 = library.import(Box::new(U32LeadingZeroes));

        let entrypoint = self.entrypoint();
        let hi_is_zero_label = format!("{entrypoint}_hi_is_zero");

        triton_asm!(
            // BEFORE: _ value_hi value_lo
            // AFTER:  _ (leading_zeros as u32)
            {entrypoint}:
                pick 1
                call {leading_zeros_u32}
                // _ value_lo leading_zeros_value_hi

                dup 0
                push 32
                eq
                skiz
                    call {hi_is_zero_label}

                // _ temp leading_zeros

                pick 1
                pop 1
                return

            {hi_is_zero_label}:
                // _ value_lo 32

                pick 1
                call {leading_zeros_u32}
                // _ 32 leading_zeros_value_lo

                addi 32
                // _ 32 leading_zeros
                return
        )
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::test_prelude::*;

    impl Closure for LeadingZeros {
        type Args = u64;

        fn rust_shadow(&self, stack: &mut Vec<BFieldElement>) {
            let arg = pop_encodable::<Self::Args>(stack);
            push_encodable(stack, &arg.leading_zeros());
        }

        fn pseudorandom_args(
            &self,
            seed: [u8; 32],
            bench_case: Option<BenchmarkCase>,
        ) -> Self::Args {
            match bench_case {
                Some(BenchmarkCase::CommonCase) => 1 << 31,
                Some(BenchmarkCase::WorstCase) => 1 << 62,
                None => StdRng::from_seed(seed).random(),
            }
        }

        fn corner_case_args(&self) -> Vec<Self::Args> {
            let small = 0..10;
            let medium = (27..35).map(|i| 1 << i);
            let large = (0..10).map(|i| u64::MAX - i);

            small.chain(medium).chain(large).collect()
        }
    }

    #[test]
    fn unit() {
        ShadowedClosure::new(LeadingZeros).test();
    }
}

#[cfg(test)]
mod benches {
    use super::*;
    use crate::test_prelude::*;

    #[test]
    fn benchmark() {
        ShadowedClosure::new(LeadingZeros).bench();
    }
}