Better legacy migration (#4436)

### Changes

1. Rename `LegacySymbolTable`, which is part of the non-LSP Dafny
server, to `SuperLegacySymbolTable`
1. Rename `SignatureAndCompletionTable` to
`LegacySignatureAndCompletionTable`
1. Let `SignatureAndCompletionTable` differentiate between positions in
different files, so it doesn't get confused between symbols in different
files when they have similar positions or ranges.
1. Only migrate state in `SignatureAndCompletionTable` that relates to
the changed file

### Testing
- There could be a test added for 3 but I didn't add it.
- Merging this PR into this other PR
(#4419), reduces the time of the
test `ManyFastEditsUsingLargeFilesAndIncludes` from 47 seconds to 8
seconds

<small>By submitting this pull request, I confirm that my contribution
is made under the terms of the [MIT
license](https://github.com/dafny-lang/dafny/blob/master/LICENSE.txt).</small>

Source/DafnyLanguageServer.Test/Synchronization/DeclarationLocationMigrationTest.cs

@@ -1,4 +1,5 @@
-using System.IO;
+using System;
+using System.IO;
 using Microsoft.Dafny.LanguageServer.IntegrationTest.Extensions;
 using Microsoft.Dafny.LanguageServer.Language.Symbols;
 using Microsoft.Dafny.LanguageServer.Workspace;
@@ -10,6 +11,7 @@
 using Xunit;
 using Xunit.Abstractions;
 using Xunit.Sdk;
+using Range = OmniSharp.Extensions.LanguageServer.Protocol.Models.Range;
 
 namespace Microsoft.Dafny.LanguageServer.IntegrationTest.Synchronization {
 
@@ -22,8 +24,8 @@ public class DeclarationLocationMigrationTest : SynchronizationTestBase {
     // TODO The "BodyEndToken" used by the CreateDeclarationDictionary.CreateDeclarationDictionary()
     //      does not incorporate the closing braces.
 
-    protected bool TryFindSymbolDeclarationByName(IdeState state, string symbolName, out SymbolLocation location) {
-      location = state.SignatureAndCompletionTable.Locations
+    protected bool TryFindSymbolDeclarationByName(IdeState state, string symbolName, out SymbolLocation location, Uri uri = null) {
+      location = state.SignatureAndCompletionTable.LocationsPerUri[uri ?? state.SignatureAndCompletionTable.LocationsPerUri.First().Key]
         .WithCancellation(CancellationToken)
         .Where(entry => entry.Key.Name == symbolName)
         .Select(entry => entry.Value)
@@ -318,25 +320,28 @@ public async Task PassingANullChangeRangeClearsSymbolsTable() {
 
     [Fact]
     public async Task PassingANullChangeRangePreservesForeignSymbols() {
+      var directory = Path.Combine(Directory.GetCurrentDirectory(), "Lookup/TestFiles");
       var source = "include \"foreign.dfy\"\nclass X {}";
-      var documentItem = CreateTestDocument(source, Path.Combine(Directory.GetCurrentDirectory(), "Lookup/TestFiles/test.dfy"));
+      var documentItem = CreateTestDocument(source, Path.Combine(directory, "test.dfy"));
+      var includePath = Path.Combine(directory, "foreign.dfy");
 
       await Client.OpenDocumentAndWaitAsync(documentItem, CancellationToken);
       var document = await Projects.GetResolvedDocumentAsyncNormalizeUri(documentItem.Uri);
       Assert.NotNull(document);
-      Assert.True(TryFindSymbolDeclarationByName(document, "A", out var _));
+      var uri = documentItem.Uri.ToUri();
+      Assert.True(TryFindSymbolDeclarationByName(document, "A", out var _, new Uri(includePath)));
 
       // Try a change that breaks resolution.  Symbols for `foreign.dfy` are kept.
       await ApplyChangeAndWaitCompletionAsync(ref documentItem, null, "; include \"foreign.dfy\"\nclass Y {}");
       document = await Projects.GetResolvedDocumentAsyncNormalizeUri(documentItem);
       Assert.NotNull(document);
-      Assert.True(TryFindSymbolDeclarationByName(document, "A", out var _));
+      Assert.True(TryFindSymbolDeclarationByName(document, "A", out var _, new Uri(includePath)));
 
       // Finally we drop the reference to `foreign.dfy` and confirm that `A` is not accessible any more.
       await ApplyChangeAndWaitCompletionAsync(ref documentItem, null, "class Y {}");
       document = await Projects.GetResolvedDocumentAsyncNormalizeUri(documentItem);
       Assert.NotNull(document);
-      Assert.False(TryFindSymbolDeclarationByName(document, "A", out var _));
+      Assert.False(TryFindSymbolDeclarationByName(document, "A", out var _, uri));
     }
 
     public DeclarationLocationMigrationTest(ITestOutputHelper output) : this(output, LogLevel.Information) {

Source/DafnyLanguageServer.Test/Synchronization/LookupMigrationTest.cs

@@ -1,4 +1,5 @@
 using System.IO;
+using System.Linq;
 using Microsoft.Dafny.LanguageServer.IntegrationTest.Extensions;
 using OmniSharp.Extensions.LanguageServer.Protocol.Models;
 using System.Threading.Tasks;
@@ -266,7 +267,7 @@ reads this
       await Client.OpenDocumentAndWaitAsync(documentItem, CancellationToken);
       var originalDocument = await Projects.GetResolvedDocumentAsyncNormalizeUri(documentItem.Uri);
       Assert.NotNull(originalDocument);
-      var lookupCountBefore = originalDocument.SignatureAndCompletionTable.LookupTree.Count;
+      var lookupCountBefore = originalDocument.SignatureAndCompletionTable.LookupTreePerUri.First().Value.Count;
 
       await ApplyChangeAndWaitCompletionAsync(
         ref documentItem,
@@ -276,7 +277,7 @@ await ApplyChangeAndWaitCompletionAsync(
       var document = await Projects.GetResolvedDocumentAsyncNormalizeUri(documentItem.Uri);
       Assert.NotNull(document);
       Assert.False(document.SignatureAndCompletionTable.TryGetSymbolAt((6, 9), out var _));
-      Assert.Equal(lookupCountBefore - 1, document.SignatureAndCompletionTable.LookupTree.Count);
+      Assert.Equal(lookupCountBefore - 1, document.SignatureAndCompletionTable.LookupTreePerUri.First().Value.Count);
     }
 
     [Fact]

Source/DafnyLanguageServer.Test/Synchronization/SymbolMigrationTest.cs

@@ -30,7 +30,7 @@ await ApplyChangeAndWaitCompletionAsync(
       var document = await Projects.GetResolvedDocumentAsyncNormalizeUri(documentItem.Uri);
       Assert.NotNull(document);
       Assert.True(document.SignatureAndCompletionTable.Resolved);
-      Assert.Equal(2, document.SignatureAndCompletionTable.Locations.Keys.OfType<FunctionSymbol>().Count());
+      Assert.Equal(2, document.SignatureAndCompletionTable.LocationsPerUri.First().Value.Keys.OfType<FunctionSymbol>().Count());
     }
 
     [Fact]
@@ -98,7 +98,7 @@ await ApplyChangeAndWaitCompletionAsync(
       var document = await Projects.GetResolvedDocumentAsyncNormalizeUri(documentItem.Uri);
       Assert.NotNull(document);
       Assert.True(document.SignatureAndCompletionTable.Resolved);
-      Assert.Single(document.SignatureAndCompletionTable.Locations.Keys.OfType<MethodSymbol>());
+      Assert.Single(document.SignatureAndCompletionTable.LocationsPerUri.First().Value.Keys.OfType<MethodSymbol>());
     }
 
     public SymbolMigrationTest(ITestOutputHelper output) : base(output) {

Source/DafnyLanguageServer.Test/Unit/GhostStateDiagnosticCollectorTest.cs

@@ -1,5 +1,6 @@
 using System;
 using System.Collections.Generic;
+using System.Collections.Immutable;
 using System.IO;
 using System.Linq;
 using System.Threading;
@@ -69,8 +70,8 @@ public void EnsureResilienceAgainstErrors() {
     var source = new CancellationTokenSource();
     source.CancelAfter(TimeSpan.FromSeconds(50));
     var ghostDiagnostics = ghostStateDiagnosticCollector.GetGhostStateDiagnostics(
-      new SignatureAndCompletionTable(null!, new CompilationUnit(rootUri, program),
-        null!, null!, new IntervalTree<Position, ILocalizableSymbol>(), true)
+      new LegacySignatureAndCompletionTable(null!, new CompilationUnit(rootUri, program),
+        null!, null!, ImmutableDictionary<Uri, IIntervalTree<Position, ILocalizableSymbol>>.Empty, true)
       , source.Token);
     Assert.Empty(ghostDiagnostics);
   }

Source/DafnyLanguageServer/Language/GhostStateDiagnosticCollector.cs

@@ -30,7 +30,7 @@ public GhostStateDiagnosticCollector(DafnyOptions options, ILogger<GhostStateDia
     }
 
     public IReadOnlyDictionary<Uri, IReadOnlyList<Range>> GetGhostStateDiagnostics(
-      SignatureAndCompletionTable signatureAndCompletionTable, CancellationToken cancellationToken) {
+      LegacySignatureAndCompletionTable signatureAndCompletionTable, CancellationToken cancellationToken) {
       if (!options.Get(ServerCommand.GhostIndicators)) {
         return ImmutableDictionary<Uri, IReadOnlyList<Range>>.Empty;
       }

Source/DafnyLanguageServer/Language/IGhostStateDiagnosticCollector.cs

@@ -20,6 +20,6 @@ public interface IGhostStateDiagnosticCollector {
     /// <exception cref="System.OperationCanceledException">Thrown when the cancellation was requested before completion.</exception>
     /// <exception cref="System.ObjectDisposedException">Thrown if the cancellation token was disposed before the completion.</exception>
     IReadOnlyDictionary<Uri, IReadOnlyList<Range>> GetGhostStateDiagnostics(
-      SignatureAndCompletionTable signatureAndCompletionTable, CancellationToken cancellationToken);
+      LegacySignatureAndCompletionTable signatureAndCompletionTable, CancellationToken cancellationToken);
   }
 }

Source/DafnyLanguageServer/Language/Symbols/ISymbolTableFactory.cs

@@ -16,7 +16,7 @@ public interface ISymbolTableFactory {
     /// <returns>A symbol table representing the provided compilation unit.</returns>
     /// <exception cref="System.OperationCanceledException">Thrown when the cancellation was requested before completion.</exception>
     /// <exception cref="System.ObjectDisposedException">Thrown if the cancellation token was disposed before the completion.</exception>
-    SignatureAndCompletionTable CreateFrom(CompilationUnit compilationUnit, CancellationToken cancellationToken);
+    LegacySignatureAndCompletionTable CreateFrom(CompilationUnit compilationUnit, CancellationToken cancellationToken);
     SymbolTable CreateFrom(Dafny.Program program, Compilation compilation, CancellationToken cancellationToken);
   }
 }

Source/DafnyLanguageServer/Language/Symbols/SignatureAndCompletionTable.cs → Source/DafnyLanguageServer/Language/Symbols/LegacySignatureAndCompletionTable.cs

@@ -2,9 +2,10 @@
 using IntervalTree;
 using OmniSharp.Extensions.LanguageServer.Protocol.Models;
 using System.Collections.Generic;
+using System.Collections.Immutable;
 using System.Diagnostics.CodeAnalysis;
+using System.Linq;
 using System.Threading;
-using Microsoft.Dafny.LanguageServer.Workspace;
 using Microsoft.Extensions.Logging;
 using Microsoft.Extensions.Logging.Abstractions;
 using AstElement = System.Object;
@@ -14,8 +15,8 @@ namespace Microsoft.Dafny.LanguageServer.Language.Symbols {
   /// <summary>
   /// Represents the symbol table
   /// </summary>
-  public class SignatureAndCompletionTable {
-    private readonly ILogger<SignatureAndCompletionTable> logger;
+  public class LegacySignatureAndCompletionTable {
+    private readonly ILogger<LegacySignatureAndCompletionTable> logger;
 
     // TODO Guard the properties from changes
     public CompilationUnit CompilationUnit { get; }
@@ -28,12 +29,12 @@ public class SignatureAndCompletionTable {
     /// <summary>
     /// Gets the dictionary allowing to resolve the location of a specified symbol. Do not modify this instance.
     /// </summary>
-    public IDictionary<ILegacySymbol, SymbolLocation> Locations { get; }
+    public ImmutableDictionary<Uri, IDictionary<ILegacySymbol, SymbolLocation>> LocationsPerUri { get; }
 
     /// <summary>
     /// Gets the interval tree backing this symbol table. Do not modify this instance.
     /// </summary>
-    public IIntervalTree<Position, ILocalizableSymbol> LookupTree { get; }
+    public ImmutableDictionary<Uri, IIntervalTree<Position, ILocalizableSymbol>> LookupTreePerUri { get; }
 
     /// <summary>
     /// <c>true</c> if the symbol table results from a successful resolution by the dafny resolver.
@@ -42,14 +43,14 @@ public class SignatureAndCompletionTable {
 
     private readonly DafnyLangTypeResolver typeResolver;
 
-    public static SignatureAndCompletionTable Empty(DafnyOptions options, DafnyProject project) {
+    public static LegacySignatureAndCompletionTable Empty(DafnyOptions options, DafnyProject project) {
       var emptyProgram = GetEmptyProgram(options, project.Uri);
-      return new SignatureAndCompletionTable(
-        NullLogger<SignatureAndCompletionTable>.Instance,
+      return new LegacySignatureAndCompletionTable(
+        NullLogger<LegacySignatureAndCompletionTable>.Instance,
         new CompilationUnit(project.Uri, emptyProgram),
         new Dictionary<object, ILocalizableSymbol>(),
-        new Dictionary<ILegacySymbol, SymbolLocation>(),
-        new IntervalTree<Position, ILocalizableSymbol>(),
+        ImmutableDictionary<Uri, IDictionary<ILegacySymbol, SymbolLocation>>.Empty,
+        ImmutableDictionary<Uri, IIntervalTree<Position, ILocalizableSymbol>>.Empty,
         symbolsResolved: false);
     }
 
@@ -68,35 +69,31 @@ public static Program GetEmptyProgram(DafnyOptions options, Uri uri) {
       return emptyProgram;
     }
 
-    public SignatureAndCompletionTable(
-        ILogger<SignatureAndCompletionTable> iLogger,
+    public LegacySignatureAndCompletionTable(
+        ILogger<LegacySignatureAndCompletionTable> iLogger,
         CompilationUnit compilationUnit,
         IDictionary<AstElement, ILocalizableSymbol> declarations,
-        IDictionary<ILegacySymbol, SymbolLocation> locations,
-        IIntervalTree<Position, ILocalizableSymbol> lookupTree,
+        ImmutableDictionary<Uri, IDictionary<ILegacySymbol, SymbolLocation>> locationsPerUri,
+        ImmutableDictionary<Uri, IIntervalTree<Position, ILocalizableSymbol>> lookupTreePerUri,
         bool symbolsResolved
     ) {
       CompilationUnit = compilationUnit;
       Declarations = declarations;
-      Locations = locations;
-      LookupTree = lookupTree;
+      LocationsPerUri = locationsPerUri;
+      LookupTreePerUri = lookupTreePerUri;
       Resolved = symbolsResolved;
       typeResolver = new DafnyLangTypeResolver(declarations);
       logger = iLogger;
-
-      // TODO IntervalTree goes out of sync after any change and "fixes" its state upon the first query. Replace it with another implementation that can be queried without potential side-effects.
-      LookupTree.Query(new Position(0, 0));
     }
 
     /// <summary>
     /// Tries to get a symbol at the specified location.
+    /// Only used for testing.
     /// </summary>
-    /// <param name="position">The requested position.</param>
-    /// <param name="symbol">The symbol that could be identified at the given position, or <c>null</c> if no symbol could be identified.</param>
-    /// <returns><c>true</c> if a symbol was found, otherwise <c>false</c>.</returns>
-    /// <exception cref="System.InvalidOperationException">Thrown if there was one more symbol at the specified position. This should never happen, unless there was an error.</exception>
     public bool TryGetSymbolAt(Position position, [NotNullWhen(true)] out ILocalizableSymbol? symbol) {
-      var symbolsAtPosition = LookupTree.Query(position);
+      var intervalTree = LookupTreePerUri.First().Value;
+
+      var symbolsAtPosition = intervalTree.Query(position);
       symbol = null;
       // Use case: function f(a: int) {}, and hover over a.
       foreach (var potentialSymbol in symbolsAtPosition) {
@@ -110,29 +107,20 @@ public bool TryGetSymbolAt(Position position, [NotNullWhen(true)] out ILocalizab
       return symbol != null;
     }
 
-    /// <summary>
-    /// Tries to get the location of the given symbol.
-    /// </summary>
-    /// <param name="symbol">The symbol to get the location of.</param>
-    /// <param name="location">The current location of the specified symbol, or <c>null</c> if no location of the given symbol is known.</param>
-    /// <returns><c>true</c> if a location was found, otherwise <c>false</c>.</returns>
-    public bool TryGetLocationOf(ILegacySymbol symbol, [NotNullWhen(true)] out SymbolLocation? location) {
-      return Locations.TryGetValue(symbol, out location);
-    }
-
     /// <summary>
     /// Resolves the innermost symbol that encloses the given position.
     /// </summary>
+    /// <param name="uri"></param>
     /// <param name="position">The position to get the innermost symbol of.</param>
     /// <param name="cancellationToken">A token to cancel the update operation before its completion.</param>
     /// <returns>The innermost symbol at the specified position.</returns>
     /// <exception cref="System.OperationCanceledException">Thrown when the cancellation was requested before completion.</exception>
     /// <exception cref="System.ObjectDisposedException">Thrown if the cancellation token was disposed before the completion.</exception>
-    public ILegacySymbol GetEnclosingSymbol(Position position, CancellationToken cancellationToken) {
+    public ILegacySymbol GetEnclosingSymbol(Uri uri, Position position, CancellationToken cancellationToken) {
       // TODO use a suitable data-structure to resolve the locations efficiently.
       ILegacySymbol innerMostSymbol = CompilationUnit;
       var innerMostRange = new Range(new Position(0, 0), new Position(int.MaxValue, int.MaxValue));
-      foreach (var (symbol, location) in Locations) {
+      foreach (var (symbol, location) in LocationsPerUri.GetValueOrDefault(uri) ?? ImmutableDictionary<ILegacySymbol, SymbolLocation>.Empty) {
         cancellationToken.ThrowIfCancellationRequested();
         var range = location.Declaration;
         if (IsEnclosedBy(innerMostRange, range) && IsInside(range, position)) {