CODEWITHSUNDEEP
performancerustinteger-arithmetic
Ben Eills
Ben Eills

Reputation: 13

Writing an efficient power function in Rust

I am writing an efficient squaring method in Rust. Let's assume that the Mul trait of AbstractNumber is a black box and that we're only allowed safe, idiomatic Rust.

Below is a first pass which uses repeated squaring for larger indices. I'm unsure how LLVM will translate Rust arithmetic method calls such as checked_next_power_of_two().

Does the following look reasonable? Would it be more efficient to split off the smaller-case branch into its own inlined function?

/// Compute an integer power of this number efficiently with repeated squaring.
pub fn pow(&self, n: u32) -> AbstractNumber {
    let optimization = 5;

    if n < optimization {
        let mut x = Complex::one();

        for _ in 0..n {
            x *= *self;
        }

        x
    } else {
        // l = floor(log_2(n)), r = n - 2^l
        let (l, r) = if n.is_power_of_two() {
            (n.trailing_zeros(), 0)
        } else {
            let p = n.checked_next_power_of_two().unwrap().trailing_zeros() - 1;
            (p, n - 2u32.pow(p))
        };

        let mut x = *self;

        for _ in 0..l {
            x *= x;
        }

        self.pow(r) * x
    }
}

Upvotes: 1

Views: 1003

Answers (1)

paholg
paholg

Reputation: 2080

Why not use num::pow::pow? In any case, here is how it is implemented:

#[inline]
pub fn pow<T: Clone + One + Mul<T, Output = T>>(mut base: T, mut exp: usize) -> T {
    if exp == 0 { return T::one() }

    while exp & 1 == 0 {
        base = base.clone() * base;
        exp >>= 1;
    }
    if exp == 1 { return base }

    let mut acc = base.clone();
    while exp > 1 {
        exp >>= 1;
        base = base.clone() * base;
        if exp & 1 == 1 {
            acc = acc * base.clone();
        }
    }
    acc
}

It requires Clone in addition to Mul (and One, but that's not needed if you're not being generic).

There's nothing wrong or unsafe about using bitwise operations in Rust, by the way.

Upvotes: 4

Related Questions