Optimize references removal (microsoft#925)

* Fix microsoft#668 (partial)

* Tests

* Revert "Tests"

This reverts commit 7ffc9db.

* Exp

* Limit concurrency

* Concurrency limit

* Drop cache after analysis

* Fix regression

* Fix test

* Optimize references removal

* Dead code

Author: Mikhail Arkhipov

src/Analysis/Ast/Impl/Modules/PythonModule.cs

@@ -403,10 +403,18 @@ public void NotifyAnalysisBegins() {
             lock (AnalysisLock) {
                 if (_updated) {
                     _updated = false;
-                    var analyzer = Services.GetService<IPythonAnalyzer>();
-                    foreach (var gs in analyzer.LoadedModules.Select(m => m.GlobalScope).OfType<IScope>().ExcludeDefault()) {
-                        foreach (var v in gs.TraverseDepthFirst(c => c.Children).SelectMany(s => s.Variables)) {
-                            v.RemoveReferences(this);
+                    // In all variables find those imported, then traverse imported modules
+                    // and remove references to this module. If variable refers to a module,
+                    // recurse into module but only process global scope.
+                    var importedVariables = ((IScope)GlobalScope)
+                        .TraverseDepthFirst(c => c.Children)
+                        .SelectMany(s => s.Variables)
+                        .Where(v => v.Source == VariableSource.Import);
+
+                    foreach (var v in importedVariables) {
+                        v.RemoveReferences(this);
+                        if(v.Value is IPythonModule module) {
+                            RemoveReferencesInModule(module);
                         }
                     }
                 }
@@ -448,6 +456,12 @@ public Task<IDocumentAnalysis> GetAnalysisAsync(int waitTime = 200, Cancellation
 
         #endregion
 
+        private void RemoveReferencesInModule(IPythonModule module) {
+            foreach (var v in module.GlobalScope.Variables) {
+                v.RemoveReferences(this);
+            }
+        }
+
         private string TryGetDocFromModuleInitFile() {
             if (string.IsNullOrEmpty(FilePath) || !FileSystem.FileExists(FilePath)) {
                 return string.Empty;

src/LanguageServer/Test/ReferencesTests.cs

@@ -241,7 +241,7 @@ def func(x):
             var uri1 = TestData.GetDefaultModuleUri();
             var uri2 = TestData.GetNextModuleUri();
 
-            var code2 = $@"
+            const string code2 = @"
 from module import x, y
 a = x
 b = y
@@ -284,6 +284,78 @@ def func(x):
             refs[0].uri.Should().Be(uri1);
         }
 
+        [TestMethod, Priority(0)]
+        public async Task RemoveReferenceNested() {
+            const string code1 = @"
+x = 1
+
+def func(x):
+    return x
+
+y = func(x)
+x = 2
+";
+            const string code2 = @"
+from module import x, y
+a = x
+b = y
+";
+            var code3 = @"
+from module1 import x, y
+a = x
+b = y
+";
+            var uri1 = TestData.GetDefaultModuleUri();
+            var uri2 = TestData.GetNextModuleUri();
+            var uri3 = TestData.GetNextModuleUri();
+
+            await CreateServicesAsync(PythonVersions.LatestAvailable3X, uri1.AbsolutePath);
+
+            var rdt = Services.GetService<IRunningDocumentTable>();
+            var doc1 = rdt.OpenDocument(uri1, code1);
+            var doc2 = rdt.OpenDocument(uri2, code2);
+            var doc3 = rdt.OpenDocument(uri3, code3);
+
+            var analysis = await GetDocumentAnalysisAsync(doc1);
+
+            var rs = new ReferenceSource(Services);
+            var refs = await rs.FindAllReferencesAsync(analysis.Document.Uri, new SourceLocation(7, 1), ReferenceSearchOptions.All);
+
+            refs.Should().HaveCount(5);
+            refs[0].range.Should().Be(6, 0, 6, 1);
+            refs[0].uri.Should().Be(uri1);
+            refs[1].range.Should().Be(1, 22, 1, 23);
+            refs[1].uri.Should().Be(uri2);
+            refs[2].range.Should().Be(3, 4, 3, 5);
+            refs[2].uri.Should().Be(uri2);
+            refs[3].range.Should().Be(1, 23, 1, 24);
+            refs[3].uri.Should().Be(uri3);
+            refs[4].range.Should().Be(3, 4, 3, 5);
+            refs[4].uri.Should().Be(uri3);
+
+            doc3.Update(new[] {
+                new DocumentChange {
+                    InsertedText = string.Empty,
+                    ReplacedSpan = new SourceSpan(4, 1, 4, 6)
+                },
+                new DocumentChange {
+                    InsertedText = string.Empty,
+                    ReplacedSpan = new SourceSpan(2, 22, 2, 25)
+                }
+            });
+            await GetDocumentAnalysisAsync(doc3);
+
+            refs = await rs.FindAllReferencesAsync(analysis.Document.Uri, new SourceLocation(7, 1), ReferenceSearchOptions.All);
+
+            refs.Should().HaveCount(3);
+            refs[0].range.Should().Be(6, 0, 6, 1);
+            refs[0].uri.Should().Be(uri1);
+            refs[1].range.Should().Be(1, 22, 1, 23);
+            refs[1].uri.Should().Be(uri2);
+            refs[2].range.Should().Be(3, 4, 3, 5);
+            refs[2].uri.Should().Be(uri2);
+        }
+
         [TestMethod, Priority(0)]
         public async Task UpdateReferencesOnEdit() {
             const string code = @"