Wire up events from RDT to module resolution

src/Analysis/Ast/Impl/Analyzer/Definitions/IPythonAnalyzer.cs

@@ -37,7 +37,12 @@ public interface IPythonAnalyzer {
         /// Invalidates current analysis for the module, assuming that AST for the new analysis will be provided later.
         /// </summary>
         void InvalidateAnalysis(IPythonModule module);
-
+
+        /// <summary>
+        /// Removes modules from the analysis.
+        /// </summary>
+        void RemoveAnalysis(IPythonModule module);
+
         /// <summary>
         /// Get most recent analysis for module. If after specified time analysis isn't available, returns previously calculated analysis.
         /// </summary>

src/Analysis/Ast/Impl/Analyzer/Evaluation/ExpressionEval.Scopes.cs

@@ -175,7 +175,7 @@ public IDisposable OpenScope(IPythonModule module, ScopeStatement node, out Scop
                     if (scope == null) {
                         scope = new Scope(node, fromScope, true);
                         fromScope.AddChildScope(scope);
-                        _scopeLookupCache[node] = fromScope;
+                        _scopeLookupCache[node] = scope;
                     }
                 }
 

src/Analysis/Ast/Impl/Analyzer/PythonAnalyzer.cs

@@ -112,6 +112,12 @@ public void InvalidateAnalysis(IPythonModule module) {
             }
         }
 
