From e960a5cf5dc027d1636cfa6d667417ea0dd6ff12 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Johanna=20S=C3=B6rng=C3=A5rd?=
 <44257381+JSorngard@users.noreply.github.com>
Date: Thu, 12 Oct 2023 13:30:47 +0200
Subject: [PATCH] =?UTF-8?q?Add=20M=C3=B6bius=20function=20(#13)?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

* Add the möbius function

* Add panics section to möbius docstring

* Make möbius take a NonZeroU64

* Move use statement to only place it's used

* Simplify isqrt test

* square --> squared in möbius docstring

* Move möbius into lib.rs

* Speed up möbius function

* improve some comments
---
 benches/prime_benches.rs | 14 +++++-
 src/imath.rs             |  5 ++-
 src/lib.rs               | 95 ++++++++++++++++++++++++++++++++++++++++
 3 files changed, 111 insertions(+), 3 deletions(-)

diff --git a/benches/prime_benches.rs b/benches/prime_benches.rs
index 2bb1f6d..0d2ab4f 100644
--- a/benches/prime_benches.rs
+++ b/benches/prime_benches.rs
@@ -1,4 +1,4 @@
-use const_primes::{are_prime, is_prime, primes};
+use const_primes::{are_prime, is_prime, moebius, primes};
 use criterion::{criterion_group, criterion_main, BatchSize, Criterion};
 use rand::prelude::*;
 use std::hint::black_box;
@@ -19,9 +19,21 @@ fn benchmarks(c: &mut Criterion) {
             BatchSize::SmallInput,
         )
     });
+
     c.bench_function("sieve 10000 integers", |b| {
         b.iter(|| black_box(are_prime::<10_000>()))
     });
+
+    let ints: Vec<_> = (1..1_000_000)
+        .map(|n| core::num::NonZeroU64::new(n).unwrap())
+        .collect();
+    c.bench_function("möbius of first 1e6 integers", |b| {
+        b.iter(|| {
+            for &i in &ints {
+                black_box(moebius(i));
+            }
+        })
+    });
 }
 
 criterion_group!(benches, benchmarks);
