Search code examples
genericsrusttraits

Creating a generic function bounded to integer types


In my attempt to learn rust I am starting with some basic exercises. I have written a simple function in what I hope is idiomatic rust to count the number of set bits in an integer.

fn bit_count(x: u32) -> u32 {
    (0..32).into_iter().map(|i| (x >> i) & 1).sum()
}

fn main() {
    println!("{} has {} set bits.", 5, bit_count(5));
}

Now I want to make the function generic so that I can pass any integer type: i32, u32, i64, u64... etc.

I am quite familiar with tmp in c++ but my attempt with rust generics have failed, so far I have this:

extern crate num;

fn bit_count<T>(x: T) -> T
where
    T: num::Integer + std::ops::BitAnd + std::ops::Shr + num::NumCast,
    std::ops::Range<T>: std::iter::IntoIterator,
{
    (T::zero()..num::NumCast::from(32).unwrap())
        .into_iter()
        .map(|i| (x >> num::NumCast::from(i)) & T::one())
        .sum()
}

fn main() {
    println!("{} has {} set bits.", 5, bit_count(5));
}

I saw the num crate advertised and it seemed like a good fit. I was expecting to have T: num::Integer and be done, however I feel like I'm spiraling down a rabbit hole here and I can't seem to get the right combination of bounds.

Any suggestions would be great! and any tips to make my code more idiomatic would also be helpful, thanks.


Solution

  • Got there in the end. Turns out I needed to use the num::PrimInt trait as my bound because it includes all of the bitwise operations and casts. The num::Integer is less constrained and models an integer in the pure mathematical sense, so no bitwise operations.

    The final code I have looks like this:

    extern crate num;
    
    fn bit_count<T>(x: T) -> T
    where
        T: num::PrimInt + std::iter::Sum,
    {    
        let n_bits = std::mem::size_of::<T>() * u8::BITS as usize;
        (0..n_bits).into_iter().map(|i| (x >> i) & T::one()).sum()
    
    }
    
    fn main() {
        println!("{} has {} set bits.", 5, bit_count(5u32));
        println!("{} has {} set bits.", 5, bit_count(5i32));
        println!("{} has {} set bits.", 5, bit_count(5i64));
    }
    

    It would be nice to not need that T::one() but it seems there's no way around it. Additionally the std::iter::Sum trait was needed in my bounds to allow the functional workflow.

    The num crate actually has a function to count the number of set bits too num::PrimInt::count_ones.