Use ITextDifferencingService.DiffSnapshotSpans instead of allocating …

…full buffer text during codeaction previews (#74114)

* Use DiffSnapshotSpans instead of allocating full buffer text during codeaction previews

This is enabled by changing ChangedSourceText to use the base class's ITextSnapshot based ctor instead of the ITextImage one. The base class then keeps a mapping of ITextImage -> ITextSnapshot, that can be used during the preview to map from a SourceText to an ITextSnapshot. Having the ITextSnapshot available allows the code to use the DiffService call that takes in SnapshotSpans instead of needing to allocate a string for the whole sourcetext.

The obvious concern here would be whether switching ChangedSourceText to have the ITextImage -> ITextSnapshot weakly kept is going to root the ITextSnapshot for longer than it would before. It appears that isn't too much of a worry though, as the SnapshotSourceText ctor that takes in the snapshot only uses it for the CWT mapping, and doesn't store it otherwise.

src/EditorFeatures/Core/Preview/AbstractPreviewFactoryService.cs

@@ -11,6 +11,7 @@
 using System.Threading;
 using System.Threading.Tasks;
 using Microsoft.CodeAnalysis.CodeActions;
+using Microsoft.CodeAnalysis.Editor.Implementation.TextDiffing;
 using Microsoft.CodeAnalysis.Editor.Shared.Extensions;
 using Microsoft.CodeAnalysis.Editor.Shared.Preview;
 using Microsoft.CodeAnalysis.Editor.Shared.Utilities;
@@ -56,6 +57,11 @@ internal abstract class AbstractPreviewFactoryService<TDifferenceViewer>(
     private readonly IDifferenceBufferFactoryService _differenceBufferService = differenceBufferService;
     private readonly ITextDocumentFactoryService _textDocumentFactoryService = textDocumentFactoryService;
 
+    private static readonly StringDifferenceOptions s_differenceOptions = new()
+    {
+        DifferenceType = StringDifferenceTypes.Word | StringDifferenceTypes.Line,
+    };
+
     protected readonly IThreadingContext ThreadingContext = threadingContext;
 
     public SolutionPreviewResult? GetSolutionPreviews(Solution oldSolution, Solution? newSolution, CancellationToken cancellationToken)
@@ -778,10 +784,8 @@ private IHierarchicalDifferenceCollection ComputeEditDifferences(TextDocument ol
             oldDocument.Project.Services.GetRequiredService<IContentTypeLanguageService>().GetDefaultContentType());
 
         diffService ??= _differenceSelectorService.DefaultTextDifferencingService;
-        return diffService.DiffStrings(oldText.ToString(), newText.ToString(), new StringDifferenceOptions()
-        {
-            DifferenceType = StringDifferenceTypes.Word | StringDifferenceTypes.Line,
-        });
+
+        return diffService.DiffSourceTexts(oldText, newText, s_differenceOptions);
     }
 
     private static NormalizedSpanCollection GetOriginalSpans(IHierarchicalDifferenceCollection diffResult, CancellationToken cancellationToken)

src/EditorFeatures/Core/TextDiffing/EditorTextDifferencingService.cs

@@ -13,7 +13,6 @@
 using Microsoft.CodeAnalysis.Editor.Shared.Extensions;
 using Microsoft.CodeAnalysis.Host.Mef;
 using Microsoft.CodeAnalysis.Text;
-using Microsoft.CodeAnalysis.Text.Shared.Extensions;
 using Microsoft.VisualStudio.Text.Differencing;
 
 namespace Microsoft.CodeAnalysis.Editor.Implementation.TextDiffing;
@@ -38,13 +37,7 @@ public async Task<ImmutableArray<TextChange>> GetTextChangesAsync(Document oldDo
 
         var differenceOptions = GetDifferenceOptions(preferredDifferenceType);
 
-        var oldTextSnapshot = oldText.FindCorrespondingEditorTextSnapshot();
-        var newTextSnapshot = newText.FindCorrespondingEditorTextSnapshot();
-        var useSnapshots = oldTextSnapshot != null && newTextSnapshot != null;
-
-        var diffResult = useSnapshots
-            ? diffService.DiffSnapshotSpans(oldTextSnapshot.GetFullSpan(), newTextSnapshot.GetFullSpan(), differenceOptions)
-            : diffService.DiffStrings(oldText.ToString(), newText.ToString(), differenceOptions);
+        var diffResult = diffService.DiffSourceTexts(oldText, newText, differenceOptions);
 
         return diffResult.Differences.Select(d =>
             new TextChange(

src/EditorFeatures/Core/TextDiffing/TextDifferencingServiceExtensions.cs

@@ -0,0 +1,25 @@
+// Licensed to the .NET Foundation under one or more agreements.
+// The .NET Foundation licenses this file to you under the MIT license.
+// See the LICENSE file in the project root for more information.
+
+using Microsoft.CodeAnalysis.Text;
+using Microsoft.CodeAnalysis.Text.Shared.Extensions;
+using Microsoft.VisualStudio.Text.Differencing;
+
+namespace Microsoft.CodeAnalysis.Editor.Implementation.TextDiffing;
+
+internal static class TextDifferencingServiceExtensions
+{
+    public static IHierarchicalDifferenceCollection DiffSourceTexts(this ITextDifferencingService diffService, SourceText oldText, SourceText newText, StringDifferenceOptions options)
+    {
+        var oldTextSnapshot = oldText.FindCorrespondingEditorTextSnapshot();
+        var newTextSnapshot = newText.FindCorrespondingEditorTextSnapshot();
+        var useSnapshots = oldTextSnapshot != null && newTextSnapshot != null;
+
+        var diffResult = useSnapshots
+            ? diffService.DiffSnapshotSpans(oldTextSnapshot!.GetFullSpan(), newTextSnapshot!.GetFullSpan(), options)
+            : diffService.DiffStrings(oldText.ToString(), newText.ToString(), options);
+
+        return diffResult;
+    }
+}

src/EditorFeatures/Text/Extensions.SnapshotSourceText.cs

@@ -35,14 +35,14 @@ private class SnapshotSourceText : SourceText
             private readonly Encoding? _encoding;
             private readonly TextBufferContainer? _container;
 
-            private SnapshotSourceText(ITextBufferCloneService? textBufferCloneService, ITextSnapshot editorSnapshot, SourceHashAlgorithm checksumAlgorithm, TextBufferContainer container)
+            private SnapshotSourceText(ITextBufferCloneService? textBufferCloneService, ITextSnapshot editorSnapshot, Encoding? encoding, SourceHashAlgorithm checksumAlgorithm, TextBufferContainer container)
                 : base(checksumAlgorithm: checksumAlgorithm)
             {
                 Contract.ThrowIfNull(editorSnapshot);
 
                 _textBufferCloneService = textBufferCloneService;
                 this.TextImage = RecordReverseMapAndGetImage(editorSnapshot);
-                _encoding = editorSnapshot.TextBuffer.GetEncodingOrUTF8();
+                _encoding = encoding ?? editorSnapshot.TextBuffer.GetEncodingOrUTF8();
                 _container = container;
             }
 
@@ -83,7 +83,7 @@ public static SourceText From(ITextBufferCloneService? textBufferCloneService, I
 
                     // Avoid capturing `textBufferCloneServiceOpt` on the fast path
                     var tempTextBufferCloneService = textBufferCloneService;
-                    snapshot = s_textSnapshotMap.GetValue(editorSnapshot, s => new SnapshotSourceText(tempTextBufferCloneService, s, SourceHashAlgorithms.OpenDocumentChecksumAlgorithm, container));
+                    snapshot = s_textSnapshotMap.GetValue(editorSnapshot, s => new SnapshotSourceText(tempTextBufferCloneService, s, encoding: null, SourceHashAlgorithms.OpenDocumentChecksumAlgorithm, container));
                 }
 
                 return snapshot;
@@ -100,7 +100,7 @@ internal static SourceText From(ITextBufferCloneService? textBufferCloneService,
                 }
 
                 Contract.ThrowIfFalse(editorSnapshot.TextBuffer == container.GetTextBuffer());
-                return s_textSnapshotMap.GetValue(editorSnapshot, s => new SnapshotSourceText(textBufferCloneService, s, SourceHashAlgorithms.OpenDocumentChecksumAlgorithm, container));
+                return s_textSnapshotMap.GetValue(editorSnapshot, s => new SnapshotSourceText(textBufferCloneService, s, encoding: null, SourceHashAlgorithms.OpenDocumentChecksumAlgorithm, container));
             }
 
             public override Encoding? Encoding
@@ -220,8 +220,8 @@ public override SourceText WithChanges(IEnumerable<TextChange> changes)
                 return new ChangedSourceText(
                     textBufferCloneService: _textBufferCloneService,
                     baseText: this,
-                    baseSnapshot: ((ITextSnapshot2)baseSnapshot).TextImage,
-                    currentSnapshot: ((ITextSnapshot2)buffer.CurrentSnapshot).TextImage);
+                    baseSnapshot: baseSnapshot,
+                    currentSnapshot: buffer.CurrentSnapshot);
             }
 
             private static ITextImage RecordReverseMapAndGetImage(ITextSnapshot editorSnapshot)
@@ -276,13 +276,13 @@ public ClosedSnapshotSourceText(ITextBufferCloneService? textBufferCloneService,
             private class ChangedSourceText : SnapshotSourceText
             {
                 private readonly SnapshotSourceText _baseText;
-                private readonly ITextImage _baseSnapshot;
+                private readonly ITextImage _baseTextImage;
 
-                public ChangedSourceText(ITextBufferCloneService? textBufferCloneService, SnapshotSourceText baseText, ITextImage baseSnapshot, ITextImage currentSnapshot)
-                    : base(textBufferCloneService, currentSnapshot, baseText.Encoding, baseText.ChecksumAlgorithm, container: null)
+                public ChangedSourceText(ITextBufferCloneService? textBufferCloneService, SnapshotSourceText baseText, ITextSnapshot baseSnapshot, ITextSnapshot currentSnapshot)
+                    : base(textBufferCloneService, currentSnapshot, baseText.Encoding, baseText.ChecksumAlgorithm, container: TextBufferContainer.From(currentSnapshot.TextBuffer))
                 {
                     _baseText = baseText;
-                    _baseSnapshot = baseSnapshot;
+                    _baseTextImage = ((ITextSnapshot2)baseSnapshot).TextImage;
                 }
 
                 public override IReadOnlyList<TextChangeRange> GetChangeRanges(SourceText oldText)
@@ -303,7 +303,7 @@ public override IReadOnlyList<TextChangeRange> GetChangeRanges(SourceText oldTex
                         return [new TextChangeRange(new TextSpan(0, oldText.Length), this.Length)];
                     }
 
-                    return GetChangeRanges(_baseSnapshot, _baseSnapshot.Length, this.TextImage);
+                    return GetChangeRanges(_baseTextImage, _baseTextImage.Length, this.TextImage);
                 }
             }
 

src/VisualStudio/Core/Def/Preview/FileChange.cs

@@ -9,6 +9,7 @@
 using System.Threading;
 using Microsoft.CodeAnalysis;
 using Microsoft.CodeAnalysis.Editor;
+using Microsoft.CodeAnalysis.Editor.Implementation.TextDiffing;
 using Microsoft.CodeAnalysis.Editor.Shared.Extensions;
 using Microsoft.CodeAnalysis.Shared.Extensions;
 using Microsoft.CodeAnalysis.Text;
@@ -32,6 +33,11 @@ internal class FileChange : AbstractChange
     private readonly ITextBuffer _buffer;
     private readonly IVsImageService2 _imageService;
 
+    private static readonly StringDifferenceOptions s_differenceOptions = new()
+    {
+        DifferenceType = StringDifferenceTypes.Line,
+    };
+
     public FileChange(TextDocument left,
         TextDocument right,
         IComponentModel componentModel,
@@ -238,12 +244,6 @@ private static IHierarchicalDifferenceCollection ComputeDiffSpans(ITextDifferenc
         var oldText = left.GetTextSynchronously(cancellationToken);
         var newText = right.GetTextSynchronously(cancellationToken);
 
-        var oldString = oldText.ToString();
-        var newString = newText.ToString();
-
-        return diffService.DiffStrings(oldString, newString, new StringDifferenceOptions()
-        {
-            DifferenceType = StringDifferenceTypes.Line,
-        });
+        return diffService.DiffSourceTexts(oldText, newText, s_differenceOptions);
     }
 }