How can I create an is_prime function that is generic over various integer types?

Question

I just took the dive into Rust and want to make some basic math functions that are generic. I have the following is_prime function:

fn is_prime(n: i64) -> bool {
    if n == 2 || n == 3 {
        return true;
    } else if n % 2 == 0 || n % 3 == 0 {
        return false;
    }

    let mut i = 5i64;
    let mut w = 2i64;
    while i*i <= n {
        if n % i == 0 {
            return false;
        }
        i += w;
        w = 6 - w;
    }
    true
}

What would it take for me to be able to pass isize, i64, usize, etc. as arguments? I have read through the Rust guide on the homepage but I am not certain how to apply the ideas of traits to my goal here.

Answer 1

Generic number types can be quite a nuisance to work with, but once you get the hang of them they don’t tend to be too bad, though a little more verbose. The standard building blocks to such methods are the traits in the num crate from crates.io, most notably Num, Zero and One, as well as the standard library's std::cmp::PartialOrd.

Numeric literals cannot be generic over any numeric type; they must be done with a trait method call; Zero::zero() and One::one() will suffice for most purposes—here the numbers that we want are 0, 1, 2, 3, 5 and 6, which are eminently achievable with these building blocks. You could also make your own trait with static methods producing these values and implement it for whatever numeric types you like, but doing it with just what’s guaranteed by Num is a better idea.

The basic procedure is to specify your generic type parameters as being based on Num (and PartialOrd if you write inequalities on values of that type, such as i * i <= n), and replace any numeric literals with ones constructed from zero and one, as the half dozen let statements at the start of the method below demonstrates. That will normally be enough.

Here’s what you end up with for this particular method:

// You’ll also need the appropriate dependencies.num addition to Cargo.toml
extern crate num;

use num::Num;

fn is_prime<N: Num + PartialOrd + Copy>(n: N) -> bool {
    let _0 = N::zero();
    let _1 = N::one();
    let _2 = _1 + _1;
    let _3 = _2 + _1;
    let _5 = _2 + _3;
    let _6 = _3 + _3;
    if n == _2 || n == _3 {
        return true;
    } else if n % _2 == _0 || n % _3 == _0 {
        return false;
    }

    let mut i = _5;
    let mut w = _2;
    while i * i <= n {
        if n % i == _0 {
            return false;
        }
        i = i + w;
        w = _6 - w;
    }
    true
}

Answer 2

To add to Chris Morgan's answer, you can use num::NumCast::from to cast to a generic number type where using Zero and One would be inappropriate. In your case:

use num::{Num, NumCast};

fn is_prime<N: Num + Ord + NumCast + Copy>(n: N) -> bool {
    let _0: N = NumCast::from(0usize).unwrap();
    let _1: N = NumCast::from(1usize).unwrap();
    let _2: N = NumCast::from(2usize).unwrap();
    let _3: N = NumCast::from(3usize).unwrap();
    let _4: N = NumCast::from(4usize).unwrap();
    let _5: N = NumCast::from(5usize).unwrap();
    let _6: N = NumCast::from(6usize).unwrap();

Answer 3

Another option using the num package is using the ToPrimitive trait to convert whatever type is passed in to an i64. Using the same function you gave, modified a little to allow us to return None when the conversion to i64 can't be done:

use num::ToPrimitive;

fn is_prime<N: ToPrimitive>(q: N) -> Option<bool> {
    let n: i64 = q.to_i64()?;
    if n == 2 || n == 3 {
        return Some(true);
    } else if n % 2 == 0 || n % 3 == 0 {
        return Some(false);
    }

    let mut i = 5i64;
    let mut w = 2i64;
    while i * i <= n {
        if n % i == 0 {
            return Some(false);
        }
        i += w;
        w = 6 - w;
    }
    Some(true)
}

fn print_is_prime<N: ToPrimitive + std::fmt::Debug + Clone>(q: N) {
    println!("Is {:?} prime? {:?}", q, is_prime(q.clone()));
}

fn main() {
    let prime: usize = 7;
    print_is_prime(prime);

    let prime: isize = 7;
    print_is_prime(prime);

    let prime: i32 = 7;
    print_is_prime(prime);

    let prime: i64 = 7;
    print_is_prime(prime);
}

This gives the output:

Is 7 prime? Some(true)
Is 7 prime? Some(true)
Is 7 prime? Some(true)
Is 7 prime? Some(true)