From 79823768e461a865019d7bdf70cfe0fcbeb99288 Mon Sep 17 00:00:00 2001
From: Stephen Toub <stoub@microsoft.com>
Date: Fri, 21 Jul 2023 09:32:48 -0400
Subject: [PATCH] Optimize MemoryMarshal.ToEnumerable for arrays and strings
 (#89274)

---
 .../tests/MemoryMarshal/ToEnumerable.cs            | 32 +++++++++++-
 .../Runtime/InteropServices/MemoryMarshal.cs       | 57 +++++++++++++++++++++-
 2 files changed, 85 insertions(+), 4 deletions(-)

diff --git a/src/libraries/System.Memory/tests/MemoryMarshal/ToEnumerable.cs b/src/libraries/System.Memory/tests/MemoryMarshal/ToEnumerable.cs
index a6eafca..61bf44b 100644
--- a/src/libraries/System.Memory/tests/MemoryMarshal/ToEnumerable.cs
+++ b/src/libraries/System.Memory/tests/MemoryMarshal/ToEnumerable.cs
@@ -1,10 +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.Collections.Generic;
-using Xunit;
 using System.Linq;
 using System.Runtime.InteropServices;
+using Xunit;
 
 namespace System.MemoryTests
 {
@@ -81,9 +82,36 @@ namespace System.MemoryTests
         {
             int[] a = { 91, 92, 93 };
             var memory = new Memory<int>(a);
-            IEnumerable<int> enumer = MemoryMarshal.ToEnumerable<int>(memory);
+            IEnumerable<int> enumer = MemoryMarshal.ToEnumerable<int>(memory.Slice(1));
             IEnumerator<int> enumerat = enumer.GetEnumerator();
             Assert.Same(enumer, enumerat);
         }
+
+        [Fact]
+        public static void ToEnumerableChars()
+        {
+            ReadOnlyMemory<char>[] memories = new[]
+            {
+                new char[] { 'a', 'b', 'c' }.AsMemory(), // array
+                "abc".AsMemory(), // string
+                new WrapperMemoryManager<char>(new char[] { 'a', 'b', 'c' }.AsMemory()).Memory // memory manager
+            };
+
+            foreach (ReadOnlyMemory<char> memory in memories)
+            {
+                Assert.Equal(new char[] { 'a', 'b', 'c' }, MemoryMarshal.ToEnumerable(memory));
+                Assert.Equal(new char[] { 'a', 'b' }, MemoryMarshal.ToEnumerable(memory.Slice(0, 2)));
+                Assert.Equal(new char[] { 'b', 'c' }, MemoryMarshal.ToEnumerable(memory.Slice(1)));
+                Assert.Same(Array.Empty<char>(), MemoryMarshal.ToEnumerable(memory.Slice(3)));
+            }
+        }
+
+        private sealed class WrapperMemoryManager<T>(Memory<T> memory) : MemoryManager<T>
+        {
+            public override Span<T> GetSpan() => memory.Span;
+            public override MemoryHandle Pin(int elementIndex = 0) => throw new NotSupportedException();
+            public override void Unpin() => throw new NotSupportedException();
+            protected override void Dispose(bool disposing) { }
+        }
     }
 }
diff --git a/src/libraries/System.Private.CoreLib/src/System/Runtime/InteropServices/MemoryMarshal.cs b/src/libraries/System.Private.CoreLib/src/System/Runtime/InteropServices/MemoryMarshal.cs
index 7625ebb..05251d7 100644
--- a/src/libraries/System.Private.CoreLib/src/System/Runtime/InteropServices/MemoryMarshal.cs
+++ b/src/libraries/System.Private.CoreLib/src/System/Runtime/InteropServices/MemoryMarshal.cs
@@ -382,8 +382,61 @@ namespace System.Runtime.InteropServices
         /// <returns>An <see cref="IEnumerable{T}"/> view of the given <paramref name="memory" /></returns>
         public static IEnumerable<T> ToEnumerable<T>(ReadOnlyMemory<T> memory)
         {
-            for (int i = 0; i < memory.Length; i++)
-                yield return memory.Span[i];
+            object? obj = memory.GetObjectStartLength(out int index, out int length);
+
+            // If the memory is empty, just return an empty array as the enumerable.
+            if (length is 0 || obj is null)
+            {
+                return Array.Empty<T>();
+            }
+
+            // If the object is a string, we can optimize. If it isn't a slice, just return the string as the
+            // enumerable. Otherwise, return an iterator dedicated to enumerating the object; while we could
+            // use the general one for any ReadOnlyMemory, that will incur a .Span access for every element.
+            if (typeof(T) == typeof(char) && obj is string str)
+            {
+                return (IEnumerable<T>)(object)(index == 0 && length == str.Length ?
+                    str :
+                    FromString(str, index, length));
+
+                static IEnumerable<char> FromString(string s, int offset, int count)
+                {
+                    for (int i = 0; i < count; i++)
+                    {
+                        yield return s[offset + i];
+                    }
+                }
+            }
+
+            // If the object is an array, we can optimize. If it isn't a slice, just return the array as the
+            // enumerable. Otherwise, return an iterator dedicated to enumerating the object.
+            if (RuntimeHelpers.ObjectHasComponentSize(obj)) // Same check as in TryGetArray to confirm that obj is a T[] or a U[] which is blittable to a T[].
+            {
+                T[] array = Unsafe.As<T[]>(obj);
+                index &= ReadOnlyMemory<T>.RemoveFlagsBitMask; // the array may be prepinned, so remove the high bit from the start index in the line below.
+                return index == 0 && length == array.Length ?
+                    array :
+                    FromArray(array, index, length);
+
+                static IEnumerable<T> FromArray(T[] array, int offset, int count)
+                {
+                    for (int i = 0; i < count; i++)
+                    {
+                        yield return array[offset + i];
+                    }
+                }
+            }
+
+            // The ROM<T> wraps a MemoryManager<T>. The best we can do is iterate, accessing .Span on each MoveNext.
+            return FromMemoryManager(memory);
+
+            static IEnumerable<T> FromMemoryManager(ReadOnlyMemory<T> memory)
+            {
+                for (int i = 0; i < memory.Length; i++)
+                {
+                    yield return memory.Span[i];
+                }
+            }
         }
 
         /// <summary>Attempts to get the underlying <see cref="string"/> from a <see cref="ReadOnlyMemory{T}"/>.</summary>
-- 
2.7.4