Merge pull request #19873 from CyrusNajmabadi/docHighlightsCrash

Fix crash in DocHighlights when delegating to a different language.

Author: CyrusNajmabadi

src/EditorFeatures/Test2/EditorServicesTest2.vbproj

@@ -249,6 +249,7 @@
     <Compile Include="IntelliSense\CompletionRulesTests.vb" />
     <Compile Include="LanguageServices\SyntaxFactsServiceTests.vb" />
     <Compile Include="NamingStyles\NamingStylesTests.vb" />
+    <Compile Include="ReferenceHighlighting\DocumentHighlightsServiceTests.vb" />
     <Compile Include="Simplification\BlockSimplificationTests.vb" />
     <Compile Include="Simplification\SimplifierAPITests.vb" />
     <Compile Include="CodeFixes\CodeFixServiceTests.vb" />

src/EditorFeatures/Test2/ReferenceHighlighting/DocumentHighlightsServiceTests.vb

@@ -0,0 +1,45 @@
+Imports System.Collections.Immutable
+Imports System.Threading
+Imports System.Threading.Tasks
+Imports Microsoft.CodeAnalysis.Editor.UnitTests.Workspaces
+
+Namespace Microsoft.CodeAnalysis.Editor.UnitTests.ReferenceHighlighting
+    Public Class DocumentHighlightsServiceTests
+
+        <WorkItem(441151, "https://devdiv.visualstudio.com/DevDiv/_workitems/edit/441151")>
+        <Fact, Trait(Traits.Feature, Traits.Features.ReferenceHighlighting)>
+        Public Async Function TestMultipleLanguagesPassedToAPI() As Task
+            Dim workspaceElement =
+                <Workspace>
+                    <Project Language="C#" CommonReferences="true">
+                        <Document>
+                            class C
+                            {
+                                $$string Blah()
+                                {
+                                    return null;
+                                }
+                            }
+                        </Document>
+                    </Project>
+                    <Project Language="Visual Basic">
+                        <Document>
+                            Class VBClass
+                            End Class
+                        </Document>
+                    </Project>
+                </Workspace>
+            Using workspace = TestWorkspace.Create(workspaceElement)
+                Dim position = workspace.DocumentWithCursor.CursorPosition.Value
+
+                Dim solution = workspace.CurrentSolution
+                Dim csharpDocument = solution.Projects.Single(Function(p) p.Language = LanguageNames.CSharp).Documents.Single()
+                Dim vbDocument = solution.Projects.Single(Function(p) p.Language = LanguageNames.VisualBasic).Documents.Single()
+
+                Dim service = csharpDocument.GetLanguageService(Of Microsoft.CodeAnalysis.DocumentHighlighting.IDocumentHighlightsService)
+                Await service.GetDocumentHighlightsAsync(
+                    csharpDocument, position, ImmutableHashSet.Create(csharpDocument, vbDocument), CancellationToken.None)
+            End Using
+        End Function
+    End Class
+End Namespace

src/Features/Core/Portable/DocumentHighlighting/AbstractDocumentHighlightsService.cs

@@ -76,8 +76,8 @@ private async Task<ImmutableArray<DocumentHighlights>> GetDocumentHighlightsInCu
 
             // Get unique tags for referenced symbols
             return await GetTagsForReferencedSymbolAsync(
-                new SymbolAndProjectId(symbol, document.Project.Id), documentsToSearch,
-                solution, cancellationToken).ConfigureAwait(false);
+                new SymbolAndProjectId(symbol, document.Project.Id),
+                document, documentsToSearch, cancellationToken).ConfigureAwait(false);
         }
 
         private static async Task<ISymbol> GetSymbolToSearchAsync(Document document, int position, SemanticModel semanticModel, ISymbol symbol, CancellationToken cancellationToken)
@@ -95,8 +95,8 @@ private static async Task<ISymbol> GetSymbolToSearchAsync(Document document, int
 
         private async Task<ImmutableArray<DocumentHighlights>> GetTagsForReferencedSymbolAsync(
             SymbolAndProjectId symbolAndProjectId,
+            Document document,
             IImmutableSet<Document> documentsToSearch,
-            Solution solution,
             CancellationToken cancellationToken)
         {
             var symbol = symbolAndProjectId.Symbol;
@@ -106,11 +106,11 @@ private async Task<ImmutableArray<DocumentHighlights>> GetTagsForReferencedSymbo
                 var progress = new StreamingProgressCollector(
                     StreamingFindReferencesProgress.Instance);
                 await SymbolFinder.FindReferencesAsync(
-                    symbolAndProjectId, solution, progress,
+                    symbolAndProjectId, document.Project.Solution, progress,
                     documentsToSearch, cancellationToken).ConfigureAwait(false);
 
                 return await FilterAndCreateSpansAsync(
-                    progress.GetReferencedSymbols(), solution, documentsToSearch, 
+                    progress.GetReferencedSymbols(), document, documentsToSearch,
                     symbol, cancellationToken).ConfigureAwait(false);
             }
 
@@ -142,10 +142,12 @@ private static bool ShouldConsiderSymbol(ISymbol symbol)
         }
 
         private async Task<ImmutableArray<DocumentHighlights>> FilterAndCreateSpansAsync(
-            IEnumerable<ReferencedSymbol> references, Solution solution, 
-            IImmutableSet<Document> documentsToSearch, ISymbol symbol, 
+            IEnumerable<ReferencedSymbol> references, Document startingDocument,
+            IImmutableSet<Document> documentsToSearch, ISymbol symbol,
             CancellationToken cancellationToken)
         {
+            var solution = startingDocument.Project.Solution;
+
             references = references.FilterToItemsToShow();
             references = references.FilterNonMatchingMethodNames(solution, symbol);
             references = references.FilterToAliasMatches(symbol as IAliasSymbol);
@@ -157,13 +159,20 @@ private async Task<ImmutableArray<DocumentHighlights>> FilterAndCreateSpansAsync
 
             var additionalReferences = new List<Location>();
 
-            foreach (var document in documentsToSearch)
+            foreach (var currentDocument in documentsToSearch)
             {
-                additionalReferences.AddRange(await GetAdditionalReferencesAsync(document, symbol, cancellationToken).ConfigureAwait(false));
+                // 'documentsToSearch' may contain documents from languages other than our own
+                // (for example cshtml files when we're searching the cs document).  Since we're
+                // delegating to a virtual method for this language type, we have to make sure
+                // we only process the document if it's also our language.
+                if (currentDocument.Project.Language == startingDocument.Project.Language)
+                {
+                    additionalReferences.AddRange(await GetAdditionalReferencesAsync(currentDocument, symbol, cancellationToken).ConfigureAwait(false));
+                }
             }
 
             return await CreateSpansAsync(
-                solution, symbol, references, additionalReferences, 
+                solution, symbol, references, additionalReferences,
                 documentsToSearch, cancellationToken).ConfigureAwait(false);
         }