diff --git a/src/LanguageServer/Impl/Sources/DocumentHighlightSource.cs b/src/LanguageServer/Impl/Sources/DocumentHighlightSource.cs index afe1efd42..38a152614 100644 --- a/src/LanguageServer/Impl/Sources/DocumentHighlightSource.cs +++ b/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 DocumentHighlightAsync(Uri uri, SourceLoc var definition = definitionSource.FindDefinition(analysis, location, out var definingMember); if (definition == null || definingMember == null) { - return Array.Empty(); + 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(); + } + } + + + internal class TokenCache { + internal class Entry { + public DateTime AccessTime; + public WeakReference> Tokens; + } + + private readonly Dictionary _cache = new Dictionary(); + private readonly int _maxEntries; + private readonly TimeSpan _expiration; + + public TokenCache(int maxEntries, TimeSpan expiration) { + _maxEntries = maxEntries; + _expiration = expiration; + } + + public IReadOnlyList GetTokens(string content, PythonLanguageVersion languageVersion) { + IReadOnlyList 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>(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)> 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))); + } + } } } diff --git a/src/LanguageServer/Test/DocumentHighlightTests.cs b/src/LanguageServer/Test/DocumentHighlightTests.cs index 4781b9db8..cfad14cfd 100644 --- a/src/LanguageServer/Test/DocumentHighlightTests.cs +++ b/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 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 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); } } }