big: Add `internal_random_prime`.

6 files changed, 180 insertions, 17 deletions

diff --git a/core/math/big/common.odin b/core/math/big/common.odin
index 243e74d0c..fb215cd69 100644
--- a/core/math/big/common.odin
+++ b/core/math/big/common.odin
@@ -89,6 +89,17 @@ MATH_BIG_USE_FROBENIUS_TEST       :: !MATH_BIG_USE_LUCAS_SELFRIDGE_TEST;
 USE_MILLER_RABIN_ONLY       := false;
 
 /*
+	How many times we'll call `internal_int_random` during random prime generation before we bail out.
+	Set to 0 or less to try indefinitely.	
+*/
+MAX_ITERATIONS_RANDOM_PRIME  := 1_000_000;
+
+/*
+	How many iterations we used for the last random prime.
+*/
+@thread_local RANDOM_PRIME_ITERATIONS_USED: int;
+
+/*
 	We don't allow these to be switched at runtime for two reasons:
 
 	1) 32-bit and 64-bit versions of procedures use different types for their storage,
@@ -175,9 +186,9 @@ Error_String :: #partial [Error]string{
 };
 
 Primality_Flag :: enum u8 {
-	Blum_Blum_Shub = 0,	/* BBS style prime */
-	Safe           = 1,	/* Safe prime (p-1)/2 == prime */
-	Second_MSB_On  = 3,  /* force 2nd MSB to 1 */
+ 	Blum_Blum_Shub = 0,  // Make prime congruent to 3 mod 4
+	Safe           = 1,  // Make sure (p-1)/2 is prime as well (implies .Blum_Blum_Shub)
+	Second_MSB_On  = 3,  // Make the 2nd highest bit one
 };
 Primality_Flags :: bit_set[Primality_Flag; u8];
 
diff --git a/core/math/big/example.odin b/core/math/big/example.odin
index 252cff0bf..45cfae778 100644
--- a/core/math/big/example.odin
+++ b/core/math/big/example.odin
@@ -37,6 +37,7 @@ Runtime tunable:
 	FACTORIAL_BINARY_SPLIT_CUTOFF         %v
 	FACTORIAL_BINARY_SPLIT_MAX_RECURSIONS %v
 	USE_MILLER_RABIN_ONLY                 %v
