Add a convenience next(range) to Rng.

2 files changed, 28 insertions, 0 deletions

diff --git a/embassy-stm32/src/rng.rs b/embassy-stm32/src/rng.rs
index 77e408d3d..e94c813de 100644
--- a/embassy-stm32/src/rng.rs
+++ b/embassy-stm32/src/rng.rs
@@ -126,6 +126,26 @@ impl<T: Instance> traits::rng::Rng for Random<T> {
             Ok(())
         }
     }
+
+    #[rustfmt::skip]
+    type NextFuture<'a> where Self: 'a = impl Future<Output=Result<u32, Self::Error>> + 'a;
+
+    fn next<'a>(&'a mut self, range: u32) -> Self::NextFuture<'a> {
+        async move {
+            let t = (-(range as i32) % (range as i32)) as u32;
+            loop {
+                let mut buf = [0; 4];
+                traits::rng::Rng::fill_bytes(self, &mut buf).await?;
+                let x = u32::from_le_bytes(buf);
+                let m = x as u64 * range as u64;
+                let l = m as u32;
+                if l < t {
+                    continue;
+                }
+                return Ok((m >> 32) as u32);
+            }
+        }
+    }
 }
 
 pub(crate) mod sealed {
diff --git a/embassy-traits/src/rng.rs b/embassy-traits/src/rng.rs
index ddc4c20e0..320d9afc2 100644
--- a/embassy-traits/src/rng.rs
+++ b/embassy-traits/src/rng.rs
@@ -4,6 +4,7 @@ use core::future::Future;
 pub trait Rng {
     type Error;
 
+    #[rustfmt::skip]
     type RngFuture<'a>: Future<Output = Result<(), Self::Error>> + 'a
     where
         Self: 'a;
@@ -14,4 +15,11 @@ pub trait Rng {
     /// filling the buffer. Upon completion, the buffer will be completely
     /// filled or an error will have been reported.
     fn fill_bytes<'a>(&'a mut self, dest: &'a mut [u8]) -> Self::RngFuture<'a>;
+
+    #[rustfmt::skip]
+    type NextFuture<'a>: Future<Output = Result<u32, Self::Error>> + 'a
+    where
+        Self: 'a;
+
+    fn next<'a>(&'a mut self, range: u32) -> Self::NextFuture<'a>;
 }