Cargo.toml

[package]
name = "necklace_matching"
version = "0.1.0"
edition = "2018"

[dependencies]
criterion = "0.3"
rand = "0.7"
num-bigint = "0.2"
num-traits = "0.2"

[[bench]]
name = "main"
harness = false

lib.rs

// src/lib.rs

use num_bigint::BigUint;
use num_traits::{Pow, Zero};

/// represents a fixed range of primes
pub struct Primes {
    /// ordered vector of primes
    numbers: Vec<usize>,

    /// sieve of prime numbers
    is_prime: Vec<bool>,

    /// maximum number tested
    n: usize,
}

impl Primes {
    /// Computes primes within range to n (inclusive).
    pub fn sieve_erato(n: usize) -> Self {
        let mut is_prime = vec![true; n];
        // set 0, 1 to false
        is_prime.iter_mut().take(2).for_each(|x| *x = false);

        for i in 0..(n as f64).sqrt() as usize + 1 {
            if is_prime[i] {
                is_prime[i * i..n]
                    .iter_mut()
                    .step_by(i)
                    .for_each(|is_p| *is_p = false)
            }
        }

        let numbers = is_prime
            .iter()
            .enumerate()
            .filter_map(|(p, &is_p)| if is_p { Some(p) } else { None })
            .collect();

        Primes {
            numbers,
            n,
            is_prime,
        }
    }

    /// Iterator of primes relativistic to `n`.
    pub fn relative(&self, n: usize) -> impl Iterator<Item = &usize> + Clone {
        debug_assert!(n < self.n);
        // minor optimization for known primes; reduces average
        // runtime by ~%10 on primes within range of `0..10_000`
        let start = if self.is_prime[n] {
            self.numbers.binary_search(&n).unwrap()
        } else {
            0
        };
        self.numbers[start..]
            .iter()
            .take_while(move |&&p| p <= n)
            .filter(move |&&p| n % p == 0)
    }

    /// Euler's totient function.
    pub fn phi(&self, n: usize) -> usize {
        let it = self.relative(n);
        let p: usize = it.clone().product();
        let p1: usize = it.map(|p| p - 1).product();
        n * p1 / p
    }

    /// Return count of `k`-ary necklace of length `n` as `u128`.
    pub fn necklaces(&self, k: usize, n: usize) -> u128 {
        let k = k as u128;
        let range = 1..(n as f64).sqrt() as usize + 1;
        let nums = range.filter(|x| n % x == 0).map(|x| {
            let (a, b) = (x, n / x);
            let div_a = self.phi(a) as u128 * k.pow(b as u32);
            let div_b = self.phi(b) as u128 * k.pow(a as u32);
            div_a + if a != b { div_b } else { 0 }
        });
        nums.sum::<u128>() / n as u128
    }

    /// Return count of `k`-ary necklace of length `n` as `BigUint`.
    pub fn necklaces_big(&self, k: usize, n: usize) -> BigUint {
        let k: BigUint = k.into();
        let range = 1..(n as f64).sqrt() as usize + 1;
        let nums = range.filter(|x| n % x == 0).map(|x| {
            let (a, b) = (x, n / x);
            let div_a = self.phi(a) * k.pow(b as u32);
            let div_b = self.phi(b) * k.pow(a as u32);
            div_a + if a != b { div_b } else { Zero::zero() }
        });
        nums.sum::<BigUint>() / n
    }
}

main.rs

// benches/main.rs
use criterion::*;
use rand::seq::SliceRandom;
use rand::Rng;
use std::time::Duration;

use necklace_matching::*;

