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.

Author: ToddGrun

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<int> data = [1, 2, 3, 4, 5, 6, 7, 8, 9, 10];
+        ImmutableArray<int> expected = [2, 4, 6, 8, 10];
+
+        var actual = data.WhereAsArray(static x => x % 2 == 0);
+        Assert.Equal<int>(expected, actual);
+    }
+
+    [Fact]
+    public void WhereAsArray_ImmutableArray_None()
+    {
+        ImmutableArray<int> data = [1, 2, 3, 4, 5, 6, 7, 8, 9, 10];
+        ImmutableArray<int> expected = [];
+
+        var actual = data.WhereAsArray(static x => false);
+        Assert.Equal<int>(expected, actual);
+    }
+
+    [Fact]
+    public void WhereAsArray_ImmutableArray_All()
+    {
+        ImmutableArray<int> data = [1, 2, 3, 4, 5, 6, 7, 8, 9, 10];
+        var expected = data;
+
+        var actual = data.WhereAsArray(static x => true);
+        Assert.Equal<int>(expected, actual);
+        Assert.Same(ImmutableCollectionsMarshal.AsArray(expected), ImmutableCollectionsMarshal.AsArray(actual));
+    }
+
+    [Fact]
+    public void WhereAsArray_ImmutableArrayBuilder()
+    {
+        var data = ImmutableArray.CreateBuilder<int>();
+        data.AddRange(new[] { 1, 2, 3, 4, 5, 6, 7, 8, 9, 10 });
+        ImmutableArray<int> expected = [2, 4, 6, 8, 10];
+
+        var actual = data.WhereAsArray(static x => x % 2 == 0);
+        Assert.Equal<int>(expected, actual);
+    }
+
+    [Fact]
+    public void WhereAsArray_ImmutableArrayBuilder_None()
+    {
+        var data = ImmutableArray.CreateBuilder<int>();
+        data.AddRange(new[] { 1, 2, 3, 4, 5, 6, 7, 8, 9, 10 });
+        ImmutableArray<int> expected = [];
+
+        var actual = data.WhereAsArray(static x => false);
+        Assert.Equal<int>(expected, actual);
+    }
+
+    [Fact]
+    public void WhereAsArray_ImmutableArrayBuilder_All()
+    {
+        var data = ImmutableArray.CreateBuilder<int>();
+        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<int>(expected, actual);
+    }
 }

src/Shared/Microsoft.AspNetCore.Razor.Utilities.Shared/ImmutableArrayExtensions.cs

@@ -146,22 +146,114 @@ public static ImmutableArray<TResult> SelectManyAsArray<TSource, TResult>(this I
 
     public static ImmutableArray<T> WhereAsArray<T>(this ImmutableArray<T> source, Func<T, bool> predicate)
     {
-        if (source is [])
+        using var builder = new PooledArrayBuilder<T>();
+        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<T>.Empty;
+        }
+        else
+        {
+            return builder.ToImmutableAndClear();
+        }
+    }
+
+    public static ImmutableArray<T> WhereAsArray<T>(this ImmutableArray<T>.Builder source, Func<T, bool> predicate)
+    {
         using var builder = new PooledArrayBuilder<T>();
+        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<T>.Empty;
+        }
+        else
+        {
+            return builder.ToImmutableAndClear();
+        }
     }
 
     /// <summary>