Modify RazorCSharpDocument to keep SourceMappings sorted by both original and generated document positions (#12772)

* Modify RazorCSharpDocument to keep sourcemappings sorted by both Original and Generated document indices

This class previously just kept an array that was sorted by generated positions. Many callers would benefit instead from an array sorted by original (html) posititions. This allows many loops to short-circuit once they've moved beyond any mappings that could be of interest to them.

This also renamed the original array to indicate it's sorted by generated positions, and to validate that upon it's input, as some callers were assuming that. By renaming this variable, it will hopefully make callers

1) Aware that the input is sorted and may be short-circuited or binary searched
2) Aware what it is sorted by (eg: formatting had some code that was assuming the array was sorted by original position)

* Be more resilient to bad input

* Use Debug.Fail

* Fix test to do different validation in debug/release

Author: ToddGrun

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

@@ -2,6 +2,7 @@
 // The .NET Foundation licenses this file to you under the MIT license.
 
 using System.Collections.Immutable;
+using System.Diagnostics;
 using Microsoft.AspNetCore.Razor.Language.CodeGeneration;
 using Microsoft.CodeAnalysis.Text;
 
@@ -12,7 +13,8 @@ public sealed class RazorCSharpDocument
     public RazorCodeDocument CodeDocument { get; }
     public SourceText Text { get; }
     public ImmutableArray<RazorDiagnostic> Diagnostics { get; }
-    public ImmutableArray<SourceMapping> SourceMappings { get; }
+    public ImmutableArray<SourceMapping> SourceMappingsSortedByGenerated { get; }
+    public ImmutableArray<SourceMapping> SourceMappingsSortedByOriginal { get; }
     public ImmutableArray<LinePragma> LinePragmas { get; }
 
     public RazorCSharpDocument(
@@ -29,7 +31,21 @@ public RazorCSharpDocument(
         Text = text;
 
         Diagnostics = diagnostics.NullToEmpty();
-        SourceMappings = sourceMappings.NullToEmpty();
+        SourceMappingsSortedByGenerated = sourceMappings.NullToEmpty();
+
+        // Verify given source mappings are ordered by their generated spans
+        for (var i = 0; i < SourceMappingsSortedByGenerated.Length - 1; i++)
+        {
+            if (SourceMappingsSortedByGenerated[i].GeneratedSpan.CompareByStartThenLength(SourceMappingsSortedByGenerated[i + 1].GeneratedSpan) > 0)
+            {
+                Debug.Fail("input not sorted");
+
+                SourceMappingsSortedByGenerated = SourceMappingsSortedByGenerated.Sort(static (m1, m2) => m1.GeneratedSpan.CompareByStartThenLength(m2.GeneratedSpan));
+                break;
+            }
+        }
+
+        SourceMappingsSortedByOriginal = SourceMappingsSortedByGenerated.Sort(static (m1, m2) => m1.OriginalSpan.CompareByStartThenLength(m2.OriginalSpan));
         LinePragmas = linePragmas.NullToEmpty();
     }
 }

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

@@ -115,4 +115,15 @@ internal readonly SourceSpan Slice(int startIndex, int length)
     {
         return !left.Equals(right);
     }
+
+    public readonly int CompareByStartThenLength(SourceSpan other)
+    {
+        var cmpByStart = AbsoluteIndex.CompareTo(other.AbsoluteIndex);
+        if (cmpByStart != 0)
+        {
+            return cmpByStart;
+        }
+
+        return Length.CompareTo(other.Length);
+    }
 }

src/Razor/benchmarks/Microsoft.AspNetCore.Razor.Microbenchmarks/LanguageServer/RazorDocumentMappingBenchmark.cs

@@ -129,7 +129,7 @@ private bool TryMapToHostDocumentPosition(RazorCSharpDocument csharpDocument, in
             throw new ArgumentNullException(nameof(csharpDocument));
         }
 