diff --git a/src/imath.rs b/src/imath.rs
index 1eb7048..0c7f8af 100644
--- a/src/imath.rs
+++ b/src/imath.rs
@@ -56,13 +56,14 @@ mod test {
     #[test]
     fn check_isqrt() {
         #[rustfmt::skip]
-        const TEST_CASES: [(u64, u64); 101] = [(0, 0), (1, 1), (2, 1), (3, 1), (4, 2), (5, 2), (6, 2), (7, 2), (8, 2), (9, 3), (10, 3), (11, 3), (12, 3), (13, 3), (14, 3), (15, 3), (16, 4), (17, 4), (18, 4), (19, 4), (20, 4), (21, 4), (22, 4), (23, 4), (24, 4), (25, 5), (26, 5), (27, 5), (28, 5), (29, 5), (30, 5), (31, 5), (32, 5), (33, 5), (34, 5), (35, 5), (36, 6), (37, 6), (38, 6), (39, 6), (40, 6), (41, 6), (42, 6), (43, 6), (44, 6), (45, 6), (46, 6), (47, 6), (48, 6), (49, 7), (50, 7), (51, 7), (52, 7), (53, 7), (54, 7), (55, 7), (56, 7), (57, 7), (58, 7), (59, 7), (60, 7), (61, 7), (62, 7), (63, 7), (64, 8), (65, 8), (66, 8), (67, 8), (68, 8), (69, 8), (70, 8), (71, 8), (72, 8), (73, 8), (74, 8), (75, 8), (76, 8), (77, 8), (78, 8), (79, 8), (80, 8), (81, 9), (82, 9), (83, 9), (84, 9), (85, 9), (86, 9), (87, 9), (88, 9), (89, 9), (90, 9), (91, 9), (92, 9), (93, 9), (94, 9), (95, 9), (96, 9), (97, 9), (98, 9), (99, 9), (u64::MAX, 4294967296)];
-        for (x, ans) in TEST_CASES {
+        const TEST_CASES: [u64; 100] = [0, 1, 1, 1, 2, 2, 2, 2, 2, 3, 3, 3, 3, 3, 3, 3, 4, 4, 4, 4, 4, 4, 4, 4, 4, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 6, 6, 6, 6, 6, 6, 6, 6, 6, 6, 6, 6, 6, 7, 7, 7, 7, 7, 7, 7, 7, 7, 7, 7, 7, 7, 7, 7, 8, 8, 8, 8, 8, 8, 8, 8, 8, 8, 8, 8, 8, 8, 8, 8, 8, 9, 9, 9, 9, 9, 9, 9, 9, 9, 9, 9, 9, 9, 9, 9, 9, 9, 9, 9];
+        for (x, ans) in TEST_CASES.into_iter().enumerate() {
             assert_eq!(isqrt(x as u64), ans);
         }
         assert_eq!(
             f64::from(u32::MAX).sqrt().floor() as u64,
             isqrt(u64::from(u32::MAX))
         );
+        assert_eq!(isqrt(u64::MAX), 4294967296);
     }
 }
diff --git a/src/lib.rs b/src/lib.rs
index 79197f7..96216d4 100644
--- a/src/lib.rs
+++ b/src/lib.rs
@@ -334,6 +334,86 @@ pub const fn are_prime<const N: usize>() -> [bool; N] {
     sieve
 }
 
+/// Returns the value of the [Möbius function](https://en.wikipedia.org/wiki/M%C3%B6bius_function).
+///
+/// This function is
+/// - 1 if `n` is a square-free integer with an even number of prime factors,  
+/// - -1 if `n` is a square-free integer with an odd number of prime factors,  
+/// - 0 if `n` has a squared prime factor.
+///
+/// Uses a small wheel to check prime factors up to `√n`` and exits early if
+/// there is a squared factor.
+///
+/// # Example
+/// ```
+/// # use const_primes::moebius;
+/// use core::num::NonZeroU64;
+/// const N: NonZeroU64 = match NonZeroU64::new(1001) {Some(i) => i, None => panic!()};
+/// const MÖBIUS1001: i8 = moebius(N);
+/// assert_eq!(MÖBIUS1001, -1);
+/// ```
+pub const fn moebius(x: core::num::NonZeroU64) -> i8 {
+    let mut x = x.get();
+
+    let mut prime_count = 0;
+
+    // Handle 2 separately
+    if x % 2 == 0 {
+        x /= 2;
+        prime_count += 1;
+        // If 2^2 is also a divisor of x
+        if x % 2 == 0 {
+            return 0;
+        }
+    }
+    // Same for 3
+    if x % 3 == 0 {
+        x /= 3;
+        prime_count += 1;
+        if x % 3 == 0 {
+            return 0;
+        }
+    }
+
+    // For the remaining factors
+    let mut i = 5;
+    let bound = isqrt(x);
+    while i <= bound {
+        // If i is a divisor of x
+        if x % i == 0 {
+            x /= i;
+            prime_count += 1;
+
+            // check if i^2 is also a divisor of x
+            if x % i == 0 {
+                return 0;
+            }
+        }
+        // Same for i + 2
+        if x % (i + 2) == 0 {
+            x /= i + 2;
+            prime_count += 1;
+
+            if x % (i + 2) == 0 {
+                return 0;
+            }
+        }
+        i += 6;
+    }
+
+    // If x is a prime it will never be divided by any factor less than its square root.
+    // In that case we can check if x is still larger than one, and then count it.
+    if x > 1 {
+        prime_count += 1;
+    }
+
+    if prime_count % 2 == 0 {
+        1
+    } else {
+        -1
+    }
+}
+
 /// Returns the largest prime smaller than or equal to `n` if there is one.
 ///
 /// Scans for primes downwards from the input with [`is_prime`].
@@ -468,6 +548,21 @@ mod test {
         );
     }
 
+    #[test]
+    fn check_möbius() {
+        use core::num::NonZeroU64;
+        #[rustfmt::skip]
+        const TEST_CASES: [i8; 50] = [1, -1, -1, 0, -1, 1, -1, 0, 0, 1, -1, 0, -1, 1, 1, 0, -1, 0, -1, 0, 1, 1, -1, 0, 0, 1, 0, 0, -1, -1, -1, 0, 1, 1, 1, 0, -1, 1, 1, 0, -1, -1, -1, 0, 0, 1, -1, 0, 0, 0];
+        for (n, ans) in TEST_CASES.into_iter().enumerate() {
+            assert_eq!(moebius(NonZeroU64::new(n as u64 + 1).unwrap()), ans);
+        }
+        #[rustfmt::skip]
+        const BIG_TEST_CASES: [i8; 51] = [0, -1, -1, 1, 0, -1, 1, 1, 0, -1, -1, 1, 0, -1, 0, -1, 0, 0, 1, -1, 0, -1, -1, -1, 0, 0, 0, 1, 0, 0, -1, -1, 0, -1, -1, 0, 0, 1, -1, -1, 0, 1, 1, 1, 0, -1, 1, 1, 0, -1, 0];
+        for (n, ans) in BIG_TEST_CASES.into_iter().enumerate() {
+            assert_eq!(moebius(NonZeroU64::new(n as u64 + 1000).unwrap()), ans);
+        }
+    }
+
     #[test]
     fn check_are_prime_below() {
         macro_rules! test_n_below_100 {