+	MAX_ITERATIONS_RANDOM_PRIME           %v
 
 `, _DIGIT_BITS,
 _LOW_MEMORY,
@@ -58,6 +59,7 @@ FACTORIAL_MAX_N,
 FACTORIAL_BINARY_SPLIT_CUTOFF,
 FACTORIAL_BINARY_SPLIT_MAX_RECURSIONS,
 USE_MILLER_RABIN_ONLY,
+MAX_ITERATIONS_RANDOM_PRIME,
 );
 
 }
@@ -84,18 +86,27 @@ print :: proc(name: string, a: ^Int, base := i8(10), print_name := true, newline
 	}
 }
 
-// printf :: fmt.printf;
+printf :: fmt.printf;
 
 demo :: proc() {
 	a, b, c, d, e, f, res := &Int{}, &Int{}, &Int{}, &Int{}, &Int{}, &Int{}, &Int{};
 	defer destroy(a, b, c, d, e, f, res);
 
-	set(a, _private_prime_table[_PRIME_TAB_SIZE - 1]);
-	print("a: ", a);
-	trials := number_of_rabin_miller_trials(internal_count_bits(a));
-	err := internal_int_prime_next_prime(a, trials, false);
-	print("a->next: ", a);
-	fmt.printf("Trials: %v, Error: %v\n", trials, err);
+	bits   := 111;
+	trials := -1;
+
+	flags := Primality_Flags{};
+	fmt.printf("Trying to generate a %v bit prime using %v Miller-Rabin trials and options %v.\n", bits, trials, flags);
+
+	err: Error;
+	{
+		SCOPED_TIMING(.random_prime);
+		err = internal_random_prime(a, bits, trials, flags);
+	}
+
+	print("a(10): ", a, 10, true, true, true);
+	fmt.printf("err: %v\n", err);
+	fmt.printf("RANDOM_PRIME_ITERATIONS_USED: %v\n", RANDOM_PRIME_ITERATIONS_USED);
 }
 
 main :: proc() {
diff --git a/core/math/big/helpers.odin b/core/math/big/helpers.odin
index ff654172c..ccd00615f 100644
--- a/core/math/big/helpers.odin
+++ b/core/math/big/helpers.odin
@@ -362,15 +362,15 @@ int_random_digit :: proc(r: ^rnd.Rand = nil) -> (res: DIGIT) {
 	return 0; // We shouldn't get here.
 }
 
-int_rand :: proc(dest: ^Int, bits: int, r: ^rnd.Rand = nil, allocator := context.allocator) -> (err: Error) {
+int_random :: proc(dest: ^Int, bits: int, r: ^rnd.Rand = nil, allocator := context.allocator) -> (err: Error) {
 	/*
 		Check that `a` is usable.
 	*/
 	assert_if_nil(dest);
-	return #force_inline internal_int_rand(dest, bits, r, allocator);
+	return #force_inline internal_int_random(dest, bits, r, allocator);
 
 }
-rand :: proc { int_rand, };
+random :: proc { int_random, };
 
 /*
 	Internal helpers.
diff --git a/core/math/big/internal.odin b/core/math/big/internal.odin
index 5b7e84a87..a0516be3d 100644
--- a/core/math/big/internal.odin
+++ b/core/math/big/internal.odin
@@ -2045,6 +2045,7 @@ internal_invmod :: proc{ internal_int_inverse_modulo, };
 
 /*
 	Helpers to extract values from the `Int`.
+	Offset is zero indexed.
 */
 internal_int_bitfield_extract_bool :: proc(a: ^Int, offset: int) -> (val: bool, err: Error) {
 	limb := offset / _DIGIT_BITS;
@@ -2116,6 +2117,34 @@ internal_int_bitfield_extract :: proc(a: ^Int, offset, count: int) -> (res: _WOR
 }
 
 /*
+	Helpers to (un)set a bit in an Int.
+	Offset is zero indexed.
+*/
+internal_int_bitfield_set_single :: proc(a: ^Int, offset: int) -> (err: Error) {
+	limb := offset / _DIGIT_BITS;
+	if limb < 0 || limb >= a.used  { return .Invalid_Argument; }
+	i := DIGIT(1 << uint((offset % _DIGIT_BITS)));
+	a.digit[limb] |= i;
+	return;
+}
+
+internal_int_bitfield_unset_single :: proc(a: ^Int, offset: int) -> (err: Error) {
+	limb := offset / _DIGIT_BITS;
+	if limb < 0 || limb >= a.used  { return .Invalid_Argument; }
+	i := DIGIT(1 << uint((offset % _DIGIT_BITS)));
+	a.digit[limb] &= _MASK - i;
+	return;
+}
+
+internal_int_bitfield_toggle_single :: proc(a: ^Int, offset: int) -> (err: Error) {
+	limb := offset / _DIGIT_BITS;
+	if limb < 0 || limb >= a.used  { return .Invalid_Argument; }
+	i := DIGIT(1 << uint((offset % _DIGIT_BITS)));
+	a.digit[limb] ~= i;
+	return;
+}
+
+/*
 	Resize backing store.
 	We don't need to pass the allocator, because the storage itself stores it.
 
@@ -2817,7 +2846,7 @@ internal_int_random_digit :: proc(r: ^rnd.Rand = nil) -> (res: DIGIT) {
 	return 0; // We shouldn't get here.
 }
 
-internal_int_rand :: proc(dest: ^Int, bits: int, r: ^rnd.Rand = nil, allocator := context.allocator) -> (err: Error) {
+internal_int_random :: proc(dest: ^Int, bits: int, r: ^rnd.Rand = nil, allocator := context.allocator) -> (err: Error) {
 	context.allocator = allocator;
 
 	bits := bits;
@@ -2842,7 +2871,7 @@ internal_int_rand :: proc(dest: ^Int, bits: int, r: ^rnd.Rand = nil, allocator :
 	dest.used = digits;
 	return nil;
 }
-internal_rand :: proc { internal_int_rand, };
+internal_random :: proc { internal_int_random, };
 
 /*
 	Internal helpers.
diff --git a/core/math/big/prime.odin b/core/math/big/prime.odin
index f2d360e3f..ae538281c 100644
--- a/core/math/big/prime.odin
+++ b/core/math/big/prime.odin
@@ -456,7 +456,7 @@ internal_int_is_prime :: proc(a: ^Int, miller_rabin_trials := int(-1), miller_ra
 		for ix := 0; ix < miller_rabin_trials; ix += 1 {
 
 			// rand() guarantees the first digit to be non-zero
-			internal_rand(b, _DIGIT_TYPE_BITS, r) or_return;
+			internal_random(b, _DIGIT_TYPE_BITS, r) or_return;
 
 			// Reduce digit before casting because DIGIT might be bigger than
 			// an unsigned int and "mask" on the other side is most probably not.
@@ -474,7 +474,7 @@ internal_int_is_prime :: proc(a: ^Int, miller_rabin_trials := int(-1), miller_ra
 				ix -= 1;
 				continue;
 			}
-			internal_rand(b, l) or_return;
+			internal_random(b, l) or_return;
 
 			// That number might got too big and the witness has to be smaller than "a"
 			l = internal_count_bits(b);
@@ -1168,6 +1168,117 @@ internal_int_prime_next_prime :: proc(a: ^Int, trials: int, bbs_style: bool, all
 	return;
 }
 
+/*
+	Makes a truly random prime of a given size (bits),
+
+	Flags are as follows:
+	 	Blum_Blum_Shub    - Make prime congruent to 3 mod 4
+		Safe              - Make sure (p-1)/2 is prime as well (implies .Blum_Blum_Shub)
+		Second_MSB_On     - Make the 2nd highest bit one
+
+	This is possibly the mother of all prime generation functions, muahahahahaha!
+*/
+internal_random_prime :: proc(a: ^Int, size_in_bits: int, trials: int, flags := Primality_Flags{}, r: ^rnd.Rand = nil, allocator := context.allocator) -> (err: Error) {
+	context.allocator = allocator;
+	flags  := flags;
+	trials := trials;
+
+	t := &Int{};
+	defer internal_destroy(t);
+
+	/*
+		Sanity check the input.
+	*/
+	if size_in_bits <= 1 || trials < -1                              { return .Invalid_Argument; }
+
+	/*
+		`.Safe` implies `.Blum_Blum_Shub`.
+	*/
+	if .Safe in flags {
+		if size_in_bits < 3 {
+			/*
+				The smallest safe prime is 5, which takes 3 bits.
+				We early out now, else we'd be locked in an infinite loop trying to generate a 2-bit Safe Prime.
+			*/
+			return .Invalid_Argument;
+		}
+		flags += { .Blum_Blum_Shub, };
+	}
+
+	/*
+		Automatically choose the number of Rabin-Miller trials?
+	*/
+	if trials == -1 {
+		trials = number_of_rabin_miller_trials(size_in_bits);
+	}
+
+	res: bool;
+	RANDOM_PRIME_ITERATIONS_USED = 0;
+
+	for {
+		if MAX_ITERATIONS_RANDOM_PRIME > 0 {
+			RANDOM_PRIME_ITERATIONS_USED += 1;
+			if RANDOM_PRIME_ITERATIONS_USED > MAX_ITERATIONS_RANDOM_PRIME {
+				return .Max_Iterations_Reached;
+			}
+		}
+
+		internal_int_random(a, size_in_bits)                         or_return;
+
+		/*
+			Make sure it's odd.
+		*/
+		if size_in_bits > 2 {
+			a.digit[0] |= 1;
+		} else {
+			/*
+				A 2-bit prime can be either 2 (0b10) or 3 (0b11).
+				So, let's force the top bit to 1 and return early.
+			*/
+			a.digit[0] |= 2;
+			return nil;
+		}
+
+		if .Blum_Blum_Shub in flags {
+			a.digit[0] |= 3;
+		}
+		if .Second_MSB_On in flags {
+			internal_int_bitfield_set_single(a, size_in_bits - 2) or_return;
+		}
+
+		/*
+			Is it prime?
+		*/
+		res = internal_int_is_prime(a, trials)                       or_return;
+
+		if (!res) {
+			continue;
+		}
+
+		if .Safe in flags {
+			/*
+				See if (a-1)/2 is prime.
+			*/
+			internal_sub(a, a, 1)                                    or_return;
+			internal_int_shr1(a, a)                                  or_return;
+
+			/*
+				Is it prime?
+			*/
+			res = internal_int_is_prime(a, trials)                   or_return;
+		}
+		if res { break; }
+	}
+
+	if .Safe in flags {
+		/*
+			Restore a to the original value.
+		*/
+		internal_int_shl1(a, a)                                      or_return;
+		internal_add(a, a, 1)                                        or_return;
+	}
+	return;
+}
 
 /*
 	Returns the number of Rabin-Miller trials needed for a given bit size.
diff --git a/core/math/big/tune.odin b/core/math/big/tune.odin
index ced8e5f5d..6ce5121b4 100644
--- a/core/math/big/tune.odin
+++ b/core/math/big/tune.odin
@@ -24,6 +24,7 @@ Category :: enum {
 	bitfield_extract,
 	rm_trials,
 	is_prime,
+	random_prime,
 };
 
 Event :: struct {