-        foreach (var mapping in csharpDocument.SourceMappings)
+        foreach (var mapping in csharpDocument.SourceMappingsSortedByGenerated)
         {
             var generatedSpan = mapping.GeneratedSpan;
             var generatedAbsoluteIndex = generatedSpan.AbsoluteIndex;

src/Razor/src/Microsoft.AspNetCore.Razor.ExternalAccess.LegacyEditor/RazorWrapperFactory.CodeDocumentWrapper.cs

@@ -55,7 +55,7 @@ public ImmutableArray<TagHelperSpan> GetTagHelperSpans()
 
         public ImmutableArray<RazorSourceMapping> GetSourceMappings()
         {
-            var mappings = Object.GetRequiredCSharpDocument().SourceMappings;
+            var mappings = Object.GetRequiredCSharpDocument().SourceMappingsSortedByGenerated;
 
             return WrapAll<SourceMapping, RazorSourceMapping>(mappings, ConvertSourceMapping);
         }

src/Razor/src/Microsoft.CodeAnalysis.Razor.Workspaces/DocumentExcerpt/DocumentExcerptHelper.cs

@@ -18,30 +18,34 @@ internal static class DocumentExcerptHelper
     public static async Task<ImmutableArray<ClassifiedSpan>.Builder> ClassifyPreviewAsync(
         TextSpan excerptSpan,
         Document generatedDocument,
-        ImmutableArray<SourceMapping> mappings,
+        ImmutableArray<SourceMapping> mappingsSortedByOriginal,
         RazorClassificationOptionsWrapper options,
         CancellationToken cancellationToken)
     {
         var builder = ImmutableArray.CreateBuilder<ClassifiedSpan>();
 
-        var sorted = mappings.Sort((x, y) => x.OriginalSpan.AbsoluteIndex.CompareTo(y.OriginalSpan.AbsoluteIndex));
-
         // The algorithm here is to iterate through the source mappings (sorted) and use the C# classifier
         // on the spans that are known to be C#. For the spans that are not known to be C# then
         // we just treat them as text since we'd don't currently have our own classifications.
 
-        var remainingSpan = excerptSpan;
-        foreach (var span in sorted)
+        if (excerptSpan.Length == 0)
         {
-            if (excerptSpan.Length == 0)
-            {
-                break;
-            }
+            return builder;
+        }
 
+        var remainingSpan = excerptSpan;
+        foreach (var span in mappingsSortedByOriginal)
+        {
             var primarySpan = span.OriginalSpan.AsTextSpan();
             if (primarySpan.Intersection(remainingSpan) is not TextSpan intersection)
             {
-                // This span is outside the area we're interested in.
+                if (primarySpan.Start > remainingSpan.End)
+                {
+                    // This span (and all following) are after the area we're interested in
+                    break;
+                }
+
+                // This span precedes the area we're interested in
                 continue;
             }
 

src/Razor/src/Microsoft.CodeAnalysis.Razor.Workspaces/DocumentMapping/AbstractDocumentMappingService.cs

@@ -264,7 +264,7 @@ public ImmutableArray<LinePositionSpan> GetCSharpSpansOverlappingRazorSpan(Razor
 
         using var builder = new PooledArrayBuilder<LinePositionSpan>();
 
-        foreach (var mapping in csharpDocument.SourceMappings)
+        foreach (var mapping in csharpDocument.SourceMappingsSortedByOriginal)
         {
             var originalSpan = mapping.OriginalSpan.ToLinePositionSpan();
 
@@ -274,18 +274,19 @@ public ImmutableArray<LinePositionSpan> GetCSharpSpansOverlappingRazorSpan(Razor
 
                 builder.Add(generatedSpan);
             }
+            else if (originalSpan.Start > razorSpan.End)
+            {
+                // This span (and all following) are after the area we're interested in
+                break;
+            }
         }
 
         return builder.ToImmutableAndClear();
     }
 
     public bool TryMapToRazorDocumentPosition(RazorCSharpDocument csharpDocument, int csharpIndex, out LinePosition razorPosition, out int razorIndex)
     {
-        var sourceMappings = csharpDocument.SourceMappings;
-
-        // We expect source mappings to be ordered by their generated document absolute index, because that is how the compiler creates them: As it
-        // outputs the generated file to the text write.
-        Debug.Assert(sourceMappings.SequenceEqual(sourceMappings.OrderBy(s => s.GeneratedSpan.AbsoluteIndex)));
+        var sourceMappings = csharpDocument.SourceMappingsSortedByGenerated;
 
         var index = sourceMappings.BinarySearchBy(csharpIndex, static (mapping, generatedDocumentIndex) =>
         {
@@ -334,7 +335,7 @@ private static bool TryMapToCSharpDocumentPositionInternal(RazorCSharpDocument c
 
         var hostDocumentLine = csharpDocument.CodeDocument.Source.Text.GetLinePosition(razorIndex).Line;
 
-        foreach (var mapping in csharpDocument.SourceMappings)
+        foreach (var mapping in csharpDocument.SourceMappingsSortedByOriginal)
         {
             var originalSpan = mapping.OriginalSpan;
             var originalAbsoluteIndex = originalSpan.AbsoluteIndex;
@@ -361,6 +362,11 @@ private static bool TryMapToCSharpDocumentPositionInternal(RazorCSharpDocument c
                 // not original span, its possible for things to be out of order.
                 nextCSharpMapping = mapping;
             }
+            else
+            {
+                // This span (and all following) are after the area we're interested in
+                break;
+            }
         }
 
         if (nextCSharpPositionOnFailure && nextCSharpMapping is not null)
@@ -434,23 +440,24 @@ private bool TryMapToRazorDocumentRangeInclusive(RazorCSharpDocument csharpDocum
         }
 
         using var _1 = ListPool<SourceMapping>.GetPooledObject(out var candidateMappings);
+        var sourceMappings = csharpDocument.SourceMappingsSortedByGenerated;
         if (startMappedDirectly)
         {
             // Start of generated range intersects with a mapping
             candidateMappings.AddRange(
-                csharpDocument.SourceMappings.Where(mapping => IntersectsWith(startIndex, mapping.GeneratedSpan)));
+                sourceMappings.Where(mapping => IntersectsWith(startIndex, mapping.GeneratedSpan)));
         }
         else if (endMappedDirectly)
         {
             // End of generated range intersects with a mapping
             candidateMappings.AddRange(
-                csharpDocument.SourceMappings.Where(mapping => IntersectsWith(endIndex, mapping.GeneratedSpan)));
+                sourceMappings.Where(mapping => IntersectsWith(endIndex, mapping.GeneratedSpan)));
         }
         else
         {
             // Our range does not intersect with any mapping; we should see if it overlaps generated locations
             candidateMappings.AddRange(
-                csharpDocument.SourceMappings
+                sourceMappings
                     .Where(mapping => Overlaps(csharpSourceText.GetTextSpan(csharpRange), mapping.GeneratedSpan)));
         }
 
@@ -503,20 +510,21 @@ private bool TryMapToRazorDocumentRangeInferred(RazorCSharpDocument csharpDocume
         var generatedRangeAsSpan = csharpSourceText.GetTextSpan(csharpRange);
         SourceMapping? mappingBeforeGeneratedRange = null;
         SourceMapping? mappingAfterGeneratedRange = null;
+        var sourceMappings = csharpDocument.SourceMappingsSortedByGenerated;
 
-        for (var i = csharpDocument.SourceMappings.Length - 1; i >= 0; i--)
+        for (var i = sourceMappings.Length - 1; i >= 0; i--)
         {
-            var sourceMapping = csharpDocument.SourceMappings[i];
+            var sourceMapping = sourceMappings[i];
             var sourceMappingEnd = sourceMapping.GeneratedSpan.AbsoluteIndex + sourceMapping.GeneratedSpan.Length;
             if (generatedRangeAsSpan.Start >= sourceMappingEnd)
             {
                 // This is the source mapping that's before us!
                 mappingBeforeGeneratedRange = sourceMapping;
 
-                if (i + 1 < csharpDocument.SourceMappings.Length)
+                if (i + 1 < sourceMappings.Length)
                 {
                     // We're not at the end of the document there's another source mapping after us
-                    mappingAfterGeneratedRange = csharpDocument.SourceMappings[i + 1];
+                    mappingAfterGeneratedRange = sourceMappings[i + 1];
                 }
 
                 break;

src/Razor/src/Microsoft.CodeAnalysis.Razor.Workspaces/DocumentMapping/RazorEditHelper.cs

@@ -22,13 +22,18 @@ internal static partial class RazorEditHelper
 {
     internal static bool TryGetMappedSpan(TextSpan span, SourceText source, RazorCSharpDocument output, out LinePositionSpan linePositionSpan, out TextSpan mappedSpan)
     {
-        foreach (var mapping in output.SourceMappings)
+        foreach (var mapping in output.SourceMappingsSortedByGenerated)
         {
-            var original = mapping.OriginalSpan.AsTextSpan();
             var generated = mapping.GeneratedSpan.AsTextSpan();
 
             if (!generated.Contains(span))
             {
+                if (generated.Start > span.End)
+                {
+                    // This span (and all following) are after the area we're interested in
+                    break;
+                }
+
                 // If the search span isn't contained within the generated span, it is not a match.
                 // A C# identifier won't cover multiple generated spans.
                 continue;
@@ -39,6 +44,7 @@ internal static bool TryGetMappedSpan(TextSpan span, SourceText source, RazorCSh
             if (leftOffset >= 0 && rightOffset <= 0)
             {
                 // This span mapping contains the span.
+                var original = mapping.OriginalSpan.AsTextSpan();
                 mappedSpan = new TextSpan(original.Start + leftOffset, (original.End + rightOffset) - (original.Start + leftOffset));
                 linePositionSpan = source.GetLinePositionSpan(mappedSpan);
                 return true;

src/Razor/src/Microsoft.CodeAnalysis.Razor.Workspaces/Extensions/RazorCodeDocumentExtensions.cs

@@ -61,7 +61,7 @@ public static bool TryGetMinimalCSharpRange(this RazorCodeDocument codeDocument,
         var csharpDoc = codeDocument.GetRequiredCSharpDocument();
 
         // We want to find the min and max C# source mapping that corresponds with our Razor range.
-        foreach (var mapping in csharpDoc.SourceMappings)
+        foreach (var mapping in csharpDoc.SourceMappingsSortedByOriginal)
         {
             var mappedTextSpan = mapping.OriginalSpan.AsTextSpan();
 
@@ -78,6 +78,11 @@ public static bool TryGetMinimalCSharpRange(this RazorCodeDocument codeDocument,
                     maxGeneratedSpan = mapping.GeneratedSpan;
                 }
             }
+            else if (mappedTextSpan.Start > textSpan.End)
+            {
+                // This span (and all following) are after the area we're interested in
+                break;
+            }
         }
 
         // Create a new projected range based on our calculated min/max source spans.

src/Razor/src/Microsoft.CodeAnalysis.Razor.Workspaces/Formatting/Passes/CSharpFormattingPass.CSharpDocumentGenerator.cs

@@ -195,7 +195,7 @@ public void Generate()
                 using var _ = StringBuilderPool.GetPooledObject(out var additionalLinesBuilder);
 
                 var root = _codeDocument.GetRequiredSyntaxRoot();
-                var sourceMappings = _codeDocument.GetRequiredCSharpDocument().SourceMappings;
+                var sourceMappings = _codeDocument.GetRequiredCSharpDocument().SourceMappingsSortedByOriginal;
                 var iMapping = 0;
                 foreach (var line in _sourceText.Lines)
                 {

src/Razor/src/Microsoft.CodeAnalysis.Razor.Workspaces/Formatting/Passes/CSharpFormattingPass.cs

@@ -33,7 +33,7 @@ public async Task<ImmutableArray<TextChange>> ExecuteAsync(FormattingContext con
         // Process changes from previous passes
         var changedText = context.SourceText.WithChanges(changes);
         var changedContext = await context.WithTextAsync(changedText, cancellationToken).ConfigureAwait(false);
-        context.Logger?.LogObject("SourceMappings", changedContext.CodeDocument.GetRequiredCSharpDocument().SourceMappings);
+        context.Logger?.LogObject("SourceMappings", changedContext.CodeDocument.GetRequiredCSharpDocument().SourceMappingsSortedByGenerated);
 
         var csharpSyntaxTrue = await changedContext.CurrentSnapshot.GetCSharpSyntaxTreeAsync(cancellationToken).ConfigureAwait(false);
         var csharpSyntaxRoot = await csharpSyntaxTrue.GetRootAsync(cancellationToken).ConfigureAwait(false);