Optimize Random{NumberGenerator}.GetItems for power-of-two choices

src/libraries/System.Private.CoreLib/src/System/Random.cs

@@ -1,8 +1,11 @@
 // Licensed to the .NET Foundation under one or more agreements.
 // The .NET Foundation licenses this file to you under the MIT license.
 
+using System.Buffers;
 using System.Diagnostics;
+using System.Numerics;
 using System.Runtime.CompilerServices;
+using System.Runtime.InteropServices;
 
 namespace System
 {
@@ -194,6 +197,44 @@ public void GetItems<T>(ReadOnlySpan<T> choices, Span<T> destination)
                 throw new ArgumentException(SR.Arg_EmptySpan, nameof(choices));
             }
 
+            // The most expensive part of this operation is the call to get random data. We can
+            // do so only once rather than per element in a common case:
+            // - The number of choices is <= 256. This let's us get a single byte per choice.
+            // - The number of choices is a power of two. This let's us use a byte and simply mask off
+            //   unnecessary bits cheaply rather than needing to use rejection sampling.
+            if (BitOperations.IsPow2(choices.Length) && choices.Length <= 256)
+            {
+                const int StackAllocThreshold = 256;
+
+                // Get a scratch buffer to fill with random bytes, either stack space if the
+                // destination is small enough, or an array pool array if it's larger.
+                // If T is unmanaged, we _could_ use the destination buffer itself, reinterpreted
+                // as a byte buffer, to hold the scratch bytes, but then we'd potentially be temporarily
+                // writing invalid Ts into the destination, which is a risk.
+                byte[]? arrayPoolArray = null;
+                Span<byte> randomBytes = destination.Length <= StackAllocThreshold ?
+                    stackalloc byte[StackAllocThreshold] :
+                    (arrayPoolArray = ArrayPool<byte>.Shared.Rent(destination.Length));
+                randomBytes = randomBytes.Slice(0, destination.Length);
+
+                NextBytes(randomBytes);
+
+                int mask = choices.Length - 1;
+                for (int i = 0; i < destination.Length; i++)
+                {
+                    destination[i] = choices[randomBytes[i] & mask];
+                }
+
+                if (arrayPoolArray is not null)
+                {
+                    ArrayPool<byte>.Shared.Return(arrayPoolArray);
+                }
+
+                return;
+            }
+
+            // Simple fallback: get each item individually, generating a new random Int32 for each
+            // item. This is slower than the above, but it works for all types and sizes of choices.
             for (int i = 0; i < destination.Length; i++)
             {
                 destination[i] = choices[Next(choices.Length)];

src/libraries/System.Security.Cryptography/src/System/Security/Cryptography/RandomNumberGenerator.cs

@@ -2,8 +2,11 @@
 // The .NET Foundation licenses this file to you under the MIT license.
 
 using System.Buffers;
+using System.Buffers.Binary;
 using System.Diagnostics;
 using System.Diagnostics.CodeAnalysis;
+using System.Numerics;
+using System.Runtime.CompilerServices;
 using System.Runtime.InteropServices;
 
 namespace System.Security.Cryptography
@@ -350,6 +353,44 @@ private static void GetHexStringCore(Span<char> destination, bool lowercase)
 
         private static void GetItemsCore<T>(ReadOnlySpan<T> choices, Span<T> destination)
         {
+            // The most expensive part of this operation is the call to get random data. We can
+            // do so only once rather than per element in a common case:
+            // - The number of choices is <= 256. This let's us get a single byte per choice.
+            // - The number of choices is a power of two. This let's us use a byte and simply mask off
+            //   unnecessary bits cheaply rather than needing to use rejection sampling.
+            if (BitOperations.IsPow2(choices.Length) && choices.Length <= 256)
+            {
+                const int StackAllocThreshold = 256;
+
+                // Get a scratch buffer to fill with random bytes, either stack space if the
+                // destination is small enough, or an array pool array if it's larger.
+                // If T is unmanaged, we _could_ use the destination buffer itself, reinterpreted
+                // as a byte buffer, to hold the scratch bytes, but then we'd potentially be temporarily
+                // writing invalid Ts into the destination, which is a risk.
+                byte[]? arrayPoolArray = null;
+                Span<byte> randomBytes = destination.Length <= StackAllocThreshold ?
+                    stackalloc byte[StackAllocThreshold] :
+                    (arrayPoolArray = ArrayPool<byte>.Shared.Rent(destination.Length));
+                randomBytes = randomBytes.Slice(0, destination.Length);
+
+                Fill(randomBytes);
+
+                int mask = choices.Length - 1;
+                for (int i = 0; i < destination.Length; i++)
+                {
+                    destination[i] = choices[randomBytes[i] & mask];
+                }
+
+                if (arrayPoolArray is not null)
+                {
+                    ArrayPool<byte>.Shared.Return(arrayPoolArray);
+                }
+
+                return;
+            }
+
+            // Simple fallback: get each item individually, generating a new random Int32 for each
+            // item. This is slower than the above, but it works for all types and sizes of choices.
             for (int i = 0; i < destination.Length; i++)
             {
                 destination[i] = choices[GetInt32(choices.Length)];