[Sonic] Support multiple C# documents in Rename

src/Razor/src/Razor/src/Microsoft.CodeAnalysis.Remote.Razor/CodeLens/RemoteCodeLensService.cs

@@ -1,4 +1,4 @@
-// Licensed to the .NET Foundation under one or more agreements.
+// Licensed to the .NET Foundation under one or more agreements.
 // The .NET Foundation licenses this file to you under the MIT license.
 
 using System.Threading;
@@ -98,8 +98,14 @@ static async ValueTask<RazorCSharpDocument> GetCSharpDocumentAsync(RemoteDocumen
             codeLens.Data = originalData;
         }
 
-        var generatedDocument = await snapshot.TryGetGeneratedDocumentAsync(declarationDocument: true, cancellationToken).ConfigureAwait(false)
-            ?? await snapshot.GetGeneratedDocumentAsync(declarationDocument: false, cancellationToken).ConfigureAwait(false);
+        // CodeLens shows information about fields, properties, methods etc. which, for components, all appear in a declaration document. For legacy documents
+        // there is no declaration document, so those things appear in the implementation document. For simplicity, we'll just attempt to resolve from the declaration
+        // document and fallback to the implementation document when it's null.
+        var generatedDocument = await snapshot.TryGetGeneratedDocumentAsync(declarationDocument: true, cancellationToken).ConfigureAwait(false);
+        if (generatedDocument is null)
+        {
+            generatedDocument = await snapshot.GetGeneratedDocumentAsync(declarationDocument: false, cancellationToken).ConfigureAwait(false);
+        }
 
         return await CodeLensResolveHandler.ResolveCodeLensAsync(codeLens, generatedDocument, cancellationToken).ConfigureAwait(false);
     }

src/Razor/src/Razor/src/Microsoft.CodeAnalysis.Remote.Razor/DocumentMapping/RazorEditService_WorkspaceEdit.cs

@@ -4,7 +4,6 @@
 using System;
 using System.Collections.Generic;
 using System.Collections.Immutable;
-using System.Runtime.InteropServices;
 using System.Threading;
 using System.Threading.Tasks;
 using Microsoft.AspNetCore.Razor.Language;
@@ -25,32 +24,85 @@ public async Task MapWorkspaceEditAsync(IDocumentSnapshot contextDocumentSnapsho
             throw new InvalidOperationException("RemoteRazorEditService can only be used with RemoteDocumentSnapshot instances.");
         }
 
+        // Collect both workspace edit shapes into TextDocumentEdits so URI coalescing and duplicate
+        // edit handling run once across the whole edit.
+        using var builder = new PooledArrayBuilder<SumType<TextDocumentEdit, CreateFile, RenameFile, DeleteFile>>();
+
         if (workspaceEdit.DocumentChanges is not null)
         {
-            using var builder = new PooledArrayBuilder<SumType<TextDocumentEdit, CreateFile, RenameFile, DeleteFile>>();
             foreach (var edit in workspaceEdit.EnumerateEdits())
             {
                 if (edit.TryGetFirst(out var textDocumentEdit))
                 {
                     await MapTextDocumentEditAsync(originSnapshot, textDocumentEdit, cancellationToken).ConfigureAwait(false);
-                    if (textDocumentEdit.Edits.Length == 0)
-                    {
-                        continue;
-                    }
                 }
 
                 builder.Add(edit);
             }
-
-            workspaceEdit.DocumentChanges = builder.ToArrayAndClear();
         }
 
         if (workspaceEdit.Changes is { } changeMap)
         {
-            workspaceEdit.Changes = await MapDocumentEditsAsync(originSnapshot, changeMap, cancellationToken).ConfigureAwait(false);
+            builder.AddRange(await MapDocumentEditsAsync(originSnapshot, changeMap, cancellationToken).ConfigureAwait(false));
+        }
+
+        var normalizedDocumentChanges = NormalizeDocumentChanges(builder.ToArrayAndClear());
+        if (workspaceEdit.DocumentChanges is not null)
+        {
+            workspaceEdit.DocumentChanges = normalizedDocumentChanges;
+        }
+        else
+        {
+            workspaceEdit.Changes = ConvertToChangeMap(normalizedDocumentChanges);
         }
     }
 
