Return ImmutableArray<T> in legacy RazorSyntaxTreeExtensions

src/Compiler/Microsoft.AspNetCore.Razor.Language/src/ClassifiedSpanVisitor.cs

@@ -3,6 +3,7 @@
 
 using System;
 using System.Collections.Generic;
+using System.Collections.Immutable;
 using System.Linq;
 using Microsoft.AspNetCore.Razor.Language.Legacy;
 using Microsoft.AspNetCore.Razor.Language.Syntax;
@@ -12,7 +13,7 @@ namespace Microsoft.AspNetCore.Razor.Language;
 internal class ClassifiedSpanVisitor : SyntaxWalker
 {
     private readonly RazorSourceDocument _source;
-    private readonly List<ClassifiedSpanInternal> _spans;
+    private readonly ImmutableArray<ClassifiedSpanInternal>.Builder _spans;
 
     private readonly Action<CSharpCodeBlockSyntax> _baseVisitCSharpCodeBlock;
     private readonly Action<CSharpStatementSyntax> _baseVisitCSharpStatement;
@@ -29,15 +30,10 @@ internal class ClassifiedSpanVisitor : SyntaxWalker
     private BlockKindInternal _currentBlockKind;
     private SyntaxNode? _currentBlock;
 
-    public ClassifiedSpanVisitor(RazorSourceDocument source)
+    public ClassifiedSpanVisitor(RazorSourceDocument source, ImmutableArray<ClassifiedSpanInternal>.Builder spans)
     {
-        if (source is null)
-        {
-            throw new ArgumentNullException(nameof(source));
-        }
-
-        _source = source;
-        _spans = new List<ClassifiedSpanInternal>();
+        _source = source ?? throw new ArgumentNullException(nameof(source));
+        _spans = spans ?? throw new ArgumentNullException(nameof(spans));
 
         _baseVisitCSharpCodeBlock = base.VisitCSharpCodeBlock;
         _baseVisitCSharpStatement = base.VisitCSharpStatement;

src/Compiler/Microsoft.AspNetCore.Razor.Language/src/ItemCollection.cs

@@ -7,6 +7,7 @@
 using System.Collections;
 using System.Collections.Concurrent;
 using System.Collections.Generic;
+using System.Diagnostics.CodeAnalysis;
 
 namespace Microsoft.AspNetCore.Razor.Language;
 
@@ -128,4 +129,17 @@ void ICollection.CopyTo(Array array, int index)
     {
         ((ICollection)_inner).CopyTo(array, index);
     }
+
+    internal bool TryGetValue<TKey, TValue>(TKey key, [MaybeNullWhen(false)] out TValue value)
+        where TKey : notnull
+    {
+        if (!_inner.TryGetValue(key, out var objValue))
+        {
+            value = default;
+            return false;
+        }
+
+        value = (TValue)objValue;
+        return true;
+    }
 }

src/Compiler/Microsoft.AspNetCore.Razor.Language/src/Legacy/RazorSyntaxTreeExtensions.cs

@@ -1,38 +1,41 @@
 // Licensed to the .NET Foundation under one or more agreements.
 // The .NET Foundation licenses this file to you under the MIT license.
 
-#nullable disable
-
 using System;
-using System.Collections.Generic;
+using System.Collections.Immutable;
+using Microsoft.AspNetCore.Razor.PooledObjects;
 
 namespace Microsoft.AspNetCore.Razor.Language.Legacy;
 
 internal static class RazorSyntaxTreeExtensions
 {
-    public static IReadOnlyList<ClassifiedSpanInternal> GetClassifiedSpans(this RazorSyntaxTree syntaxTree)
+    public static ImmutableArray<ClassifiedSpanInternal> GetClassifiedSpans(this RazorSyntaxTree syntaxTree)
     {
         if (syntaxTree == null)
         {
             throw new ArgumentNullException(nameof(syntaxTree));
         }
 
-        var visitor = new ClassifiedSpanVisitor(syntaxTree.Source);
+        using var _ = ArrayBuilderPool<ClassifiedSpanInternal>.GetPooledObject(out var builder);
+
+        var visitor = new ClassifiedSpanVisitor(syntaxTree.Source, builder);
         visitor.Visit(syntaxTree.Root);
 
-        return visitor.ClassifiedSpans;
+        return builder.DrainToImmutable();
     }
 
-    public static IReadOnlyList<TagHelperSpanInternal> GetTagHelperSpans(this RazorSyntaxTree syntaxTree)
+    public static ImmutableArray<TagHelperSpanInternal> GetTagHelperSpans(this RazorSyntaxTree syntaxTree)
     {
         if (syntaxTree == null)
         {
             throw new ArgumentNullException(nameof(syntaxTree));
         }
 
-        var visitor = new TagHelperSpanVisitor(syntaxTree.Source);
+        using var _ = ArrayBuilderPool<TagHelperSpanInternal>.GetPooledObject(out var builder);
+
+        var visitor = new TagHelperSpanVisitor(syntaxTree.Source, builder);
         visitor.Visit(syntaxTree.Root);
 
-        return visitor.TagHelperSpans;
+        return builder.DrainToImmutable();
     }
 }

src/Compiler/Microsoft.AspNetCore.Razor.Language/src/TagHelperSpanVisitor.cs

@@ -1,9 +1,7 @@
 // Licensed to the .NET Foundation under one or more agreements.
 // The .NET Foundation licenses this file to you under the MIT license.
 
-#nullable disable
-
-using System.Collections.Generic;
+using System.Collections.Immutable;
 using Microsoft.AspNetCore.Razor.Language.Legacy;
 using Microsoft.AspNetCore.Razor.Language.Syntax;
 
@@ -12,16 +10,14 @@ namespace Microsoft.AspNetCore.Razor.Language;
 internal class TagHelperSpanVisitor : SyntaxWalker
 {
     private readonly RazorSourceDocument _source;
-    private readonly List<TagHelperSpanInternal> _spans;
+    private readonly ImmutableArray<TagHelperSpanInternal>.Builder _spans;
 
-    public TagHelperSpanVisitor(RazorSourceDocument source)
+    public TagHelperSpanVisitor(RazorSourceDocument source, ImmutableArray<TagHelperSpanInternal>.Builder spans)
     {
         _source = source;
-        _spans = new List<TagHelperSpanInternal>();
+        _spans = spans;
     }
 
-    public IReadOnlyList<TagHelperSpanInternal> TagHelperSpans => _spans;
-
     public override void VisitMarkupTagHelperElement(MarkupTagHelperElementSyntax node)
     {
         var span = new TagHelperSpanInternal(node.GetSourceSpan(_source), node.TagHelperInfo.BindingResult);

src/Razor/src/Microsoft.AspNetCore.Razor.LanguageServer/RazorDocumentMappingService.cs

@@ -3,6 +3,7 @@
 
 using System;
 using System.Collections.Generic;
+using System.Collections.Immutable;
 using System.Diagnostics;
 using System.Diagnostics.CodeAnalysis;
 using System.Linq;
@@ -12,6 +13,7 @@
 using Microsoft.AspNetCore.Razor.Language.Legacy;
 using Microsoft.AspNetCore.Razor.LanguageServer.Extensions;
 using Microsoft.AspNetCore.Razor.LanguageServer.Protocol;
+using Microsoft.AspNetCore.Razor.PooledObjects;
 using Microsoft.CodeAnalysis.Razor.Workspaces;
 using Microsoft.CodeAnalysis.Razor.Workspaces.Extensions;
 using Microsoft.CodeAnalysis.Text;
@@ -491,13 +493,13 @@ public async Task<WorkspaceEdit> RemapWorkspaceEditAsync(WorkspaceEdit workspace
 
     // Internal for testing
     internal static RazorLanguageKind GetLanguageKindCore(
-        IReadOnlyList<ClassifiedSpanInternal> classifiedSpans,
-        IReadOnlyList<TagHelperSpanInternal> tagHelperSpans,
+        ImmutableArray<ClassifiedSpanInternal> classifiedSpans,
+        ImmutableArray<TagHelperSpanInternal> tagHelperSpans,
         int hostDocumentIndex,
         int hostDocumentLength,
         bool rightAssociative)
     {
-        for (var i = 0; i < classifiedSpans.Count; i++)
+        for (var i = 0; i < classifiedSpans.Length; i++)
         {
             var classifiedSpan = classifiedSpans[i];
             var span = classifiedSpan.Span;
@@ -522,7 +524,7 @@ internal static RazorLanguageKind GetLanguageKindCore(
                         // of, if we're also at the start of the next one
                         if (rightAssociative)
                         {
-                            if (i < classifiedSpans.Count - 1 && classifiedSpans[i + 1].Span.AbsoluteIndex == hostDocumentIndex)
+                            if (i < classifiedSpans.Length - 1 && classifiedSpans[i + 1].Span.AbsoluteIndex == hostDocumentIndex)
                             {
                                 // If we're at the start of the next span, then use that span
                                 return GetLanguageFromClassifiedSpan(classifiedSpans[i + 1]);
@@ -538,9 +540,8 @@ internal static RazorLanguageKind GetLanguageKindCore(
             }
         }
 
-        for (var i = 0; i < tagHelperSpans.Count; i++)
+        foreach (var tagHelperSpan in tagHelperSpans)
         {
-            var tagHelperSpan = tagHelperSpans[i];
             var span = tagHelperSpan.Span;
 
             if (span.AbsoluteIndex <= hostDocumentIndex)
@@ -562,7 +563,7 @@ internal static RazorLanguageKind GetLanguageKindCore(
 
         // Use the language of the last classified span if we're at the end
         // of the document.
-        if (classifiedSpans.Count != 0 && hostDocumentIndex == hostDocumentLength)
+        if (classifiedSpans.Length != 0 && hostDocumentIndex == hostDocumentLength)
         {
             var lastClassifiedSpan = classifiedSpans.Last();
             return GetLanguageFromClassifiedSpan(lastClassifiedSpan);
@@ -952,39 +953,39 @@ private static SourceText GetGeneratedSourceText(IRazorGeneratedDocument generat
         return codeDocument.GetGeneratedSourceText(generatedDocument);
     }
 
-    private static IReadOnlyList<ClassifiedSpanInternal> GetClassifiedSpans(RazorCodeDocument document)
+    private static ImmutableArray<ClassifiedSpanInternal> GetClassifiedSpans(RazorCodeDocument document)
     {
         // Since this service is called so often, we get a good performance improvement by caching these values
         // for this code document. If the document changes, as the user types, then the document instance will be
         // different, so we don't need to worry about invalidating the cache.
-        var classifiedSpans = (IReadOnlyList<ClassifiedSpanInternal>)document.Items[typeof(ClassifiedSpanInternal)];
-        if (classifiedSpans is null)
+        if (!document.Items.TryGetValue(typeof(ClassifiedSpanInternal), out ImmutableArray<ClassifiedSpanInternal> classifiedSpans))
         {
             var syntaxTree = document.GetSyntaxTree();
 
-            var visitor = new ClassifiedSpanVisitor(syntaxTree.Source);
+            using var _ = ArrayBuilderPool<ClassifiedSpanInternal>.GetPooledObject(out var builder);
+            var visitor = new ClassifiedSpanVisitor(syntaxTree.Source, builder);
             visitor.Visit(syntaxTree.Root);
-            classifiedSpans = visitor.ClassifiedSpans;
+            classifiedSpans = builder.DrainToImmutable();
 
             document.Items[typeof(ClassifiedSpanInternal)] = classifiedSpans;
         }
 
         return classifiedSpans;
     }
 
-    private static IReadOnlyList<TagHelperSpanInternal> GetTagHelperSpans(RazorCodeDocument document)
+    private static ImmutableArray<TagHelperSpanInternal> GetTagHelperSpans(RazorCodeDocument document)
     {
         // Since this service is called so often, we get a good performance improvement by caching these values
         // for this code document. If the document changes, as the user types, then the document instance will be
         // different, so we don't need to worry about invalidating the cache.
-        var tagHelperSpans = (IReadOnlyList<TagHelperSpanInternal>)document.Items[typeof(TagHelperSpanInternal)];
-        if (tagHelperSpans is null)
+        if (!document.Items.TryGetValue(typeof(TagHelperSpanInternal), out ImmutableArray<TagHelperSpanInternal> tagHelperSpans))
         {
             var syntaxTree = document.GetSyntaxTree();
 
-            var visitor = new TagHelperSpanVisitor(syntaxTree.Source);
+            using var _ = ArrayBuilderPool<TagHelperSpanInternal>.GetPooledObject(out var builder);
+            var visitor = new TagHelperSpanVisitor(syntaxTree.Source, builder);
             visitor.Visit(syntaxTree.Root);
-            tagHelperSpans = visitor.TagHelperSpans;
+            tagHelperSpans = builder.DrainToImmutable();
 
             document.Items[typeof(TagHelperSpanInternal)] = tagHelperSpans;
         }

src/Razor/src/Microsoft.AspNetCore.Razor.ProjectEngineHost/Serialization/TagHelperDeltaResult.cs

@@ -26,7 +26,7 @@ public ImmutableArray<TagHelperDescriptor> Apply(ImmutableArray<TagHelperDescrip
         // 1. This TagHelperDeltaResult.Apply where we don't iterate / Contains check the "base" collection.
         // 2. The rest of the Razor project system. Everything there is always indexed / iterated as a list.
         using var _ = ArrayBuilderPool<TagHelperDescriptor>.GetPooledObject(out var newTagHelpers);
-        newTagHelpers.SetCapacityIfNeeded(baseTagHelpers.Length + Added.Length - Removed.Length);
+        newTagHelpers.SetCapacityIfLarger(baseTagHelpers.Length + Added.Length - Removed.Length);
         newTagHelpers.AddRange(Added);
 
         foreach (var existingTagHelper in baseTagHelpers)

src/Razor/test/Microsoft.AspNetCore.Razor.LanguageServer.Test/RazorDocumentMappingServiceTest.cs

@@ -3,6 +3,7 @@
 
 using System;
 using System.Collections.Generic;
+using System.Collections.Immutable;
 using System.Linq;
 using Microsoft.AspNetCore.Razor.Language;
 using Microsoft.AspNetCore.Razor.Language.CodeGeneration;
@@ -868,22 +869,18 @@ public void GetLanguageKindCore_GetsLastClassifiedSpanLanguageIfAtEndOfDocument(
     {
         // Arrange
         var text = $"<strong>Something</strong>{Environment.NewLine}<App>";
-        var classifiedSpans = new List<ClassifiedSpanInternal>()
-        {
-           new ClassifiedSpanInternal(
-               new SourceSpan(0, 0),
+        var classifiedSpans = ImmutableArray.Create<ClassifiedSpanInternal>(
+           new(new SourceSpan(0, 0),
                blockSpan: new SourceSpan(absoluteIndex: 0, lineIndex: 0, characterIndex: 0, length: text.Length),
                SpanKindInternal.Transition,
                blockKind: default,
                acceptedCharacters: default),
-           new ClassifiedSpanInternal(
-               new SourceSpan(0, 26),
+           new(new SourceSpan(0, 26),
                blockSpan: default,
                SpanKindInternal.Markup,
                blockKind: default,
-               acceptedCharacters: default)
-        };
-        var tagHelperSpans = Array.Empty<TagHelperSpanInternal>();
+               acceptedCharacters: default));
+        var tagHelperSpans = ImmutableArray<TagHelperSpanInternal>.Empty;
 
         // Act
         var languageKind = RazorDocumentMappingService.GetLanguageKindCore(classifiedSpans, tagHelperSpans, text.Length, text.Length, rightAssociative: false);
@@ -1068,7 +1065,7 @@ public void GetLanguageKindCore_TagHelperInCSharpRightAssociative()
         Assert.Equal(RazorLanguageKind.Html, languageKind);
     }
 
-    private static (IReadOnlyList<ClassifiedSpanInternal> classifiedSpans, IReadOnlyList<TagHelperSpanInternal> tagHelperSpans) GetClassifiedSpans(string text, IReadOnlyList<TagHelperDescriptor>? tagHelpers = null)
+    private static (ImmutableArray<ClassifiedSpanInternal> classifiedSpans, ImmutableArray<TagHelperSpanInternal> tagHelperSpans) GetClassifiedSpans(string text, IReadOnlyList<TagHelperDescriptor>? tagHelpers = null)
     {
         var codeDocument = CreateCodeDocument(text, tagHelpers);
         var syntaxTree = codeDocument.GetSyntaxTree();

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

@@ -1,6 +1,8 @@
 // Copyright (c) .NET Foundation. All rights reserved.
 // Licensed under the MIT license. See License.txt in the project root for license information.
 
+using Microsoft.AspNetCore.Razor.PooledObjects;
+
 namespace System.Collections.Immutable;
 
 /// <summary>
@@ -16,11 +18,69 @@ public static ImmutableArray<T> NullToEmpty<T>(this ImmutableArray<T> array)
         return array.IsDefault ? ImmutableArray<T>.Empty : array;
     }
 
-    public static void SetCapacityIfNeeded<T>(this ImmutableArray<T>.Builder builder, int newCapacity)
+    public static void SetCapacityIfLarger<T>(this ImmutableArray<T>.Builder builder, int newCapacity)
     {
         if (builder.Capacity < newCapacity)
         {
             builder.Capacity = newCapacity;
         }
     }
+
+    /// <summary>
+    ///  Returns the current contents as an <see cref="ImmutableArray{T}"/> and sets
+    ///  the collection to a zero length array.
+    /// </summary>
+    /// <remarks>
+    ///  If <see cref="ImmutableArray{T}.Builder.Capacity"/> equals
+    ///  <see cref="ImmutableArray{T}.Builder.Count"/>, the internal array will be extracted
+    ///  as an <see cref="ImmutableArray{T}"/> without copying the contents. Otherwise, the
+    ///  contents will be copied into a new array. The collection will then be set to a
+    ///  zero-length array.
+    /// </remarks>
+    /// <returns>An immutable array.</returns>
+    public static ImmutableArray<T> DrainToImmutable<T>(this ImmutableArray<T>.Builder builder)
+    {
+#if NET8_0_OR_GREATER
+        return builder.DrainToImmutable();
+#else
+        if (builder.Count == 0)
+        {
+            return ImmutableArray<T>.Empty;
+        }
+
+        if (builder.Count == builder.Capacity)
+        {
+            return builder.MoveToImmutable();
+        }
+
+        var result = builder.ToImmutable();
+        builder.Clear();
+        return result;
+#endif
+    }
+
+    public static ImmutableArray<TResult> SelectAsArray<T, TResult>(this ImmutableArray<T> source, Func<T, TResult> selector)
+    {
+        return source switch
+        {
+            [] => ImmutableArray<TResult>.Empty,
+            [var item] => ImmutableArray.Create(selector(item)),
+            [var item1, var item2] => ImmutableArray.Create(selector(item1), selector(item2)),
+            [var item1, var item2, var item3] => ImmutableArray.Create(selector(item1), selector(item2), selector(item3)),
+            [var item1, var item2, var item3, var item4] => ImmutableArray.Create(selector(item1), selector(item2), selector(item3), selector(item4)),
+            var items => BuildResult(items, selector)
+        };
+
+        static ImmutableArray<TResult> BuildResult(ImmutableArray<T> items, Func<T, TResult> selector)
+        {
+            using var _ = ArrayBuilderPool<TResult>.GetPooledObject(out var result);
+
+            foreach (var item in items)
+            {
+                result.Add(selector(item));
+            }
+
+            return result.DrainToImmutable();
+        }
+    }
 }