pub fn bench_necklaces(c: &mut Criterion) {
    let mut group = c.benchmark_group("necklaces");
    group
        .warm_up_time(Duration::new(1, 0))
        .sample_size(15)
        .measurement_time(Duration::new(10, 0));

    group.bench_function("sieve_erato-100", |b| b.iter(|| Primes::sieve_erato(100)));

    group.bench_function("sieve_erato-1000", |b| b.iter(|| Primes::sieve_erato(1000)));

    let primes = Primes::sieve_erato(1001);
    let mut rng = rand::thread_rng();
    let mut n: Vec<usize> = (1..100_000).map(|x| x % 99 + 1).collect();
    n.shuffle(&mut rng);
    let mut n = n.into_iter().cycle();

    group.bench_function("relative", |b| {
        b.iter(|| primes.relative(n.next().unwrap()).sum::<usize>())
    });

    group.bench_function("phi", |b| b.iter(|| primes.phi(n.next().unwrap())));

    group.bench_function("necklaces-n-100", |b| {
        b.iter(|| primes.necklaces(1, n.next().unwrap()))
    });

    group.bench_function("necklaces-n-1000", |b| {
        b.iter(|| primes.necklaces(1, n.next().unwrap() * 10))
    });

    let mut k: Vec<usize> = (1..100_000).map(|x| x % 99 + 1).collect();
    k.shuffle(&mut rng);
    let mut k = k.into_iter().cycle();

    group.bench_function("necklaces-k-100", |b| {
        b.iter(|| primes.necklaces(k.next().unwrap(), 1))
    });

    group.bench_function("necklaces-k-1000", |b| {
        b.iter(|| primes.necklaces(k.next().unwrap() * 10, 1))
    });

    group.bench_function("necklaces-k-n", |b| {
        b.iter(|| primes.necklaces(k.next().unwrap(), n.next().unwrap()))
    });

    group.bench_function("necklaces-big-k-n", |b| {
        b.iter(|| primes.necklaces_big(k.next().unwrap(), n.next().unwrap()))
    });
    
    group.bench_function("necklaces-big-3-90", |b| { 
        b.iter(|| primes.necklaces_big(3, 90))
    });

    group.bench_function("necklaces-123-18", |b| b.iter(|| primes.necklaces(123, 18)));
    
    group
        .warm_up_time(Duration::new(10, 0))
        .sample_size(15)
        .measurement_time(Duration::new(100, 0));

    group.bench_function("necklaces-big-1024-512", |b| {
        b.iter(|| primes.necklaces_big(1024, 512))
    });
}

criterion_group!(benches, bench_necklaces);
criterion_main!(benches);

tests.rs

// src/tests.rs
use super::*;

#[test]
fn necklaces_small_test() {
    let prime = Primes::sieve_erato(1000);

    assert_eq!(prime.necklaces(3, 10), 5934);
    assert_eq!(prime.necklaces(2, 12), 352);
    assert_eq!(prime.necklaces(3, 7), 315);
    assert_eq!(prime.necklaces(9, 4), 1665);
    assert_eq!(prime.necklaces(21, 3), 3101);
    assert_eq!(prime.necklaces(99, 2), 4950);
}

#[test]
fn necklaces_u128_test() {
    let prime = Primes::sieve_erato(19);
    assert_eq!(
        prime.necklaces(12345678910, 3),
        627225458787209496560873442940_u128
    );
    assert_eq!(
        prime.necklaces(1234567, 6),
        590115108867910855092196771880677924_u128
    );
    assert_eq!(
        prime.necklaces(123, 18),
        2306850769218800390268044415272597042_u128,
    );
}

#[test]
fn necklaces_large_test() {
    let prime = Primes::sieve_erato(91);
    assert_eq!(
        prime.necklaces_big(3, 90),
        BigUint::parse_bytes(b"96977372978752360287715019917722911297222", 10).unwrap()
    );
    assert_eq!(
        prime.necklaces_big(123, 18),
        2306850769218800390268044415272597042_u128.into()
    );
    assert_eq!(
        prime.necklaces_big(1234567, 6),
        590115108867910855092196771880677924_u128.into()
    );
    assert_eq!(
        prime.necklaces_big(12345678910, 3),
        627225458787209496560873442940_u128.into()
    );
}

#[test]
fn phi_test() {
    let prime = Primes::sieve_erato(98);
    assert_eq!(prime.phi(97), 96);

    // hard coded test 1
    let phis = vec![1, 1, 2, 2, 4, 2, 6, 4, 6, 4, 10, 4, 12, 6, 8, 8];
    for (i, &phi) in phis.iter().enumerate() {
        assert_eq!(prime.phi(i + 1), phi);
    }

    // primes test
    for &number in prime.numbers.iter() {
        assert_eq!(prime.phi(number), number - 1);
    }

    // hard coded test 2
    assert_eq!(prime.phi(20), 8);
    assert_eq!(prime.phi(36), 12);
    assert_eq!(prime.phi(81), 54);
    assert_eq!(prime.phi(90), 24);
}

#[test]
fn primes_numbers() {
    let p: Vec<_> = Primes::sieve_erato(100).numbers;
    assert_eq!(&p[..5], &[2, 3, 5, 7, 11]);
}