Reduce allocations in TagHelperBinder.ProcessDescriptors

src/Compiler/Microsoft.CodeAnalysis.Razor.Compiler/src/Language/TagHelperBinder.cs

@@ -6,6 +6,7 @@
 using System.Collections.Immutable;
 using System.Collections.ObjectModel;
 using System.Diagnostics;
+using System.Runtime.InteropServices;
 using Microsoft.AspNetCore.Razor.PooledObjects;
 
 namespace Microsoft.AspNetCore.Razor.Language;
@@ -42,48 +43,124 @@ private static void ProcessDescriptors(
         out ImmutableArray<TagHelperDescriptor> catchAllDescriptors)
     {
         using var catchAllBuilder = new PooledArrayBuilder<TagHelperDescriptor>();
-        using var pooledMap = StringDictionaryPool<ImmutableArray<TagHelperDescriptor>.Builder>.OrdinalIgnoreCase.GetPooledObject(out var mapBuilder);
-        using var pooledSet = HashSetPool<TagHelperDescriptor>.GetPooledObject(out var distinctSet);
+        using var toVisit = new PooledArrayBuilder<(string, TagHelperDescriptor)>();
+        using var _ = StringDictionaryPool<TagHelperDescriptorArrayBuilder>.OrdinalIgnoreCase.GetPooledObject(out var builderMap);
+
+        // The initial pass does three things:
+        // 1: Fills out builderMap with a structure that can be used to create the array we will return
+        // 2: Fills out toVisit with all the (tagName, descriptor) pairs we need to process in the next pass
+        // 3: Fills out catchAllBuilder with all the catch-all descriptors
+        PerformInitialPass(descriptors, tagNamePrefix, builderMap, ref toVisit.AsRef(), ref catchAllBuilder.AsRef());
+
+        // Perform a pass over toVisit to populate each TagHelperDescriptorArrayBuilder entry in builderMap with descriptors
+        // that match the tag name
+        PopulateArrayBuilders(in toVisit, builderMap);
 
-        // Build a map of tag name -> tag helpers.
-        foreach (var descriptor in descriptors)
+        // Build the final dictionary with immutable arrays
+        var map = new Dictionary<string, ImmutableArray<TagHelperDescriptor>>(capacity: builderMap.Count, StringComparer.OrdinalIgnoreCase);
+
+        foreach (var (tagName, builder) in builderMap)
         {
-            if (!distinctSet.Add(descriptor))
-            {
-                // We're already seen this descriptor, skip it.
-                continue;
-            }
+            map.Add(tagName, builder.ToImmutable());
+        }
+
+        tagNameToDescriptorsMap = new ReadOnlyDictionary<string, ImmutableArray<TagHelperDescriptor>>(map);
+
+        // Build the catch all descriptors array.
+        catchAllDescriptors = catchAllBuilder.ToImmutableAndClear();
 
-            foreach (var rule in descriptor.TagMatchingRules)
+        static void PerformInitialPass(
+            ImmutableArray<TagHelperDescriptor> descriptors,
+            string? tagNamePrefix,
+            Dictionary<string, TagHelperDescriptorArrayBuilder> builderMap,
+            ref PooledArrayBuilder<(string, TagHelperDescriptor)> toVisit,
+            ref PooledArrayBuilder<TagHelperDescriptor> catchAllBuilder)
+        {
+            using var _ = HashSetPool<TagHelperDescriptor>.GetPooledObject(out var distinctSet);
+
+            foreach (var descriptor in descriptors)
             {
-                if (rule.TagName == TagHelperMatchingConventions.ElementCatchAllName)
+                if (!distinctSet.Add(descriptor))
                 {
-                    // This is a catch-all descriptor, we can keep track of it separately.
-                    catchAllBuilder.Add(descriptor);
+                    // We're already seen this descriptor, skip it.
+                    continue;
                 }
-                else
+
+                foreach (var rule in descriptor.TagMatchingRules)
                 {
-                    // This is a specific tag name, we need to add it to the map.
-                    var tagName = tagNamePrefix + rule.TagName;
-                    var builder = mapBuilder.GetOrAdd(tagName, _ => ImmutableArray.CreateBuilder<TagHelperDescriptor>());
+                    if (rule.TagName == TagHelperMatchingConventions.ElementCatchAllName)
+                    {
+                        // This is a catch-all descriptor, we can keep track of it separately.
+                        catchAllBuilder.Add(descriptor);
+                    }
+                    else
+                    {
+                        // This is a specific tag name, we need to add it to the map.
+                        var tagName = tagNamePrefix + rule.TagName;
 
-                    builder.Add(descriptor);
+                        if (!builderMap.TryGetValue(tagName, out var builder))
+                        {
+                            builder = default;
+                        }
+
+                        builder.IncreaseSize();
+
+                        // Copy back to the dictionary.
+                        builderMap[tagName] = builder;
+
+                        // Ensure we visit tagName and descriptor in the next pass.
+                        toVisit.Add((tagName, descriptor));
+                    }
                 }
             }
         }
 
-        // Build the final dictionary with immutable arrays.
-        var map = new Dictionary<string, ImmutableArray<TagHelperDescriptor>>(capacity: mapBuilder.Count, StringComparer.OrdinalIgnoreCase);
+        static void PopulateArrayBuilders(
+            ref readonly PooledArrayBuilder<(string, TagHelperDescriptor)> toVisit,
+            Dictionary<string, TagHelperDescriptorArrayBuilder> builderMap)
+        {
+            foreach (var (tagName, descriptor) in toVisit)
+            {
+                var builder = builderMap[tagName];
+
+                builder.Add(descriptor);
+
+                builderMap[tagName] = builder;
+            }
+        }
+    }
+
+    private struct TagHelperDescriptorArrayBuilder
+    {
+        private int _size;
+        private TagHelperDescriptor[]? _array;
+        private int _index;
 
-        foreach (var (key, value) in mapBuilder)
+        public void IncreaseSize()
         {
-            map.Add(key, value.ToImmutableAndClear());
+            Debug.Assert(_array is null, "Can't increase size once the array has been created");
+            _size++;
         }
 
-        tagNameToDescriptorsMap = new ReadOnlyDictionary<string, ImmutableArray<TagHelperDescriptor>>(map);
+        public void Add(TagHelperDescriptor value)
+        {
+            Debug.Assert(_index < _size, "Can't add past the end of the array");
+            var array = _array ??= new TagHelperDescriptor[_size];
+            array[_index++] = value;
+        }
 
-        // Build the catch all descriptors array.
-        catchAllDescriptors = catchAllBuilder.ToImmutableAndClear();
+        public readonly ImmutableArray<TagHelperDescriptor> ToImmutable()
+        {
+            if (_size == 0)
+            {
+                return [];
+            }
+
+            Debug.Assert(_array is not null);
+            Debug.Assert(_index == _size, "Can't produce the final array until it's filled.");
+
+            return ImmutableCollectionsMarshal.AsImmutableArray(_array);
+        }
     }
 
     /// <summary>