From 3de46b16226b03aff4f36c45ee4ef7e2d96c3270 Mon Sep 17 00:00:00 2001
From: Arthur Gautier <arthur.gautier@arista.com>
Date: Fri, 14 Feb 2025 16:11:28 -0800
Subject: [PATCH] rand_core: introduce an UnwrapMut wrapper

This is used when a method takes in a `&mut impl TryRngCore` but the inner
method requires an `&mut impl RngCore`.

An `UnwrapErr` can not be used as it requires ownership of the rng.
---
 rand_core/CHANGELOG.md |  5 +++++
 rand_core/src/lib.rs   | 29 +++++++++++++++++++++++++++++
 2 files changed, 34 insertions(+)

diff --git a/rand_core/CHANGELOG.md b/rand_core/CHANGELOG.md
index f13de92d1a..66cd46f49c 100644
--- a/rand_core/CHANGELOG.md
+++ b/rand_core/CHANGELOG.md
@@ -4,6 +4,11 @@ All notable changes to this project will be documented in this file.
 The format is based on [Keep a Changelog](http://keepachangelog.com/en/1.0.0/)
 and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0.html).
 
+## Unreleased
+
+### API changes
+- Add `TryRngCore::unwrap_ref` to only take a mutable reference of the rng (#1589)
+
 ## [0.9.0] - 2025-01-27
 ### Dependencies and features
 - Bump the MSRV to 1.63.0 (#1207, #1246, #1269, #1341, #1416, #1536); note that 1.60.0 may work for dependents when using `--ignore-rust-version`
diff --git a/rand_core/src/lib.rs b/rand_core/src/lib.rs
index 874a086504..0336c5bc2a 100644
--- a/rand_core/src/lib.rs
+++ b/rand_core/src/lib.rs
@@ -236,6 +236,11 @@ pub trait TryRngCore {
         UnwrapErr(self)
     }
 
+    /// Wrap RNG with the [`UnwrapMut`] wrapper.
+    fn unwrap_mut(&mut self) -> UnwrapMut<'_, Self> {
+        UnwrapMut(self)
+    }
+
     /// Convert an [`RngCore`] to a [`RngReadAdapter`].
     #[cfg(feature = "std")]
     fn read_adapter(&mut self) -> RngReadAdapter<'_, Self>
@@ -311,6 +316,30 @@ impl<R: TryRngCore> RngCore for UnwrapErr<R> {
 
 impl<R: TryCryptoRng> CryptoRng for UnwrapErr<R> {}
 
+/// Wrapper around [`TryRngCore`] implementation which implements [`RngCore`]
+/// by panicking on potential errors.
+#[derive(Debug, Eq, PartialEq, Hash)]
+pub struct UnwrapMut<'r, R: TryRngCore + ?Sized>(pub &'r mut R);
+
+impl<R: TryRngCore> RngCore for UnwrapMut<'_, R> {
+    #[inline]
+    fn next_u32(&mut self) -> u32 {
+        self.0.try_next_u32().unwrap()
+    }
+
+    #[inline]
+    fn next_u64(&mut self) -> u64 {
+        self.0.try_next_u64().unwrap()
+    }
+
+    #[inline]
+    fn fill_bytes(&mut self, dst: &mut [u8]) {
+        self.0.try_fill_bytes(dst).unwrap()
+    }
+}
+
+impl<R: TryCryptoRng> CryptoRng for UnwrapMut<'_, R> {}
+
 /// A random number generator that can be explicitly seeded.
 ///
 /// This trait encapsulates the low-level functionality common to all