tasm_lib/arithmetic/u64/
popcount.rs

1use std::collections::HashMap;
2
3use triton_vm::prelude::*;
4
5use crate::prelude::*;
6use crate::traits::basic_snippet::Reviewer;
7use crate::traits::basic_snippet::SignOffFingerprint;
8
9/// Count [the number of ones](u64::count_ones) in the binary representation of
10/// the argument, also known as its population count.
11///
12/// ### Behavior
13///
14/// ```text
15/// BEFORE: _ [x: u64]
16/// AFTER:  _ [result: u32]
17/// ```
18///
19/// ### Preconditions
20///
21/// - the input argument is properly [`BFieldCodec`] encoded
22///
23/// ### Postconditions
24///
25/// - the output equals the number of 1s in the binary representation of the
26///   input argument
27/// - the output is properly [`BFieldCodec`] encoded
28#[derive(Debug, Copy, Clone, Eq, PartialEq, Hash)]
29pub struct PopCount;
30
31impl BasicSnippet for PopCount {
32    fn parameters(&self) -> Vec<(DataType, String)> {
33        vec![(DataType::U64, "x".to_string())]
34    }
35
36    fn return_values(&self) -> Vec<(DataType, String)> {
37        vec![(DataType::U32, "pop_count(x)".to_string())]
38    }
39
40    fn entrypoint(&self) -> String {
41        "tasmlib_arithmetic_u64_popcount".to_string()
42    }
43
44    fn code(&self, _: &mut Library) -> Vec<LabelledInstruction> {
45        triton_asm!(
46            {self.entrypoint()}:
47                pop_count
48                pick 1
49                pop_count
50                add
51                return
52        )
53    }
54
55    fn sign_offs(&self) -> HashMap<Reviewer, SignOffFingerprint> {
56        let mut sign_offs = HashMap::new();
57        sign_offs.insert(Reviewer("ferdinand"), 0xa23391ac8286ed10.into());
58        sign_offs
59    }
60}
61
62#[cfg(test)]
63mod tests {
64    use super::*;
65    use crate::test_prelude::*;
66
67    impl Closure for PopCount {
68        type Args = u64;
69
70        fn rust_shadow(&self, stack: &mut Vec<BFieldElement>) -> Result<(), RustShadowError> {
71            let x = pop_encodable::<Self::Args>(stack)?;
72            push_encodable(stack, &x.count_ones());
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 << 25,
83                Some(BenchmarkCase::WorstCase) => u64::MAX,
84                None => StdRng::from_seed(seed).random(),
85            }
86        }
87
88        fn corner_case_args(&self) -> Vec<Self::Args> {
89            [1, 1 << 32, u64::MAX]
90                .into_iter()
91                .flat_map(|x| [x.checked_sub(1), Some(x), x.checked_add(1)])
92                .flatten()
93                .collect()
94        }
95    }
96
97    #[macro_rules_attr::apply(test)]
98    fn rust_shadow_test() {
99        ShadowedClosure::new(PopCount).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(PopCount).bench();
111    }
112}
tasm_lib/arithmetic/u64/popcount.rs

tasm_lib/arithmetic/u64/
popcount.rs
