Skip to content

Commit

Permalink
Optimize MemoryMarshal.ToEnumerable for arrays and strings (#89274)
Browse files Browse the repository at this point in the history
  • Loading branch information
stephentoub committed Jul 21, 2023
1 parent c5b4bb0 commit 7982376
Show file tree
Hide file tree
Showing 2 changed files with 85 additions and 4 deletions.
32 changes: 30 additions & 2 deletions src/libraries/System.Memory/tests/MemoryMarshal/ToEnumerable.cs
Original file line number Diff line number Diff line change
@@ -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
{
Expand Down Expand Up @@ -81,9 +82,36 @@ public static void ToEnumerableSameAsIEnumerator()
{
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) { }
}
}
}
Original file line number Diff line number Diff line change
Expand Up @@ -382,8 +382,61 @@ public static bool TryGetArray<T>(ReadOnlyMemory<T> memory, out ArraySegment<T>
/// <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>
Expand Down

0 comments on commit 7982376

Please sign in to comment.