Put partial references underneath definitions.

src/EditorFeatures/Core/Implementation/FindReferences/FindReferencesCommandHandler.cs

@@ -65,11 +65,11 @@ public void ExecuteCommand(FindReferencesCommandArgs args, Action nextHandler)
                 // See if we're running on a host that can provide streaming results.
                 // We'll both need a FAR service that can stream results to us, and 
                 // a presenter that can accept streamed results.
-                if (streamingService != null && streamingPresenter != null)
-                {
-                    StreamingFindReferences(document, streamingService, streamingPresenter, caretPosition);
-                    return;
-                }
+                //if (streamingService != null && streamingPresenter != null)
+                //{
+                //    StreamingFindReferences(document, streamingService, streamingPresenter, caretPosition);
+                //    return;
+                //}
 
                 // Otherwise, either the language doesn't support streaming results,
                 // or the host has no way to present results in a sreaming manner.

src/Features/Core/Portable/FindReferences/DefinitionItem.cs

@@ -23,10 +23,15 @@ internal sealed class DefinitionItem
         public ImmutableArray<TaggedText> DisplayParts { get; }
 
         /// <summary>
-        /// The locations to present in the UI.  A definition may have multiple locations for cases
-        /// like partial types/members.
+        /// The main locations to present in the UI.
         /// </summary>
-        public ImmutableArray<DefinitionLocation> Locations { get; }
+        public DefinitionLocation MainLocation { get; }
+
+        /// <summary>
+        /// Additional locations to present in the UI.  A definition may have multiple locations 
+        /// for cases like partial types/members.
+        /// </summary>
+        public ImmutableArray<DocumentLocation> AdditionalLocations { get; }
 
         /// <summary>
         /// Whether or not this definition should be presented if we never found any references to
@@ -43,12 +48,14 @@ internal sealed class DefinitionItem
         public DefinitionItem(
             ImmutableArray<string> tags,
             ImmutableArray<TaggedText> displayParts,
-            ImmutableArray<DefinitionLocation> locations,
+            DefinitionLocation mainLocation,
+            ImmutableArray<DocumentLocation> additionalLocations,
             bool displayIfNoReferences)
         {
             Tags = tags;
             DisplayParts = displayParts;
-            Locations = locations;
+            MainLocation = mainLocation;
+            AdditionalLocations = additionalLocations;
             DisplayIfNoReferences = displayIfNoReferences;
         }
     }

src/Features/Core/Portable/FindReferences/IDefinitionsAndReferencesFactory.cs

@@ -146,20 +146,22 @@ internal static class DefinitionItemExtensions
             this ISymbol definition,
             Solution solution)
         {
-            var definitionLocations = ConvertDefinitionLocations(solution, definition);
+            var locations = ConvertDefinitionLocations(solution, definition);
             var displayParts = definition.ToDisplayParts(s_definitionDisplayFormat).ToTaggedText();
 
             return new DefinitionItem(
                 GlyphTags.GetTags(definition.GetGlyph()),
                 displayParts,
-                definitionLocations,
+                locations.Item1,
+                locations.Item2,
                 definition.ShouldShowWithNoReferenceLocations());
         }
 
