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 inputs(&self) -> Vec<(DataType, String)> {
33        vec![(DataType::U64, "x".to_string())]
34    }
35
36    fn outputs(&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"), 0xeaae261c6df2bcbb.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>) {
71            let x = pop_encodable::<Self::Args>(stack);
72            push_encodable(stack, &x.count_ones());
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 << 25,
82                Some(BenchmarkCase::WorstCase) => u64::MAX,
83                None => StdRng::from_seed(seed).random(),
84            }
85        }
86
87        fn corner_case_args(&self) -> Vec<Self::Args> {
88            [1, 1 << 32, u64::MAX]
89                .into_iter()
90                .flat_map(|x| [x.checked_sub(1), Some(x), x.checked_add(1)])
91                .flatten()
92                .collect()
93        }
94    }
95
96    #[test]
97    fn rust_shadow_test() {
98        ShadowedClosure::new(PopCount).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(PopCount).bench();
110    }
111}
tasm_lib/arithmetic/u64/popcount.rs

tasm_lib/arithmetic/u64/
popcount.rs
