From d5c4f85391efe1a9293dc628b8aeffc9e153368e Mon Sep 17 00:00:00 2001
From: Brian Smith <brian@briansmith.org>
Date: Wed, 19 Oct 2022 09:54:59 -0700
Subject: [PATCH] Add `getrandom_arrays`.

Implement `getrandom_array. It requires Rust 1.51 and the user must
enable the "array" feature explicitly.
---
 Cargo.toml  |  4 ++-
 src/lib.rs  | 88 +++++++++++++++++++++++++++++++++++++++++++++++++++++
 src/util.rs | 14 +++++++++
 3 files changed, 105 insertions(+), 1 deletion(-)

diff --git a/Cargo.toml b/Cargo.toml
index 2024c8f9..44d40521 100644
--- a/Cargo.toml
+++ b/Cargo.toml
@@ -30,6 +30,8 @@ js-sys = { version = "0.3", optional = true }
 wasm-bindgen-test = "0.3.18"
 
 [features]
+# Feature to enable `getrandom_array`, which requires Rust 1.51 or later.
+array = []
 # Implement std-only traits for getrandom::Error
 std = []
 # Feature to enable fallback RDRAND-based implementation on x86/x86_64
@@ -49,5 +51,5 @@ rustc-dep-of-std = [
 test-in-browser = []
 
 [package.metadata.docs.rs]
-features = ["std", "custom"]
+features = ["array", "custom", "std"]
 rustdoc-args = ["--cfg", "docsrs"]
diff --git a/src/lib.rs b/src/lib.rs
index c4809d6b..00200fe6 100644
--- a/src/lib.rs
+++ b/src/lib.rs
@@ -327,3 +327,91 @@ pub fn getrandom_uninit(dest: &mut [MaybeUninit<u8>]) -> Result<&mut [u8], Error
     // since it returned `Ok`.
     Ok(unsafe { slice_assume_init_mut(dest) })
 }
+
+#[cfg(feature = "array")]
+pub use crate::util::from_bytes::ArrayItem;
+
+/// Returns a value of type `T` where all bytes of the value were
+/// initialized using `getrandom_uninit`. Feature: `rust-1-55`.
+///
+/// This works on byte arrays, without any type annotations:
+/// ```
+/// # use getrandom::getrandom_array;
+/// fn tls_hello_random() -> Result<[u8; 32], getrandom::Error> {
+///     getrandom_array()
+/// }
+/// ```
+///
+/// This also works on nested byte arrays. This nested array support can be
+/// used to safely and efficiently (in optimized builds) construct random
+/// arrays of other types, though this requires type annotations to help the
+/// compiler figure out the nested array size:
+/// ```
+/// # use getrandom::{getrandom_array, Error};
+/// # fn u32_array_example() -> Result<(), Error> {
+/// let random_u32s: [u32; 4] =
+///     getrandom_array()?.map(u32::from_ne_bytes);
+/// # Ok(())
+/// # }
+/// ```
+///
+/// Multiple levels of array nesting can be used to construct more complicated
+/// types:
+/// ```
+/// # // TODO: Use `std::simd::Simd` when the `portable_simd` feature is
+/// # // available; until then, here is a minimal polyfill for it that
+/// # // allows the examle to work in stable Rust.
+/// # struct Simd<T, const N: usize>([T; N]);
+/// # impl<T, const N: usize> From<[T; N]> for Simd<T, N> {
+/// #     fn from(value: [T; N]) -> Self {
+/// #        Self(value)
+/// #     }
+/// # }
+/// # use getrandom::{getrandom_array, Error};
+/// # fn simd_array_example() -> Result<(), Error> {
+/// let random_vectors: [Simd<u32, 4>; 16] =
+///     getrandom_array()?
+///         .map(|bytes: [_; 4]| bytes.map(u32::from_ne_bytes))
+///         .map(Simd::from);
+/// # Ok(())
+/// # }
+/// ```
+///
+/// The number of nested arrays is limited only by whatever limits the compiler
+/// imposes:
+/// ```
+/// # use getrandom::{getrandom_array, Error};
+/// # fn more_nesting_example() -> Result<(), Error> {
+/// let more_nesting: [[[[[[u8; 1]; 2]; 3]; 4]; 5]; 6] = getrandom_array()?;
+/// # Ok(())
+/// }
+/// ```
+///
+/// The patterns above allows us to avoid implementing
+/// `ArrayItem` for an endless number of types.
+#[cfg(feature = "array")]
+#[inline(always)]
+pub fn getrandom_array<I: ArrayItem, const N: usize>() -> Result<[I; N], Error> {
+    // This is `inline(always)` because the code generated by the compiler is
+    // terrible when it isn't inlined.
+    //
+    // This function only requires Rust 1.51 but `[T]::map()` used in the
+    // doctests is stable only in Rust 1.55.
+
+    let mut uninit: MaybeUninit<[I; N]> = MaybeUninit::uninit();
+    // TODO: `uninit.as_bytes_mut()` when that is available.
+    {
+        // SAFETY: MaybeUninit<u8> is always valid, even for padding bytes.
+        // The compiler will ensure that `B` isn't too large.
+        let as_bytes_mut = unsafe {
+            core::slice::from_raw_parts_mut(
+                uninit.as_mut_ptr() as *mut MaybeUninit<u8>,
+                core::mem::size_of::<[I; N]>(),
+            )
+        };
+        getrandom_uninit(as_bytes_mut)?;
+    }
+    // SAFETY: `dest` has been fully initialized by `getrandom_uninit`
+    // since it returned `Ok`.
+    Ok(unsafe { uninit.assume_init() })
+}
diff --git a/src/util.rs b/src/util.rs
index 3162afad..f36c0c25 100644
--- a/src/util.rs
+++ b/src/util.rs
@@ -99,3 +99,17 @@ pub unsafe fn slice_as_uninit_mut<T>(slice: &mut [T]) -> &mut [MaybeUninit<T>] {
     // SAFETY: `MaybeUninit<T>` is guaranteed to be layout-compatible with `T`.
     &mut *(slice as *mut [T] as *mut [MaybeUninit<T>])
 }
+
+#[cfg(feature = "array")]
+pub(crate) mod from_bytes {
+    mod sealed {
+        pub trait Sealed {}
+
+        impl Sealed for u8 {}
+        impl<B: Sealed, const N: usize> Sealed for [B; N] {}
+    }
+
+    /// A type supported by [`getrandom_array`](crate::getrandom_array).
+    pub unsafe trait ArrayItem: sealed::Sealed {}
+    unsafe impl<B: sealed::Sealed> ArrayItem for B {}
+}