tasm_lib/arithmetic/u64/
trailing_zeros.rs

1use std::collections::HashMap;
2
3use triton_vm::prelude::*;
4
5use crate::arithmetic::u32::trailing_zeros::TrailingZeros as U32TrailingZeros;
6use crate::prelude::*;
7use crate::traits::basic_snippet::Reviewer;
8use crate::traits::basic_snippet::SignOffFingerprint;
9
10/// Returns the number of trailing zeros in the binary representation of the
11/// input argument. Behaves like [`u64::trailing_zeros`].
12///
13/// ### Behavior
14///
15/// ```text
16/// BEFORE: _ [arg: u64]
17/// AFTER:  _ [u64::trailing_zeros(arg): u32]
18/// ```
19///
20/// ### Preconditions
21///
22/// - the input argument is properly [`BFieldCodec`] encoded
23///
24/// ### Postconditions
25///
26/// - the output is the number of trailing zeros in the binary representation
27///   of `arg`
28#[derive(Debug, Copy, Clone, Eq, PartialEq, Hash)]
29pub struct TrailingZeros;
30
31impl BasicSnippet for TrailingZeros {
32    fn inputs(&self) -> Vec<(DataType, String)> {
33        vec![(DataType::U64, "arg".to_string())]
34    }
35
36    fn outputs(&self) -> Vec<(DataType, String)> {
37        vec![(DataType::U32, "trailing_zeros(arg)".to_string())]
38    }
39
40    fn entrypoint(&self) -> String {
41        "tasmlib_arithmetic_u64_trailing_zeros".to_string()
42    }
43
44    fn code(&self, library: &mut Library) -> Vec<LabelledInstruction> {
45        let u32_trailing_zeros = library.import(Box::new(U32TrailingZeros));
46
47        let entrypoint = self.entrypoint();
48        let hi_bits_dominate = format!("{entrypoint}_hi_bits_dominate");
49        let lo_bits_dominate = format!("{entrypoint}_lo_bits_dominate");
50
51        triton_asm! {
52            // BEFORE: _ arg_hi arg_lo
53            // AFTER:  _ trailing_zeros(arg)
54            {entrypoint}:
55                push 1
56                dup 1
57                push 0
58                eq
59                // _ arg_hi arg_lo 1 (arg_lo == 0)
60
61                skiz call {hi_bits_dominate}
62                skiz call {lo_bits_dominate}
63                // _ trailing_zeros(arg)
64
65                return
66
67            // BEFORE: _ arg_hi arg_lo 1
68            // AFTER:  _ trailing_zeros(arg) 0
69            // where arg_lo == 0
70            {hi_bits_dominate}:
71                pop 2
72                call {u32_trailing_zeros}
73                addi 32
74                push 0
75                return
76
77            // BEFORE: _ arg_hi arg_lo
78            // AFTER:  _ trailing_zeros(arg)
79            // where arg_lo != 0
80            {lo_bits_dominate}:
81                pick 1
82                pop 1
83                call {u32_trailing_zeros}
84                return
85        }
86    }
87
88    fn sign_offs(&self) -> HashMap<Reviewer, SignOffFingerprint> {
89        let mut sign_offs = HashMap::new();
90        sign_offs.insert(Reviewer("ferdinand"), 0x6819affeab8f551c.into());
91        sign_offs
92    }
93}
94
95#[cfg(test)]
96mod tests {
97    use super::*;
98    use crate::test_prelude::*;
99
100    impl Closure for TrailingZeros {
101        type Args = u64;
102
103        fn rust_shadow(&self, stack: &mut Vec<BFieldElement>) {
104            let arg = pop_encodable::<Self::Args>(stack);
105            push_encodable(stack, &arg.trailing_zeros());
106        }
107
108        fn pseudorandom_args(
109            &self,
110            seed: [u8; 32],
111            bench_case: Option<BenchmarkCase>,
112        ) -> Self::Args {
113            match bench_case {
114                Some(BenchmarkCase::CommonCase) => 0b1111_1111 << 15,
115                Some(BenchmarkCase::WorstCase) => 1 << 63,
116                None => StdRng::from_seed(seed).random(),
117            }
118        }
119
120        fn corner_case_args(&self) -> Vec<Self::Args> {
121            [1, 1 << 31, 1 << 32, 1 << 33, 1 << 63, u64::MAX - 1]
122                .into_iter()
123                .flat_map(|i| [i - 1, i, i + 1])
124                .collect()
125        }
126    }
127
128    #[test]
129    fn unit() {
130        ShadowedClosure::new(TrailingZeros).test();
131    }
132}
133
134#[cfg(test)]
135mod benches {
136    use super::*;
137    use crate::test_prelude::*;
138
139    #[test]
140    fn benchmark() {
141        ShadowedClosure::new(TrailingZeros).bench()
142    }
143}