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.

Kapoor

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
}

Wrongheaded

Instead of using the Num trait as a constraint, it's also possible to use the basic traits that are actually required: N: PartialEq + PartialOrd + Add<N, N> + Sub<N, N> + Mul<N, N> + Rem<N, N> + One + Zero. Num is simply a convenient shortcut. – Jambeau

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();

Estriol

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)

Kif

Answer 4

So you can also use macro's for this, but truth be told it's better to write your code with a specific type first and then turn it in to a macro by replacing all relevant type occurrences with the macro type.

(Looks like that is what the num crate does btw.)

If you stick to top level functions you'll have to use the paste macro crate to compose your name.

use paste::paste;

macro_rules! is_prime_for_x {
    ($T:ty) => {
        paste! {
            named_is_prime_for_x!([<is_prime_ $T>],$T);
        }
    };
}
// 2 stage macro allows composing the name, with a macro, but doesn't brake formatting
macro_rules! named_is_prime_for_x {
    ($name:ident,$T:ty) => {
        fn $name(n: $T) -> bool {
            ...
        }
    };
}

is_prime_for_x!(u32) // creates fn is_prime_u32(n:u32)
is_prime_for_x!(u64) // creates fn is_prime_u64(n:u64)
is_prime_for_x!(u128) // creates fn is_prime_u128(n:u128)

Using associated functions of the generic impl you can even go without the past create.

struct Primes<T> {
    primes: Vec<T>, // could store known primes here ;)
}

macro_rules! impl_primes {
    ($T:ty) => {
        impl Primes<$T> {
            pub fn is_prime(n: $T) -> bool {
                ...
            }
        }
    };
}

impl_primes!(u32);
impl_primes!(u64);
impl_primes!(u128);

fn main() {
    Primes::<u64>::is_prime(123);
}

In the binary it will be multiple implementations, but seem the main goal is dont-repeat-yourself (DRY), right?

Aishaaisle