-        private static ImmutableArray<DefinitionLocation> ConvertDefinitionLocations(
+        private static ValueTuple<DefinitionLocation, ImmutableArray<DocumentLocation>> ConvertDefinitionLocations(
             Solution solution, ISymbol definition)
         {
-            var result = ImmutableArray.CreateBuilder<DefinitionLocation>();
+            var additionalLocations = ImmutableArray.CreateBuilder<DocumentLocation>();
+            DefinitionLocation mainLocation = null;
 
             // If it's a namespace, don't create any normal lcoation.  Namespaces
             // come from many different sources, but we'll only show a single 
@@ -170,29 +172,36 @@ internal static class DefinitionItemExtensions
                 {
                     if (location.IsInMetadata)
                     {
-                        result.Add(DefinitionLocation.CreateSymbolLocation(solution, definition));
+                        mainLocation = DefinitionLocation.CreateSymbolLocation(solution, definition);
                     }
                     else if (location.IsVisibleSourceLocation())
                     {
                         var document = solution.GetDocument(location.SourceTree);
                         if (document != null)
                         {
-                            result.Add(DefinitionLocation.CreateDocumentLocation(
-                                new DocumentLocation(document, location.SourceSpan)));
+                            var documentLocation = new DocumentLocation(document, location.SourceSpan);
+                            if (mainLocation == null)
+                            {
+                                mainLocation = DefinitionLocation.CreateDocumentLocation(documentLocation);
+                            }
+                            else
+                            {
+                                additionalLocations.Add(documentLocation);
+                            }
                         }
                     }
                 }
             }
 
-            if (result.Count == 0)
+            if (mainLocation == null)
             {
                 // If we got no definition locations, then create a sentinel one
                 // that we can display but which will not allow navigation.
-                result.Add(DefinitionLocation.CreateNonNavigatingLocation(
-                    DefinitionLocation.GetOriginationParts(definition)));
+                mainLocation = DefinitionLocation.CreateNonNavigatingLocation(
+                    DefinitionLocation.GetOriginationParts(definition));
             }
 
-            return result.ToImmutable();
+            return ValueTuple.Create(mainLocation, additionalLocations.ToImmutable());
         }
 
         public static SourceReferenceItem TryCreateSourceReferenceItem(

src/VisualStudio/Core/Def.Next/FindReferences/StreamingFindReferencesPresenter.ReferenceEntry.cs

@@ -32,7 +32,7 @@ private class ReferenceEntry
             private readonly VisualStudioWorkspaceImpl _workspace;
 
             private readonly RoslynDefinitionBucket _definitionBucket;
-            private readonly SourceReferenceItem _sourceReferenceItem;
+            private readonly DocumentLocation _documentLocation;
 
             private readonly object _boxedProjectGuid;
             private readonly SourceText _sourceText;
@@ -42,7 +42,7 @@ private class ReferenceEntry
                 TableDataSourceFindReferencesContext context,
                 VisualStudioWorkspaceImpl workspace,
                 RoslynDefinitionBucket definitionBucket,
-                SourceReferenceItem sourceReferenceItem,
+                DocumentLocation documentLocation,
                 Guid projectGuid,
                 SourceText sourceText,
                 TaggedTextAndHighlightSpan taggedLineParts)
@@ -51,7 +51,7 @@ private class ReferenceEntry
 
                 _workspace = workspace;
                 _definitionBucket = definitionBucket;
-                _sourceReferenceItem = sourceReferenceItem;
+                _documentLocation = documentLocation;
 
                 _boxedProjectGuid = projectGuid;
                 _sourceText = sourceText;
@@ -66,9 +66,8 @@ public bool TryGetValue(string keyName, out object content)
                 return content != null;
             }
 