+    private static SumType<TextDocumentEdit, CreateFile, RenameFile, DeleteFile>[] NormalizeDocumentChanges(
+        SumType<TextDocumentEdit, CreateFile, RenameFile, DeleteFile>[] documentChanges)
+    {
+        // Multiple generated C# documents can map edits back to the same Razor document.
+        // Keep one TextDocumentEdit per URI so duplicate mapped edits are applied against the same source text.
+        using var _ = DictionaryPool<DocumentUri, TextDocumentEdit>.GetPooledObject(out var textDocumentEditsByUri);
+        using var builder = new PooledArrayBuilder<SumType<TextDocumentEdit, CreateFile, RenameFile, DeleteFile>>(documentChanges.Length);
+
+        foreach (var documentChange in documentChanges)
+        {
+            if (!documentChange.TryGetFirst(out var textDocumentEdit))
+            {
+                builder.Add(documentChange);
+                continue;
+            }
+
+            if (textDocumentEdit.Edits.Length == 0)
+            {
+                continue;
+            }
+
+            var uri = textDocumentEdit.TextDocument.DocumentUri;
+            if (!textDocumentEditsByUri.TryGetValue(uri, out var existingTextDocumentEdit))
+            {
+                textDocumentEditsByUri.Add(uri, textDocumentEdit);
+                builder.Add(textDocumentEdit);
+                continue;
+            }
+
+            existingTextDocumentEdit.Edits = [.. existingTextDocumentEdit.Edits, .. textDocumentEdit.Edits];
+        }
+
+        // After coalescing by URI, collapse exact duplicate edits that can be produced by the
+        // implementation and declaration generated documents mapping to the same Razor span.
+        var normalizedDocumentChanges = builder.ToArrayAndClear();
+        foreach (var documentChange in normalizedDocumentChanges)
+        {
+            if (documentChange.TryGetFirst(out var textDocumentEdit))
+            {
+                DeduplicateTextDocumentEdit(textDocumentEdit);
+            }
+        }
+
+        return normalizedDocumentChanges;
+    }
+
     private async Task MapTextDocumentEditAsync(RemoteDocumentSnapshot contextDocumentSnapshot, TextDocumentEdit entry, CancellationToken cancellationToken)
     {
         var generatedDocumentUri = entry.TextDocument.DocumentUri.GetRequiredSystemUri();
@@ -104,9 +156,11 @@ private async Task MapTextDocumentEditAsync(RemoteDocumentSnapshot contextDocume
         entry.Edits = mappedEdits.SelectAsPlainArray(static e => new SumType<TextEdit, AnnotatedTextEdit>(e));
     }
 
-    private async Task<Dictionary<string, TextEdit[]>> MapDocumentEditsAsync(RemoteDocumentSnapshot contextDocumentSnapshot, Dictionary<string, TextEdit[]> changes, CancellationToken cancellationToken)
+    private async Task<SumType<TextDocumentEdit, CreateFile, RenameFile, DeleteFile>[]> MapDocumentEditsAsync(RemoteDocumentSnapshot contextDocumentSnapshot, Dictionary<string, TextEdit[]> changes, CancellationToken cancellationToken)
     {
-        var mappedChanges = new Dictionary<string, TextEdit[]>(capacity: changes.Count);
+        // Map legacy Changes into TextDocumentEdits so MapWorkspaceEditAsync can normalize both shapes together.
+        using var builder = new PooledArrayBuilder<SumType<TextDocumentEdit, CreateFile, RenameFile, DeleteFile>>(changes.Count);
+        var solution = contextDocumentSnapshot.TextDocument.Project.Solution;
 
         foreach (var (uriString, edits) in changes)
         {
@@ -116,17 +170,17 @@ private async Task<Dictionary<string, TextEdit[]>> MapDocumentEditsAsync(RemoteD
             if (_filePathService.IsVirtualHtmlFile(generatedDocumentUri))
             {
                 var razorUri = _filePathService.GetRazorDocumentUri(generatedDocumentUri);
-                mappedChanges[razorUri.AbsoluteUri] = edits;
+                builder.Add(CreateTextDocumentEdit(razorUri.CreateDocumentUriFromSystemUri(), edits.AsSpan()));
+                continue;
             }
 
             // Check if the edit is actually for a generated document, because if not we don't need to do anything
             if (!_filePathService.IsVirtualCSharpFile(generatedDocumentUri))
             {
-                mappedChanges[uriString] = edits;
+                builder.Add(CreateTextDocumentEdit(new(uriString), edits.AsSpan()));
                 continue;
             }
 
-            var solution = contextDocumentSnapshot.TextDocument.Project.Solution;
             var razorDocument = await _snapshotManager.TryGetRazorDocumentAsync(solution, generatedDocumentUri, cancellationToken).ConfigureAwait(false);
             if (razorDocument is null)
             {
@@ -147,10 +201,48 @@ private async Task<Dictionary<string, TextEdit[]>> MapDocumentEditsAsync(RemoteD
                 continue;
             }
 
-            mappedChanges[razorDocument.CreateSystemUri().AbsoluteUri] = ImmutableCollectionsMarshal.AsArray(mappedEdits)!;
+            builder.Add(CreateTextDocumentEdit(razorDocument.GetURI(), mappedEdits.AsSpan()));
         }
 
-        return mappedChanges;
+        return builder.ToArrayAndClear();
+    }
+
+    private static TextDocumentEdit CreateTextDocumentEdit(DocumentUri documentUri, ReadOnlySpan<TextEdit> edits)
+    {
+        var textEdits = new SumType<TextEdit, AnnotatedTextEdit>[edits.Length];
+        for (var i = 0; i < edits.Length; i++)
+        {
+            textEdits[i] = edits[i];
+        }
+
+        return new TextDocumentEdit
+        {
+            TextDocument = new OptionalVersionedTextDocumentIdentifier { DocumentUri = documentUri },
+            Edits = textEdits
+        };
+    }
+
+    private static Dictionary<string, TextEdit[]> ConvertToChangeMap(SumType<TextDocumentEdit, CreateFile, RenameFile, DeleteFile>[] documentChanges)
+    {
+        var changes = new Dictionary<string, TextEdit[]>(capacity: documentChanges.Length);
+
+        foreach (var documentChange in documentChanges)
+        {
+            if (documentChange.TryGetFirst(out var textDocumentEdit))
+            {
+                var textEdits = new TextEdit[textDocumentEdit.Edits.Length];
+                for (var i = 0; i < textDocumentEdit.Edits.Length; i++)
+                {
+                    textEdits[i] = (TextEdit)textDocumentEdit.Edits[i];
+                }
+
+                changes.Add(
+                    textDocumentEdit.TextDocument.DocumentUri.GetRequiredSystemUri().AbsoluteUri,
+                    textEdits);
+            }
+        }
+
+        return changes;
     }
 
     private async Task<ImmutableArray<TextEdit>> GetMappedTextEditsAsync(RemoteDocumentSnapshot snapshot, RazorCSharpDocument csharpDocument, TextEdit[] edits, CancellationToken cancellationToken)
@@ -162,4 +254,35 @@ private async Task<ImmutableArray<TextEdit>> GetMappedTextEditsAsync(RemoteDocum
 
         return mappedEdits.SelectAsArray(razorSourceText.GetTextEdit);
     }
+
+    private static void DeduplicateTextDocumentEdit(TextDocumentEdit entry)
+    {
+        var edits = entry.Edits;
+
+        if (edits.Length <= 1)
+        {
+            return;
+        }
+
+        using var _ = HashSetPool<(int StartLine, int StartCharacter, int EndLine, int EndCharacter, string? NewText)>.GetPooledObject(out var seenEdits);
+        using var builder = new PooledArrayBuilder<SumType<TextEdit, AnnotatedTextEdit>>(edits.Length);
+
+        foreach (var edit in edits)
+        {
+            var textEdit = (TextEdit)edit;
+            var range = textEdit.Range;
+            if (seenEdits.Add((range.Start.Line, range.Start.Character, range.End.Line, range.End.Character, textEdit.NewText)))
+            {
+                builder.Add(edit);
+            }
+        }
+
+        if (builder.Count == edits.Length)
+        {
+            // No duplicates, return original array to avoid unnecessary allocations.
+            return;
+        }
+
+        entry.Edits = builder.ToArrayAndClear();
+    }
 }

src/Razor/src/Razor/src/Microsoft.CodeAnalysis.Remote.Razor/DocumentSymbols/RemoteDocumentSymbolService.cs

@@ -32,8 +32,17 @@ protected override IRemoteDocumentSymbolService CreateService(in ServiceArgs arg
     private async ValueTask<SumType<DocumentSymbol[], SymbolInformation[]>?> GetDocumentSymbolsAsync(RemoteDocumentContext context, bool useHierarchicalSymbols, CancellationToken cancellationToken)
     {
         var codeDocument = await context.GetCodeDocumentAsync(cancellationToken).ConfigureAwait(false);
-        var csharpDocument = codeDocument.GetCSharpDocument(declarationDocument: true)
-            ?? codeDocument.GetRequiredCSharpDocument(declarationDocument: false);
+
+        // We only care about fields, properties, methods etc. in document symbols, and for components those will exist in the declaration document.
+        // For legacy documents, there is no declaration document, so we use the implementation document. An edge case is components that have no
+        // declarations, and therefore no declaration document, where we want to use the implementation document so we at least get the class name
+        // and have something to base our render method symbol off of. Therefore, for simplicity, we'll just attempt to get the declaration document,
+        // and fallback to impl if it doesn't exist.
+        var csharpDocument = codeDocument.GetCSharpDocument(declarationDocument: true);
+        if (csharpDocument is null)
+        {
+            csharpDocument = codeDocument.GetRequiredCSharpDocument(declarationDocument: false);
+        }
 
         var generatedDocument = await context.Snapshot
             .GetGeneratedDocumentAsync(csharpDocument.IsDeclarationDocument, cancellationToken)

src/Razor/src/Razor/src/Microsoft.CodeAnalysis.Remote.Razor/Rename/RemoteRenameService.cs

@@ -66,7 +66,7 @@ protected override IRemoteRenameService CreateService(in ServiceArgs args)
         var positionInfo = GetPositionInfo(codeDocument, hostDocumentIndex, preferCSharpOverHtml: true);
 
         var generatedDocument = await context.Snapshot
-            .GetGeneratedDocumentAsync(cancellationToken)
+            .GetGeneratedDocumentAsync(positionInfo.InDeclDocument, cancellationToken)
             .ConfigureAwait(false);
 
         var razorEdit = await _renameService
@@ -134,7 +134,7 @@ protected override IRemoteRenameService CreateService(in ServiceArgs args)
             return RemoteResponse<LspRange?>.CallHtml;
         }
 
-        var generatedDocument = await context.Snapshot.GetGeneratedDocumentAsync(cancellationToken).ConfigureAwait(false);
+        var generatedDocument = await context.Snapshot.GetGeneratedDocumentAsync(positionInfo.InDeclDocument, cancellationToken).ConfigureAwait(false);
 
         var csharpRange = await PrepareRenameHandler.GetRenameRangeAsync(generatedDocument, positionInfo.Position.ToLinePosition(), cancellationToken).ConfigureAwait(false);
 
@@ -143,7 +143,7 @@ protected override IRemoteRenameService CreateService(in ServiceArgs args)
             return RemoteResponse<LspRange?>.NoFurtherHandling;
         }
 
-        if (!DocumentMappingService.TryMapToRazorDocumentRange(codeDocument.GetRequiredImplCSharpDocument(), csharpRange, out var mappedRange))
+        if (!DocumentMappingService.TryMapToRazorDocumentRange(codeDocument.GetRequiredCSharpDocument(positionInfo.InDeclDocument), csharpRange, out var mappedRange))
         {
             return RemoteResponse<LspRange?>.NoFurtherHandling;
         }
@@ -209,7 +209,7 @@ protected override IRemoteRenameService CreateService(in ServiceArgs args)
                 continue;
             }
 
-            var documentEdit = await GetEditsAsync(documentContext, newFileName, cancellationToken).ConfigureAwait(false);
+            var documentEdit = await GetFileRenameEditAsync(documentContext, newFileName, cancellationToken).ConfigureAwait(false);
             response = response.Concat(documentEdit);
         }
 
