Hightlight of unknown items

Author: Mikhail Arkhipov

src/LanguageServer/Impl/Sources/DocumentHighlightSource.cs

@@ -14,20 +14,17 @@
 // permissions and limitations under the License.
 
 using System;
-using System.Collections.Generic;
 using System.IO;
 using System.Linq;
 using System.Threading;
 using System.Threading.Tasks;
 using Microsoft.Python.Analysis;
-using Microsoft.Python.Analysis.Documents;
-using Microsoft.Python.Analysis.Modules;
-using Microsoft.Python.Analysis.Types;
 using Microsoft.Python.Core;
-using Microsoft.Python.Core.IO;
 using Microsoft.Python.Core.Text;
 using Microsoft.Python.LanguageServer.Documents;
 using Microsoft.Python.LanguageServer.Protocol;
+using Microsoft.Python.Parsing;
+using Microsoft.Python.Parsing.Ast;
 
 namespace Microsoft.Python.LanguageServer.Sources {
     internal sealed class DocumentHighlightSource {
@@ -40,25 +37,51 @@ public DocumentHighlightSource(IServiceContainer services) {
 
         public async Task<DocumentHighlight[]> DocumentHighlightAsync(Uri uri, SourceLocation location, CancellationToken cancellationToken = default) {
             if (uri == null) {
-                return Array.Empty<DocumentHighlight>();
+                return null;
             }
 
             var analysis = await Document.GetAnalysisAsync(uri, _services, DocumentHighlightAnalysisTimeout, cancellationToken);
             var definitionSource = new DefinitionSource(_services);
 
             var definition = definitionSource.FindDefinition(analysis, location, out var definingMember);
             if (definition == null || definingMember == null) {
-                return Array.Empty<DocumentHighlight>();
+                return FromTokens(analysis, location);
             }
 
             var rootDefinition = definingMember.GetRootDefinition();
 
             var result = rootDefinition.References
                 .Where(r => r.DocumentUri.Equals(uri))
-                .Select((r, i) => new DocumentHighlight { kind = (i == 0) ? DocumentHighlightKind.Write : DocumentHighlightKind.Read, range = r.Span })
+                .Select((r, i) => new DocumentHighlight {
+                    kind = i == 0 ? DocumentHighlightKind.Write : DocumentHighlightKind.Read, range = r.Span
+                })
                 .ToArray();
 
             return result;
         }
+
+        private static DocumentHighlight[] FromTokens(IDocumentAnalysis analysis, SourceLocation location) {
+            var position = analysis.Ast.LocationToIndex(location);
+            var content = analysis.Document.Content;
+
+            var tokenizer = new Tokenizer(analysis.Document.Interpreter.LanguageVersion);
+            tokenizer.Initialize(null, new StringReader(content), SourceLocation.MinValue);
+            var tokens = tokenizer.ReadTokens(content.Length);
+            
+            var t = tokens.FirstOrDefault(x => x.SourceSpan.Start.Index <= position && position < x.SourceSpan.End.Index);
+            if (t.Category != TokenCategory.None) {
+                var length = t.SourceSpan.End.Index - t.SourceSpan.Start.Index;
+                return tokens
+                    .Where(x =>
+                        x.SourceSpan.End.Index - x.SourceSpan.Start.Index == length &&
+                        string.Compare(content, x.SourceSpan.Start.Index, content, t.SourceSpan.Start.Index, length) == 0)
+                    .Select(s => new DocumentHighlight {
+                        kind = DocumentHighlightKind.Text,
+                        range = s.SourceSpan
+                    }).ToArray();
+            }
+
+            return null;
+        }
     }
 }