How is count_ones implemented in Rust?

user_12482789 0
Viewed 805

I tried looking for how Rust implements count_ones(). I'm curious because it seems to vastly outperform my own naive implementation (no kidding), and I would really like to see why it is so performant. My guess is that Rust is using some asm to do the work. For completeness, here was my attempt:

/*
* my attempt to implement count_ones for i32 types
* but this is much slower than the default
* implementation.
*/
fn count_ones(num: i32) -> u32 {
    let mut ans: u32 = 0;
    let mut _num = num;

    while _num > 0 {
        if _num & 0x1 == 0x1 {
            ans += 1;
        }
        _num >>= 1;
    }

    ans
}

I found this on the rust repo, but I can't make sense of it (still new to Rust!) (reproduced below).

#[inline]
fn count_ones(self) -> u32 {
    unsafe { $ctpop(self as $ActualT) as u32 }
}
rust
1 Answers

Let's follow the code step-by-step.

First, looking at the snippet you've posted - it contains several macro variables (identifiers with a dollar sign prepended), so it is assumed that this code is, in fact, a part of macro definition. Scrolling up, we get the following:

macro_rules! uint_impl {
    ($T:ty = $ActualT:ty, $BITS:expr,
     $ctpop:path,
     $ctlz:path,
     $cttz:path,
     $bswap:path,
     $add_with_overflow:path,
     $sub_with_overflow:path,
     $mul_with_overflow:path) => {
        #[stable(feature = "rust1", since = "1.0.0")]
        #[allow(deprecated)]
        impl Int for $T {
             // skipped
        }
    }
}

Now, to see that are the variable values here, we should find where this macro is invoked. In general, this might be hard, due to the macro scoping rules, but here we'll just search the same file, and here it is:

uint_impl! { u8 = u8, 8,
    intrinsics::ctpop8,
    intrinsics::ctlz8,
    intrinsics::cttz8,
    bswap8,
    intrinsics::u8_add_with_overflow,
    intrinsics::u8_sub_with_overflow,
    intrinsics::u8_mul_with_overflow }

(and multiple another invocations). Comparing this with the macro definition, we see that the function we're looking for will be expanded to the following:

#[inline]
fn count_ones(self) -> u32 {
    unsafe { intrinsics::ctpop8(self as u8) as u32 }
}

And, finally, intrinsics::ctpop8 is, as Stargateur mentioned in comment, an LLVM intrinsic, i.e. this call is directly converted into LLVM instruction.

However, there's a little better way to find out what is what.

Let's now look for the function we're interested in in the std documentation. Searching for count_ones brings together a bunch of functions, for each primitive number type independently; we'll take a look on the implementation for u8. Clicking the src link on the function brings us to the code:

doc_comment! {
    concat!("Returns the number of ones in the binary representation of `self`.

# Examples

Basic usage:

```
", $Feature, "let n = 0b01001100", stringify!($SelfT), ";

assert_eq!(n.count_ones(), 3);", $EndFeature, "
```"),
    #[stable(feature = "rust1", since = "1.0.0")]
    #[rustc_const_stable(feature = "const_math", since = "1.32.0")]
    #[inline]
    pub const fn count_ones(self) -> u32 {
        intrinsics::ctpop(self as $ActualT) as u32
    }
}

...which just directly calls the intrinsics::ctpop function we've found before.

Now you might wonder, why these two searches yielded different pieces of code. Reason is simple: the commit you're referring to is from the fairly old version of rustc - pre-1.0, if I understand correctly; at that time, numerical operations were implemented as part of Num trait, not directly on primitive types. If you check out the implementation for version 1.44.1, which is the current one at the time of writing, you'll see the same code I've quoted above from the docs.

user_30034010
Related