@@ -221,14 +221,16 @@ protected override IRemoteRenameService CreateService(in ServiceArgs args)
         return response;
     }
 
-    private async Task<WorkspaceEdit?> GetEditsAsync(RemoteDocumentContext context, string newFileName, CancellationToken cancellationToken)
+    private async Task<WorkspaceEdit?> GetFileRenameEditAsync(RemoteDocumentContext context, string newFileName, CancellationToken cancellationToken)
     {
         if (!context.Snapshot.FileKind.IsComponent())
         {
             return null;
         }
 
-        var generatedDocument = await context.Snapshot.GetGeneratedDocumentAsync(cancellationToken).ConfigureAwait(false);
+        // We're renaming the class declaration of a generated C# class, which exists in both decl and impl documents, so we can just work from
+        // the impl document since that will always exist. Decl may not.
+        var generatedDocument = await context.Snapshot.GetGeneratedDocumentAsync(declarationDocument: false, cancellationToken).ConfigureAwait(false);
         var text = await generatedDocument.GetTextAsync(cancellationToken).ConfigureAwait(false);
         var tree = await generatedDocument.GetSyntaxRootAsync(cancellationToken).ConfigureAwait(false);
         var declaration = tree.AssumeNotNull().DescendantNodes().OfType<ClassDeclarationSyntax>().FirstOrDefault();

src/Razor/src/Razor/test/Microsoft.CodeAnalysis.Razor.CohostingShared.UnitTests/Endpoints/CohostPrepareRenameEndpointTest.cs

@@ -11,7 +11,7 @@ namespace Microsoft.VisualStudio.Razor.LanguageClient.Cohost;
 
 public class CohostPrepareRenameEndpointTest(ITestOutputHelper testOutputHelper) : CohostEndpointTestBase(testOutputHelper)
 {
-    [Fact(Skip = "PROTOTYPE(sonic): cohosting feature not yet decl/impl split aware; see PR #83887")]
+    [Fact]
     public Task CSharp_Method()
         => VerifyPrepareRenameAsync(
             input: """