From e99e981f27f1b35803a18299f9d6cb26e9491401 Mon Sep 17 00:00:00 2001 From: Todd Grunke Date: Wed, 23 Jul 2025 11:32:08 -0700 Subject: [PATCH] Reduce allocations in WhereAsArray IA/IA.Builder extension methods Very heavily inspired by Roslyn's implementation (https://github.com/dotnet/roslyn/blob/fadd60c39d63743cdaeacd2c3f8aa003f5bb7f58/src/Dependencies/Collections/Extensions/ImmutableArrayExtensions.cs#L546-L613) I needed the IA.Builder version of this for local changes, but separated this out into a separate PR. --- .../ImmutableArrayExtensionsTests.cs | 65 +++++++++++ .../ImmutableArrayExtensions.cs | 104 +++++++++++++++++- 2 files changed, 163 insertions(+), 6 deletions(-) diff --git a/src/Shared/Microsoft.AspNetCore.Razor.Utilities.Shared.Test/ImmutableArrayExtensionsTests.cs b/src/Shared/Microsoft.AspNetCore.Razor.Utilities.Shared.Test/ImmutableArrayExtensionsTests.cs index c121dc7efb0..16c3ecc577e 100644 --- a/src/Shared/Microsoft.AspNetCore.Razor.Utilities.Shared.Test/ImmutableArrayExtensionsTests.cs +++ b/src/Shared/Microsoft.AspNetCore.Razor.Utilities.Shared.Test/ImmutableArrayExtensionsTests.cs @@ -4,6 +4,7 @@ using System; using System.Collections.Generic; using System.Collections.Immutable; +using System.Runtime.InteropServices; using Xunit; namespace Microsoft.AspNetCore.Razor.Utilities.Shared.Test; @@ -232,4 +233,68 @@ public void InsertRange_WithReferenceTypes_InsertsCorrectly() Assert.Equal(4, builder.Count); Assert.Equal(["apple", "cherry", "date", "banana"], builder.ToArray()); } + + [Fact] + public void WhereAsArray_ImmutableArray() + { + ImmutableArray data = [1, 2, 3, 4, 5, 6, 7, 8, 9, 10]; + ImmutableArray expected = [2, 4, 6, 8, 10]; + + var actual = data.WhereAsArray(static x => x % 2 == 0); + Assert.Equal(expected, actual); + } + + [Fact] + public void WhereAsArray_ImmutableArray_None() + { + ImmutableArray data = [1, 2, 3, 4, 5, 6, 7, 8, 9, 10]; + ImmutableArray expected = []; + + var actual = data.WhereAsArray(static x => false); + Assert.Equal(expected, actual); + } + + [Fact] + public void WhereAsArray_ImmutableArray_All() + { + ImmutableArray data = [1, 2, 3, 4, 5, 6, 7, 8, 9, 10]; + var expected = data; + + var actual = data.WhereAsArray(static x => true); + Assert.Equal(expected, actual); + Assert.Same(ImmutableCollectionsMarshal.AsArray(expected), ImmutableCollectionsMarshal.AsArray(actual)); + } + + [Fact] + public void WhereAsArray_ImmutableArrayBuilder() + { + var data = ImmutableArray.CreateBuilder(); + data.AddRange(new[] { 1, 2, 3, 4, 5, 6, 7, 8, 9, 10 }); + ImmutableArray expected = [2, 4, 6, 8, 10]; + + var actual = data.WhereAsArray(static x => x % 2 == 0); + Assert.Equal(expected, actual); + } + + [Fact] + public void WhereAsArray_ImmutableArrayBuilder_None() + { + var data = ImmutableArray.CreateBuilder(); + data.AddRange(new[] { 1, 2, 3, 4, 5, 6, 7, 8, 9, 10 }); + ImmutableArray expected = []; + + var actual = data.WhereAsArray(static x => false); + Assert.Equal(expected, actual); + } + + [Fact] + public void WhereAsArray_ImmutableArrayBuilder_All() + { + var data = ImmutableArray.CreateBuilder(); + data.AddRange(new[] { 1, 2, 3, 4, 5, 6, 7, 8, 9, 10 }); + var expected = data; + + var actual = data.WhereAsArray(static x => true); + Assert.Equal(expected, actual); + } } diff --git a/src/Shared/Microsoft.AspNetCore.Razor.Utilities.Shared/ImmutableArrayExtensions.cs b/src/Shared/Microsoft.AspNetCore.Razor.Utilities.Shared/ImmutableArrayExtensions.cs index 5ddeaabe06f..56984b4cd10 100644 --- a/src/Shared/Microsoft.AspNetCore.Razor.Utilities.Shared/ImmutableArrayExtensions.cs +++ b/src/Shared/Microsoft.AspNetCore.Razor.Utilities.Shared/ImmutableArrayExtensions.cs @@ -146,22 +146,114 @@ public static ImmutableArray SelectManyAsArray(this I public static ImmutableArray WhereAsArray(this ImmutableArray source, Func predicate) { - if (source is []) + using var builder = new PooledArrayBuilder(); + var none = true; + var all = true; + + var n = source.Length; + for (var i = 0; i < n; i++) { - return []; + var a = source[i]; + + if (predicate(a)) + { + none = false; + if (all) + { + continue; + } + + Debug.Assert(i > 0); + builder.Add(a); + } + else + { + if (none) + { + all = false; + continue; + } + + Debug.Assert(i > 0); + if (all) + { + all = false; + for (var j = 0; j < i; j++) + { + builder.Add(source[j]); + } + } + } } + if (all) + { + return source; + } + else if (none) + { + return ImmutableArray.Empty; + } + else + { + return builder.ToImmutableAndClear(); + } + } + + public static ImmutableArray WhereAsArray(this ImmutableArray.Builder source, Func predicate) + { using var builder = new PooledArrayBuilder(); + var none = true; + var all = true; - foreach (var item in source) + var n = source.Count; + for (var i = 0; i < n; i++) { - if (predicate(item)) + var a = source[i]; + + if (predicate(a)) + { + none = false; + if (all) + { + continue; + } + + Debug.Assert(i > 0); + builder.Add(a); + } + else { - builder.Add(item); + if (none) + { + all = false; + continue; + } + + Debug.Assert(i > 0); + if (all) + { + all = false; + for (var j = 0; j < i; j++) + { + builder.Add(source[j]); + } + } } } - return builder.ToImmutableAndClear(); + if (all) + { + return source.ToImmutable(); + } + else if (none) + { + return ImmutableArray.Empty; + } + else + { + return builder.ToImmutableAndClear(); + } } ///