-            private DocumentLocation Location => _sourceReferenceItem.Location;
-            private Document Document => Location.Document;
-            private TextSpan SourceSpan => Location.SourceSpan;
+            private Document Document => _documentLocation.Document;
+            private TextSpan SourceSpan => _documentLocation.SourceSpan;
 
             private object GetValue(string keyName)
             {
@@ -93,7 +92,7 @@ private object GetValue(string keyName)
                 case StandardTableKeyNames.Column:
                     return _sourceText.Lines.GetLinePosition(SourceSpan.Start).Character;
                 case StandardTableKeyNames.ProjectName:
-                    return Location.Document.Project.Name;
+                    return Document.Project.Name;
                 case StandardTableKeyNames.ProjectGuid:
                     return _boxedProjectGuid;
 
@@ -217,14 +216,14 @@ private void SetHighlightSpansOnBuffer()
                 var key = PredefinedPreviewTaggerKeys.ReferenceHighlightingSpansKey;
                 textBuffer.Properties.RemoveProperty(key);
                 textBuffer.Properties.AddProperty(key, new NormalizedSnapshotSpanCollection(
-                    _sourceReferenceItem.Location.SourceSpan.ToSnapshotSpan(textBuffer.CurrentSnapshot)));
+                    SourceSpan.ToSnapshotSpan(textBuffer.CurrentSnapshot)));
             }
 
             private Span GetRegionSpanForReference()
             {
                 const int AdditionalLineCountPerSide = 3;
 
-                var referenceSpan = this._sourceReferenceItem.Location.SourceSpan;
+                var referenceSpan = this.SourceSpan;
                 var lineNumber = _sourceText.Lines.GetLineFromPosition(referenceSpan.Start).LineNumber;
                 var firstLineNumber = Math.Max(0, lineNumber - AdditionalLineCountPerSide);
                 var lastLineNumber = Math.Min(_sourceText.Lines.Count - 1, lineNumber + AdditionalLineCountPerSide);

src/VisualStudio/Core/Def.Next/FindReferences/StreamingFindReferencesPresenter.RoslynDefinitionBucket.cs

@@ -37,15 +37,7 @@ private class RoslynDefinitionBucket : DefinitionBucket, ISupportsNavigation
 
             public bool TryNavigateTo()
             {
-                foreach (var location in DefinitionItem.Locations)
-                {
-                    if (location.TryNavigateTo())
-                    {
-                        return true;
-                    }
-                }
-
-                return false;
+                return DefinitionItem.MainLocation.TryNavigateTo();
             }
 
             public override bool TryGetValue(string key, out object content)

src/VisualStudio/Core/Def.Next/FindReferences/StreamingFindReferencesPresenter.TableDataSourceFindReferencesContext.cs

@@ -149,7 +149,7 @@ public async override void OnReferenceFound(SourceReferenceItem reference)
                     // know so that it doesn't verify results until this completes.
                     using (var token = Presenter._asyncListener.BeginAsyncOperation(nameof(OnReferenceFound)))
                     {
-                        await OnReferenceFoundAsync(reference).ConfigureAwait(false);
+                        await OnReferenceFoundAsync(reference.Definition, reference.Location).ConfigureAwait(false);
                     }
                 }
                 catch (Exception e) when (FatalError.ReportWithoutCrashUnlessCanceled(e))
@@ -157,14 +157,15 @@ public async override void OnReferenceFound(SourceReferenceItem reference)
                 }
             }
 
