Merge pull request #19 from jorenham/prime

Author: jorenham

src/ffi.rs

@@ -66,4 +66,7 @@ unsafe extern "C" {
     pub fn math_legendre_p_prime(l: c_int, x: f64) -> f64;
     pub fn math_legendre_p_zeros(l: c_int, out: *mut f64);
     pub fn math_legendre_q(l: c_uint, x: f64) -> f64;
+
+    // boost/math/special_functions/prime.hpp
+    pub fn math_prime(n: c_uint) -> u32;
 }

src/math/mod.rs

@@ -10,3 +10,4 @@ pub use special_functions::factorials::*;
 pub use special_functions::gamma::*;
 pub use special_functions::jacobi::*;
 pub use special_functions::legendre::*;
+pub use special_functions::prime::*;

src/math/special_functions/mod.rs

@@ -6,3 +6,4 @@ pub mod factorials;
 pub mod gamma;
 pub mod jacobi;
 pub mod legendre;
+pub mod prime;

src/math/special_functions/prime.rs

@@ -0,0 +1,38 @@
+use crate::ffi;
+use core::ffi::c_uint;
+
+/// Corresponds to `static const unsigned max_prime` in `boost/math/special_functions/prime.hpp`
+const MAX_PRIME: u32 = 9_999;
+
+/// Fast table lookup to the first 10,000 prime numbers
+///
+/// The first prime number (`n=0`) is `2` (as `1` isn't terribly useful in practice).
+///
+/// The function will panic if `n >= 10_000`.
+///
+/// Corresponds to `boost::math::prime(n)` in C++.
+/// <https://boost.org/doc/libs/latest/libs/math/doc/html/math_toolkit/number_series/primes.html>
+pub fn prime(n: u32) -> u32 {
+    if n > MAX_PRIME {
+        panic!("Argument n out of range: got {n}");
+    }
+    unsafe { ffi::math_prime(n as c_uint) }
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    #[test]
+    fn test_prime() {
+        assert_eq!(prime(0), 2);
+        assert_eq!(prime(54), 257);
+        assert_eq!(prime(9_999), 104_729);
+    }
+
+    #[test]
+    #[should_panic(expected = "Argument n out of range: got 10000")]
+    fn test_prime_out_of_range() {
+        let _ = prime(10_000);
+    }
+}

wrapper.cpp

@@ -22,6 +22,7 @@
 #include <boost/math/special_functions/gamma.hpp>
 #include <boost/math/special_functions/jacobi.hpp>
 #include <boost/math/special_functions/legendre.hpp>
+#include <boost/math/special_functions/prime.hpp>
 
 using namespace boost::math;
 
@@ -101,4 +102,7 @@ void math_legendre_p_zeros(int l, double* out) {
 }
 double math_legendre_q(unsigned l, double x) { return legendre_q(l, x); }
 
+// boost/math/special_functions/prime.hpp
+std::uint32_t math_prime(unsigned n) { return prime(n); }
+
 } // extern "C"