From 2c7904d2b3457ce6d1572ab6a06993c5fcd985a1 Mon Sep 17 00:00:00 2001
From: Imbus <>
Date: Wed, 12 Feb 2025 20:37:45 +0100
Subject: [PATCH] Miller-rabin primality test for orders of magnitude lower
 computational complexity

---
 main.c | 14 ++++++++++++++
 rsa.c  | 34 +++++++++++++++++++++++++++++++++-
 rsa.h  | 12 +++++++++++-
 3 files changed, 58 insertions(+), 2 deletions(-)

diff --git a/main.c b/main.c
index f9d48d3..ccfb0fe 100644
--- a/main.c
+++ b/main.c
@@ -41,6 +41,20 @@ int main() {
 
     uint64_t n = p * q;
 
+    int before = SysTick->CNT;
+    bool check_a = miller_rabin(p, 10);
+    int miller = SysTick->CNT - before;
+
+    before = SysTick->CNT;
+    bool check_b = is_prime(p);
+    int isprime = SysTick->CNT - before;
+
+    printf("Is prime: %s %s\n", check_a ? "true" : "false",
+           check_b ? "true" : "false");
+
+    printf("Miller took %d ticks\n", miller);
+    printf("Is_prime took %d ticks\n", isprime);
+
     // Make these work by patching printf
     printf("P: %u\n", (uint32_t)p);
     printf("Q: %u\n", (uint32_t)q);
diff --git a/rsa.c b/rsa.c
index 07e06bd..b9c623a 100644
--- a/rsa.c
+++ b/rsa.c
@@ -54,7 +54,7 @@ uint64_t modexp(uint64_t a, uint64_t b, uint64_t m) {
 
 uint64_t gen_prime(uint64_t min, uint64_t max) {
     uint64_t cand = 0;
-    while (!is_prime(cand)) cand = prand_range(min, max);
+    while (!miller_rabin(cand, 10)) cand = prand_range(min, max);
 
     return cand;
 }
@@ -71,6 +71,38 @@ bool is_prime(int n) {
     return true;
 }
 
+bool miller_rabin(int n, int k) {
+    if (n < 2)
+        return false;
+
+    int d = n - 1;
+    int s = 0;
+
+    while (d % 2 == 0) {
+        d /= 2;
+        s++;
+    }
+
+    for (int i = 0; i < k; i++) {
+        int a = prand_range(2, n - 2);
+        int x = modexp(a, d, n);
+
+        if (x == 1 || x == n - 1)
+            continue;
+
+        for (int r = 1; r < s; r++) {
+            x = modexp(x, 2, n);
+            if (x == n - 1)
+                break;
+        }
+
+        if (x != n - 1)
+            return false; // Not prime
+    }
+
+    return true; // Likely prime
+}
+
 int mod_inverse(int e, int phi) {
     for (int d = 0; d < phi; d++) {
         if ((d * e) % phi == 1)
diff --git a/rsa.h b/rsa.h
index b84c43d..9111722 100644
--- a/rsa.h
+++ b/rsa.h
@@ -1,6 +1,6 @@
 #pragma once
-#include <stdint.h>
 #include <stdbool.h>
+#include <stdint.h>
 
 /**
  * @brief Calculates greatest common divider of two integers using the euclidean
@@ -55,3 +55,13 @@ uint64_t gen_prime(uint64_t min, uint64_t max);
  * @return true if n is prime, false otherwise.
  */
 bool is_prime(int n);
+
+/**
+ * @brief Performs the Miller-Rabin primality test to check if a number is
+ * probably prime.
+ *
+ * @param n The number to test for primality.
+ * @param k The number of rounds of testing to perform.
+ * @return true if n is probably prime, false if n is composite.
+ */
+bool miller_rabin(int n, int k);