-            private async Task OnReferenceFoundAsync(SourceReferenceItem referenceItem)
+            private async Task OnReferenceFoundAsync(
+                DefinitionItem definition, DocumentLocation referenceLocation)
             {
                 var cancellationToken = _cancellationTokenSource.Token;
                 cancellationToken.ThrowIfCancellationRequested();
 
                 // First find the bucket corresponding to our definition. If we can't find/create 
                 // one, then don't do anything for this reference.
-                var definitionBucket = GetOrCreateDefinitionBucket(referenceItem.Definition);
+                var definitionBucket = GetOrCreateDefinitionBucket(definition);
                 if (definitionBucket == null)
                 {
                     return;
@@ -174,7 +175,7 @@ private async Task OnReferenceFoundAsync(SourceReferenceItem referenceItem)
                 //var entryData = await this.TryCreateNavigableItemEntryData(
                 //    referenceItem, isDefinition: false, cancellationToken: cancellationToken).ConfigureAwait(false);
                 var referenceEntry = await this.CreateReferenceEntryAsync(
-                    definitionBucket, referenceItem, cancellationToken).ConfigureAwait(false);
+                    definitionBucket, referenceLocation, cancellationToken).ConfigureAwait(false);
                 if (referenceEntry == null)
                 {
                     return;
@@ -227,10 +228,9 @@ private ITextBuffer CreateNewBuffer(Document document, SourceText sourceText)
 
 
             private async Task<ReferenceEntry> CreateReferenceEntryAsync(
-                RoslynDefinitionBucket definitionBucket, SourceReferenceItem referenceItem, CancellationToken cancellationToken)
+                RoslynDefinitionBucket definitionBucket, DocumentLocation documentLocation, CancellationToken cancellationToken)
             {
-                var location = referenceItem.Location;
-                var document = location.Document;
+                var document = documentLocation.Document;
 
                 // The FAR system needs to know the guid for the project that a def/reference is 
                 // from.  So we only support this for documents from a VSWorkspace.
@@ -248,13 +248,13 @@ private ITextBuffer CreateNewBuffer(Document document, SourceText sourceText)
 
                 var sourceText = await document.GetTextAsync(cancellationToken).ConfigureAwait(false);
 
-                var referenceSpan = referenceItem.Location.SourceSpan;
+                var referenceSpan = documentLocation.SourceSpan;
                 var lineSpan = GetLineSpanForReference(sourceText, referenceSpan);
 
                 var taggedLineParts = await GetTaggedTextForReferenceAsync(document, referenceSpan, lineSpan, cancellationToken).ConfigureAwait(false);
 
                 return new ReferenceEntry(
-                    this, workspace, definitionBucket, referenceItem, 
+                    this, workspace, definitionBucket, documentLocation, 
                     projectGuid.Value, sourceText, taggedLineParts);
             }
 

src/VisualStudio/Core/Def/Implementation/FindReferences/VisualStudioDefinitionsAndReferencesFactory.cs

@@ -40,13 +40,14 @@ public VisualStudioDefinitionsAndReferencesFactory(SVsServiceProvider servicePro
             }
 
             var displayParts = GetDisplayParts(filePath, lineNumber, charOffset);
-            var definitionLocation = new ExternalDefinitionLocation(
+            var mainLocation = new ExternalDefinitionLocation(
                 _serviceProvider, filePath, lineNumber, charOffset);
 
             return new DefinitionItem(
                 GlyphTags.GetTags(definition.GetGlyph()),
                 displayParts,
-                ImmutableArray.Create<DefinitionLocation>(definitionLocation),
+                mainLocation,
+                additionalLocations: ImmutableArray<DocumentLocation>.Empty,
                 displayIfNoReferences: true);
         }
 

src/VisualStudio/Core/Def/Implementation/Library/FindResults/LibraryManager_FindReferences.cs

@@ -7,6 +7,7 @@
 using Microsoft.CodeAnalysis.FindReferences;
 using Microsoft.VisualStudio.LanguageServices.Implementation.Utilities;
 using Roslyn.Utilities;
+using Microsoft.CodeAnalysis.Editor.Shared.Extensions;
 
 namespace Microsoft.VisualStudio.LanguageServices.Implementation.Library.FindResults
 {
@@ -24,19 +25,21 @@ public void PresentDefinitionsAndReferences(DefinitionsAndReferences definitions
         internal IList<AbstractTreeItem> CreateFindReferencesItems(
             DefinitionsAndReferences definitionsAndReferences)
         {
-            var documents = definitionsAndReferences.References.Select(r => r.Location.Document)
-                                                               .WhereNotNull()
-                                                               .ToSet();
-            var commonPathElements = CountCommonPathElements(documents);
+            var definitionDocuments =
+                definitionsAndReferences.Definitions.SelectMany(d => d.AdditionalLocations)
+                                        .Select(loc => loc.Document);
+            var referenceDocuments =
+                definitionsAndReferences.References.Select(r => r.Location.Document);
 
-            var query = from d in definitionsAndReferences.Definitions
-                        from i in CreateDefinitionItems(d, definitionsAndReferences, commonPathElements)
-                        select (AbstractTreeItem)i;
+            var documents = definitionDocuments.Concat(referenceDocuments).WhereNotNull().ToSet();
+            var commonPathElements = CountCommonPathElements(documents);
 
-            return query.ToList();
+            return definitionsAndReferences.Definitions
+                .Select(d => CreateDefinitionItem(d, definitionsAndReferences, commonPathElements))
+                .ToList<AbstractTreeItem>();
         }
 
-        private ImmutableArray<DefinitionTreeItem> CreateDefinitionItems(
+        private DefinitionTreeItem CreateDefinitionItem(
             DefinitionItem definitionItem,
             DefinitionsAndReferences definitionsAndReferences,
             int commonPathElements)
@@ -44,29 +47,28 @@ public void PresentDefinitionsAndReferences(DefinitionsAndReferences definitions
             var referenceItems = CreateReferenceItems(
                 definitionItem, definitionsAndReferences, commonPathElements);
 
-            return ConvertToDefinitionTreeItems(definitionItem, referenceItems);
+            return ConvertToDefinitionTreeItem(definitionItem, referenceItems, commonPathElements);
         }
 
-        private ImmutableArray<DefinitionTreeItem> ConvertToDefinitionTreeItems(
+        private DefinitionTreeItem ConvertToDefinitionTreeItem(
             DefinitionItem definitionItem,
-            ImmutableArray<SourceReferenceTreeItem> referenceItems)
+            ImmutableArray<SourceReferenceTreeItem> referenceItems,
+            int commonPathElements)
         {
-            var result = ImmutableArray.CreateBuilder<DefinitionTreeItem>();
+            var finalReferenceItems = ImmutableArray.CreateBuilder<SourceReferenceTreeItem>();
 
-            for (int i = 0, n = definitionItem.Locations.Length; i < n; i++)
+            foreach (var additionalLocation in definitionItem.AdditionalLocations)
             {
-                var location = definitionItem.Locations[i];
-
-                // Each definition item may end up as several top nodes (because of partials).
-                // Add the references to the last item actually in the list.
-                var childItems = i == n - 1
-                    ? referenceItems
-                    : ImmutableArray<SourceReferenceTreeItem>.Empty;
-
-                result.Add(new DefinitionTreeItem(definitionItem, location, childItems));
+                finalReferenceItems.Add(new SourceReferenceTreeItem(
+                    additionalLocation.Document,
+                    additionalLocation.SourceSpan,
+                    definitionItem.Tags.GetGlyph().GetGlyphIndex(),
+                    commonPathElements));
             }
 
-            return result.ToImmutable();
+            finalReferenceItems.AddRange(referenceItems);
+
+            return new DefinitionTreeItem(definitionItem, finalReferenceItems.ToImmutable());
         }
 
         private ImmutableArray<SourceReferenceTreeItem> CreateReferenceItems(

src/VisualStudio/Core/Def/Implementation/Library/FindResults/TreeItems/DefinitionTreeItem.cs

@@ -12,16 +12,13 @@ namespace Microsoft.VisualStudio.LanguageServices.Implementation.Library.FindRes
     internal class DefinitionTreeItem : AbstractTreeItem
     {
         private readonly DefinitionItem _definitionItem;
-        private readonly DefinitionLocation _definitionLocation;
 
         public DefinitionTreeItem(
             DefinitionItem definitionItem,
-            DefinitionLocation definitionLocation,
             ImmutableArray<SourceReferenceTreeItem> referenceItems)
             : base(definitionItem.Tags.GetGlyph().GetGlyphIndex())
         {
             _definitionItem = definitionItem;
-            _definitionLocation = definitionLocation;
 
             this.Children.AddRange(referenceItems);
             this.DisplayText = CreateDisplayText();
@@ -41,9 +38,9 @@ private string CreateDisplayText()
             // results that tell us about their definition location, but not any additional
             // reference.  We don't want to say '0' references in that case as that can
             // be misleading.
-            var hasOrigination = _definitionLocation.OriginationParts.Length > 0;
+            var hasOrigination = _definitionItem.MainLocation.OriginationParts.Length > 0;
             return hasOrigination
-                ? $"[{_definitionLocation.OriginationParts.JoinText()}] {displayString} ({referenceCountDisplay})"
+                ? $"[{_definitionItem.MainLocation.OriginationParts.JoinText()}] {displayString} ({referenceCountDisplay})"
                 : referenceCount > 0
                     ? $"{displayString} ({referenceCountDisplay})"
                     : displayString;
@@ -51,14 +48,14 @@ private string CreateDisplayText()
 
         public override int GoToSource()
         {
-            return _definitionLocation.TryNavigateTo()
+            return _definitionItem.MainLocation.TryNavigateTo()
                 ? VSConstants.S_OK
                 : VSConstants.E_FAIL;
         }
 
         public override bool CanGoToDefinition()
         {
-            return _definitionLocation.CanNavigateTo();
+            return _definitionItem.MainLocation.CanNavigateTo();
         }
     }
 }
\ No newline at end of file