+        public void RemoveAnalysis(IPythonModule module) {
+            lock (_syncObj) {
+                _analysisEntries.Remove(new AnalysisModuleKey(module));
+            }
+        }
+
         public void EnqueueDocumentForAnalysis(IPythonModule module, ImmutableArray<IPythonModule> analysisDependencies) {
             var key = new AnalysisModuleKey(module);
             PythonAnalyzerEntry entry;
@@ -269,7 +275,7 @@ private async Task<int> AnalyzeAffectedEntriesAsync(IDependencyChainWalker<Analy
 
             return remaining;
         }
-        
+
         private void LoadMissingDocuments(IPythonInterpreter interpreter, ImmutableArray<AnalysisModuleKey> missingKeys) {
             foreach (var (moduleName, _, isTypeshed) in missingKeys) {
                 var moduleResolution = isTypeshed ? interpreter.TypeshedResolution : interpreter.ModuleResolution;
@@ -356,7 +362,7 @@ private void AnalyzeEntry(PythonAnalyzerEntry entry, IPythonModule module, Pytho
             // Python analyzer to call NotifyAnalysisComplete.
             walker.Complete();
             cancellationToken.ThrowIfCancellationRequested();
-            var analysis = new DocumentAnalysis((IDocument) module, version, walker.GlobalScope, walker.Eval);
+            var analysis = new DocumentAnalysis((IDocument)module, version, walker.GlobalScope, walker.Eval);
 
             (module as IAnalyzable)?.NotifyAnalysisComplete(analysis);
             entry.TrySetAnalysis(analysis, version);

src/Analysis/Ast/Impl/Analyzer/PythonInterpreter.cs

@@ -47,18 +47,19 @@ private async Task LoadBuiltinTypesAsync(string root, IServiceManager sm, Cancel
 
             sm.AddService(this);
             _moduleResolution = new MainModuleResolution(root, sm);
-            await _moduleResolution.InitializeAsync(cancellationToken);
 
-            _stubResolution = new TypeshedResolution(sm);
-            await _stubResolution.InitializeAsync(cancellationToken);
-
-            var builtinModule = _moduleResolution.BuiltinsModule;
             lock (_lock) {
+                var builtinModule = _moduleResolution.CreateBuiltinsModule();
                 _builtinTypes[BuiltinTypeId.NoneType]
                     = new PythonType("NoneType", builtinModule, string.Empty, LocationInfo.Empty, BuiltinTypeId.NoneType);
                 _builtinTypes[BuiltinTypeId.Unknown]
                     = UnknownType = new PythonType("Unknown", builtinModule, string.Empty, LocationInfo.Empty);
             }
+            await _moduleResolution.InitializeAsync(cancellationToken);
+
+            _stubResolution = new TypeshedResolution(sm);
+            await _stubResolution.InitializeAsync(cancellationToken);
+
             await _moduleResolution.LoadBuiltinTypesAsync(cancellationToken);
         }
 

src/Analysis/Ast/Impl/Dependencies/DependencyResolver.cs

@@ -21,7 +21,6 @@
 
 namespace Microsoft.Python.Analysis.Dependencies {
     internal sealed class DependencyResolver<TKey, TValue> : IDependencyResolver<TKey, TValue> {
-        private readonly IDependencyFinder<TKey, TValue> _dependencyFinder;
         private readonly DependencyGraph<TKey, TValue> _vertices = new DependencyGraph<TKey, TValue>();
         private readonly Dictionary<TKey, DependencyVertex<TKey, TValue>> _changedVertices = new Dictionary<TKey, DependencyVertex<TKey, TValue>>();
         private readonly object _syncObj = new object();

src/Analysis/Ast/Impl/Documents/Definitions/IRunningDocumentTable.cs

@@ -55,6 +55,20 @@ public interface IRunningDocumentTable: IEnumerable<IDocument> {
         /// </summary>
         IDocument GetDocument(string name);
 
+        /// <summary>
+        /// Increase reference count of the document.
+        /// </summary>
+        /// <param name="uri"></param>
+        /// <returns>New lock count or -1 if document was not found.</returns>
+        int LockDocument(Uri uri);
+
+        /// <summary>
+        /// Decrease reference count of the document.
+        /// </summary>
+        /// <param name="uri"></param>
+        /// <returns>New lock count or -1 if document was not found.</returns>
+        int UnlockDocument(Uri uri);
+
         /// <summary>
         /// Fires when document is opened.
         /// </summary>
@@ -64,5 +78,10 @@ public interface IRunningDocumentTable: IEnumerable<IDocument> {
         /// Fires when document is closed.
         /// </summary>
         event EventHandler<DocumentEventArgs> Closed;
+
+        /// <summary>
+        /// Fires when document is removed.
+        /// </summary>
+        event EventHandler<DocumentEventArgs> Removed;
     }
 }

src/Analysis/Ast/Impl/Documents/RunningDocumentTable.cs

@@ -19,9 +19,9 @@
 using System.Diagnostics;
 using System.IO;
 using System.Linq;
+using Microsoft.Python.Analysis.Analyzer;
 using Microsoft.Python.Analysis.Modules;
 using Microsoft.Python.Core;
-using Microsoft.Python.Core.Diagnostics;
 
 namespace Microsoft.Python.Analysis.Documents {
     /// <summary>
@@ -52,6 +52,7 @@ public RunningDocumentTable(string workspaceRoot, IServiceContainer services) {
 
         public event EventHandler<DocumentEventArgs> Opened;
         public event EventHandler<DocumentEventArgs> Closed;
+        public event EventHandler<DocumentEventArgs> Removed;
 
         /// <summary>
         /// Adds file to the list of available documents.
@@ -60,17 +61,19 @@ public RunningDocumentTable(string workspaceRoot, IServiceContainer services) {
         /// <param name="content">Document content</param>
         /// <param name="filePath">Optional file path, if different from the URI.</param>
         public IDocument OpenDocument(Uri uri, string content, string filePath = null) {
-            var justOpened = false;
+            bool justOpened;
             DocumentEntry entry;
             lock (_lock) {
                 entry = FindDocument(null, uri);
                 if (entry == null) {
+                    var resolver = _services.GetService<IPythonInterpreter>().ModuleResolution.CurrentPathResolver;
+                    var moduleType = resolver.IsLibraryFile(uri.ToAbsolutePath()) ? ModuleType.Library : ModuleType.User;
                     var mco = new ModuleCreationOptions {
                         ModuleName = Path.GetFileNameWithoutExtension(uri.LocalPath),
                         Content = content,
                         FilePath = filePath,
                         Uri = uri,
-                        ModuleType = ModuleType.User
+                        ModuleType = moduleType
                     };
                     entry = CreateDocument(mco);
                 }
@@ -114,39 +117,61 @@ public IDocument GetDocument(string name) {
             }
         }
 
+        public int LockDocument(Uri uri) {
+            lock (_lock) {
+                if (_documentsByUri.TryGetValue(uri, out var entry)) {
+                    return ++entry.LockCount;
+                }
+                return -1;
+            }
+        }
+
+        public int UnlockDocument(Uri uri) {
+            lock (_lock) {
+                if (_documentsByUri.TryGetValue(uri, out var entry)) {
+                    return --entry.LockCount;
+                }
+                return -1;
+            }
+        }
+
         public IEnumerator<IDocument> GetEnumerator() => _documentsByUri.Values.Select(e => e.Document).GetEnumerator();
 
         public void CloseDocument(Uri documentUri) {
-            var justClosed = false;
+            var closed = false;
+            var removed = false;
             DocumentEntry entry;
             lock (_lock) {
                 if (_documentsByUri.TryGetValue(documentUri, out entry)) {
                     Debug.Assert(entry.LockCount >= 1);
 
                     if (entry.Document.IsOpen) {
                         entry.Document.IsOpen = false;
-                        justClosed = true;
+                        closed = true;
                     }
 
                     entry.LockCount--;
 
                     if (entry.LockCount == 0) {
                         _documentsByUri.Remove(documentUri);
                         _documentsByName.Remove(entry.Document.Name);
+                        removed = true;
                         entry.Document.Dispose();
                     }
-                    // TODO: Remove from module resolution?
                 }
             }
-            if(justClosed) {
+            if (closed) {
                 Closed?.Invoke(this, new DocumentEventArgs(entry.Document));
             }
+            if (removed) {
+                Removed?.Invoke(this, new DocumentEventArgs(entry.Document));
+            }
         }
 
         IEnumerator IEnumerable.GetEnumerator() => _documentsByUri.Values.GetEnumerator();
 
         public void Dispose() {
-            lock(_lock) {
+            lock (_lock) {
                 foreach (var d in _documentsByUri.Values.OfType<IDisposable>()) {
                     d.Dispose();
                 }

src/Analysis/Ast/Impl/Modules/Definitions/IModuleCache.cs

@@ -13,10 +13,6 @@
 // See the Apache Version 2.0 License for specific language governing
 // permissions and limitations under the License.
 
-using System.Threading;
-using System.Threading.Tasks;
-using Microsoft.Python.Analysis.Documents;
-
 namespace Microsoft.Python.Analysis.Modules {
     public interface IModuleCache {
         string GetCacheFilePath(string filePath);

src/Analysis/Ast/Impl/Modules/Definitions/IModuleResolution.cs

@@ -13,7 +13,6 @@
 // See the Apache Version 2.0 License for specific language governing
 // permissions and limitations under the License.
 
-using System.Collections.Generic;
 using System.Threading;
 using System.Threading.Tasks;
 using Microsoft.Python.Analysis.Core.DependencyResolution;

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

@@ -18,9 +18,7 @@
 using System.IO;
 using System.Security.Cryptography;
 using System.Text;
-using System.Threading;
 using System.Threading.Tasks;
-using Microsoft.Python.Analysis.Documents;
 using Microsoft.Python.Core;
 using Microsoft.Python.Core.IO;
 using Microsoft.Python.Core.Logging;
@@ -43,7 +41,7 @@ public ModuleCache(IPythonInterpreter interpreter, IServiceContainer services) {
             _log = services.GetService<ILogger>();
             _skipCache = string.IsNullOrEmpty(_interpreter.Configuration.DatabasePath);
         }
-        
+
         public string GetCacheFilePath(string filePath) {
             if (string.IsNullOrEmpty(filePath) || !PathEqualityComparer.IsValidPath(ModuleCachePath)) {
                 if (!_loggedBadDbPath) {

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

@@ -29,6 +29,7 @@
 using Microsoft.Python.Analysis.Values;
 using Microsoft.Python.Core;
 using Microsoft.Python.Core.Diagnostics;
+using Microsoft.Python.Core.Disposables;
 using Microsoft.Python.Core.IO;
 using Microsoft.Python.Core.Logging;
 using Microsoft.Python.Core.Text;
@@ -53,7 +54,7 @@ private enum State {
         }
 
         private readonly DocumentBuffer _buffer = new DocumentBuffer();
-        private readonly CancellationTokenSource _allProcessingCts = new CancellationTokenSource();
+        private readonly DisposeToken _disposeToken = DisposeToken.Create< PythonModule>();
         private IReadOnlyList<DiagnosticsEntry> _parseErrors = Array.Empty<DiagnosticsEntry>();
         private readonly IDiagnosticsService _diagnosticsService;
 
@@ -242,8 +243,7 @@ private void LoadContent(string content) {
 
         protected virtual void Dispose(bool disposing) {
             _diagnosticsService?.Remove(Uri);
-            _allProcessingCts.Cancel();
-            _allProcessingCts.Dispose();
+            _disposeToken.TryMarkDisposed();
         }
         #endregion
 
@@ -310,7 +310,7 @@ public void Update(IEnumerable<DocumentChange> changes) {
                 _parseCts = new CancellationTokenSource();
 
                 _linkedParseCts?.Dispose();
-                _linkedParseCts = CancellationTokenSource.CreateLinkedTokenSource(_allProcessingCts.Token, _parseCts.Token);
+                _linkedParseCts = CancellationTokenSource.CreateLinkedTokenSource(_disposeToken.CancellationToken, _parseCts.Token);
 
                 _buffer.Update(changes);
                 Parse();
@@ -332,7 +332,7 @@ private void Parse() {
             _parseCts = new CancellationTokenSource();
 
             _linkedParseCts?.Dispose();
-            _linkedParseCts = CancellationTokenSource.CreateLinkedTokenSource(_allProcessingCts.Token, _parseCts.Token);
+            _linkedParseCts = CancellationTokenSource.CreateLinkedTokenSource(_disposeToken.CancellationToken, _parseCts.Token);
 
             ContentState = State.Parsing;
             _parsingTask = Task.Run(() => Parse(_linkedParseCts.Token), _linkedParseCts.Token);
@@ -383,7 +383,7 @@ private void Parse(CancellationToken cancellationToken) {
                 ContentState = State.Analyzing;
 
                 var analyzer = Services.GetService<IPythonAnalyzer>();
-                analyzer.EnqueueDocumentForAnalysis(this, ast, version, _allProcessingCts.Token);
+                analyzer.EnqueueDocumentForAnalysis(this, ast, version, _disposeToken.CancellationToken);
             }
 
             lock (AnalysisLock) {