[PATCH] random: get_random_u64_below()

[Kent Overstreet <kent.overstreet@linux.dev>]

bcachefs needs this, for sampling devices to read from based on squared
device latencies.

this uses the same algorithm as get_random_u32_below: since the multiply
uses the top and bottom halves separately, it works out fairly well.

Cc: "Theodore Ts'o" <tytso@mit.edu> (maintainer:RANDOM NUMBER DRIVER)
Cc: "Jason A. Donenfeld" <Jason@zx2c4.com> (maintainer:RANDOM NUMBER DRIVER)
Signed-off-by: Kent Overstreet <kent.overstreet@linux.dev>
---
 drivers/char/random.c  | 22 ++++++++++++++++++++++
 include/linux/random.h | 22 ++++++++++++++++++++++
 2 files changed, 44 insertions(+)

diff --git a/drivers/char/random.c b/drivers/char/random.c
index 2581186fa61b..84808300044c 100644
--- a/drivers/char/random.c
+++ b/drivers/char/random.c
@@ -588,6 +588,28 @@ u32 __get_random_u32_below(u32 ceil)
 }
 EXPORT_SYMBOL(__get_random_u32_below);
 
+u64 __get_random_u64_below(u64 ceil)
+{
+	if (unlikely(!ceil))
+		return get_random_u64();
+	if (ceil <= U32_MAX)
+		return __get_random_u32_below(ceil);
+
+	u64 rand = get_random_u64();
+	u64 mult = ceil * rand;
+
+	if (unlikely(mult < ceil)) {
+		u64 bound = -ceil % ceil;
+		while (unlikely(mult < bound)) {
+			rand = get_random_u64();
+			mult = ceil * rand;
+		}
+	}
+
+	return mul_u64_u64_shr(ceil, rand, 64);
+}
+EXPORT_SYMBOL(__get_random_u64_below);
+
 #ifdef CONFIG_SMP
 /*
  * This function is called when the CPU is coming up, with entry
diff --git a/include/linux/random.h b/include/linux/random.h
index 333cecfca93f..b025bf3d8f27 100644
--- a/include/linux/random.h
+++ b/include/linux/random.h
@@ -6,6 +6,7 @@
 #include <linux/bug.h>
 #include <linux/kernel.h>
 #include <linux/list.h>
+#include <linux/math64.h>
 
 #include <uapi/linux/random.h>
 
@@ -90,6 +91,27 @@ static inline u32 get_random_u32_below(u32 ceil)
 	}
 }
 
+u64 __get_random_u64_below(u64 ceil);
+
+static inline u64 get_random_u64_below(u32 ceil)
+{
+	if (!__builtin_constant_p(ceil))
+		return __get_random_u64_below(ceil);
+
+	BUILD_BUG_ON_MSG(!ceil, "get_random_u64_below() must take ceil > 0");
+	if (ceil <= 1)
+		return 0;
+	if (ceil <= U32_MAX)
+		return get_random_u32_below(ceil);
+
+	for (;;) {
+		u64 rand = get_random_u64();
+		u64 mult = ceil * rand;
+		if (likely(mult >= -ceil % ceil))
+			return mul_u64_u64_shr(ceil, rand, 64);
+	}
+}
+
 /*
  * Returns a random integer in the interval (floor, U32_MAX], with uniform
  * distribution, suitable for all uses. Fastest when floor is a constant, but
-- 
2.47.2