Provide token-based highlight if references are not available

src/LanguageServer/Impl/Sources/DocumentHighlightSource.cs

@@ -17,21 +17,21 @@
 using System.Collections.Generic;
 using System.IO;
 using System.Linq;
+using System.Security.Cryptography;
+using System.Text;
 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;
 
 namespace Microsoft.Python.LanguageServer.Sources {
     internal sealed class DocumentHighlightSource {
-        private const int DocumentHighlightAnalysisTimeout = 10000;
+        private const int DocumentHighlightAnalysisTimeout = 1000;
+        private static TokenCache _tokenCache = new TokenCache(10, TimeSpan.FromMinutes(30));
         private readonly IServiceContainer _services;
 
         public DocumentHighlightSource(IServiceContainer services) {
@@ -48,17 +48,108 @@ public async Task<DocumentHighlight[]> DocumentHighlightAsync(Uri uri, SourceLoc
 
             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 tokens = _tokenCache.GetTokens(analysis.Document.Content, analysis.Document.Interpreter.LanguageVersion);
+            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 Array.Empty<DocumentHighlight>();
+        }
+    }
+
+
+    internal class TokenCache {
+        internal class Entry {
+            public DateTime AccessTime;
+            public WeakReference<IReadOnlyList<TokenInfo>> Tokens;
+        }
+
+        private readonly Dictionary<long, Entry> _cache = new Dictionary<long, Entry>();
+        private readonly int _maxEntries;
+        private readonly TimeSpan _expiration;
+
+        public TokenCache(int maxEntries, TimeSpan expiration) {
+            _maxEntries = maxEntries;
+            _expiration = expiration;
+        }
+
+        public IReadOnlyList<TokenInfo> GetTokens(string content, PythonLanguageVersion languageVersion) {
+            IReadOnlyList<TokenInfo> tokens;
+
+            var hash = GetHash(content);
+            if (_cache.TryGetValue(hash, out var entry)) {
+                if (entry.Tokens.TryGetTarget(out tokens)) {
+                    entry.AccessTime = DateTime.Now;
+                    return tokens;
+                }
+            }
+
+            var tokenizer = new Tokenizer(languageVersion);
+            using (var sr = new StringReader(content)) {
+                tokenizer.Initialize(null, sr, SourceLocation.MinValue);
+                tokens = tokenizer.ReadTokens(content.Length);
+                _cache[hash] = new Entry {
+                    AccessTime = DateTime.Now,
+                    Tokens = new WeakReference<IReadOnlyList<TokenInfo>>(tokens)
+                };
+            }
+
+            var byTime = _cache.OrderByDescending(kvp => (DateTime.Now - kvp.Value.AccessTime).TotalSeconds).ToArray();
+
+            var expired = byTime.TakeWhile(kvp => DateTime.Now - kvp.Value.AccessTime > _expiration).ToArray();
+            foreach (var e in expired) {
+                _cache.Remove(e.Key);
+            }
+
+            if (_cache.Count > _maxEntries) {
+                var (key, _) = byTime.FirstOrDefault();
+                if (key != default) {
+                    _cache.Remove(key);
+                }
+            }
+
+            return tokens;
+        }
+
+        // For tests
+        internal IEnumerable<(DateTime, IReadOnlyList<TokenInfo>)> Entries
+            => _cache.Values.Select(kvp => {
+                kvp.Tokens.TryGetTarget(out var t);
+                return (kvp.AccessTime, t);
+            });
+
+        internal static long GetHash(string content) {
+            using (var sha = SHA1.Create()) {
+                return BitConverter.ToInt64(sha.ComputeHash(Encoding.UTF32.GetBytes(content)));
+            }
+        }
     }
 }

src/LanguageServer/Test/DocumentHighlightTests.cs

@@ -14,16 +14,15 @@
 // permissions and limitations under the License.
 
 using System;
-using System.IO;
+using System.Collections.Generic;
+using System.Linq;
 using System.Threading.Tasks;
 using FluentAssertions;
-using Microsoft.Python.Analysis.Analyzer;
-using Microsoft.Python.Analysis.Documents;
 using Microsoft.Python.Core.Text;
 using Microsoft.Python.LanguageServer.Protocol;
 using Microsoft.Python.LanguageServer.Sources;
 using Microsoft.Python.LanguageServer.Tests.FluentAssertions;
-using Microsoft.Python.Parsing.Tests;
+using Microsoft.Python.Parsing;
 using Microsoft.VisualStudio.TestTools.UnitTesting;
 using TestUtilities;
 
@@ -76,10 +75,80 @@ def func(x):
 
         [TestMethod, Priority(0)]
         public async Task HighlightEmptyDocument() {
-            await GetAnalysisAsync(string.Empty);
+            var analysis = await GetAnalysisAsync(string.Empty);
             var dhs = new DocumentHighlightSource(Services);
-            var references = await dhs.DocumentHighlightAsync(null, new SourceLocation(1, 1));
-            references.Should().BeEmpty();
+            var highlights = await dhs.DocumentHighlightAsync(analysis.Document.Uri, new SourceLocation(1, 1));
+            highlights.Should().BeEmpty();
+        }
+
+        [TestMethod, Priority(0)]
+        public async Task HighlightNonReference() {
+            const string code = @"
+x = y = 0
+assert x == 1
+assert y != 3
+";
+            var analysis = await GetAnalysisAsync(code);
+            var dhs = new DocumentHighlightSource(Services);
+            var highlights = await dhs.DocumentHighlightAsync(analysis.Document.Uri, new SourceLocation(3, 5));
+
+            highlights.Should().HaveCount(2);
+            highlights[0].range.Should().Be(2, 0, 2, 6);
+            highlights[0].kind.Should().Be(DocumentHighlightKind.Text);
+            highlights[1].range.Should().Be(3, 0, 3, 6);
+            highlights[1].kind.Should().Be(DocumentHighlightKind.Text);
+        }
+
+        [TestMethod, Priority(0)]
+        public async Task HighlightUndefined() {
+            const string code = @"
+assert x == 1
+assert x != 3
+";
+            var analysis = await GetAnalysisAsync(code);
+            var dhs = new DocumentHighlightSource(Services);
+            var highlights = await dhs.DocumentHighlightAsync(analysis.Document.Uri, new SourceLocation(2, 8));
+
+            highlights.Should().HaveCount(2);
+            highlights[0].range.Should().Be(1, 7, 1, 8);
+            highlights[0].kind.Should().Be(DocumentHighlightKind.Text);
+            highlights[1].range.Should().Be(2, 7, 2, 8);
+            highlights[1].kind.Should().Be(DocumentHighlightKind.Text);
+        }
+
+        [TestMethod, Priority(0)]
+        public void TokenCacheEntriesLimit() {
+            var c = new TokenCache(3, TimeSpan.FromMinutes(1));
+            c.GetTokens("1", PythonLanguageVersion.V38);
+            c.GetTokens("12", PythonLanguageVersion.V38);
+            c.GetTokens("123", PythonLanguageVersion.V38);
+            c.GetTokens("1234", PythonLanguageVersion.V38);
+
+            (DateTime AccessTime, IReadOnlyList<TokenInfo> Tokens)[] e = c.Entries.ToArray();
+
+            e.Should().HaveCount(3);
+            var byTime = e.OrderBy(x => x.AccessTime).ToArray();
+            (byTime[0].Tokens[0].SourceSpan.End.Column - byTime[0].Tokens[0].SourceSpan.Start.Column).Should().Be(2);
+            (byTime[1].Tokens[0].SourceSpan.End.Column - byTime[1].Tokens[0].SourceSpan.Start.Column).Should().Be(3);
+            (byTime[2].Tokens[0].SourceSpan.End.Column - byTime[2].Tokens[0].SourceSpan.Start.Column).Should().Be(4);
+        }
+
+        [TestMethod, Priority(0)]
+        public async Task TokenCacheExpiration() {
+            var c = new TokenCache(5, TimeSpan.FromMilliseconds(50));
+            c.GetTokens("1", PythonLanguageVersion.V38);
+            await Task.Delay(10);
+            c.GetTokens("12", PythonLanguageVersion.V38);
+            await Task.Delay(10);
+            c.GetTokens("123", PythonLanguageVersion.V38);
+            await Task.Delay(40);
+            c.GetTokens("1234", PythonLanguageVersion.V38);
+
+            (DateTime AccessTime, IReadOnlyList<TokenInfo> Tokens)[] e = c.Entries.ToArray();
+            e.Should().HaveCount(2);
+            var byTime = e.OrderBy(x => x.AccessTime).ToArray();
+            (byTime[0].Tokens[0].SourceSpan.End.Column - byTime[1].Tokens[0].SourceSpan.Start.Column).Should().Be(3);
+            (byTime[1].Tokens[0].SourceSpan.End.Column - byTime[1].Tokens[0].SourceSpan.Start.Column).Should().Be(4);
         }
     }
 }