Async Indexer for new LanguageServer

src/Analysis/Ast/Impl/Indexing/IIndexManager.cs

@@ -0,0 +1,13 @@
+using System;
+using System.Threading;
+using System.Threading.Tasks;
+using Microsoft.Python.Analysis.Documents;
+
+namespace Microsoft.Python.Analysis.Indexing {
+    internal interface IIndexManager : IDisposable {
+        Task AddRootDirectory(CancellationToken workspaceCancellationToken = default);
+        void ProcessFile(Uri uri, IDocument doc);
+        Task ProcessClosedFile(Uri uri, CancellationToken fileCancellationToken = default);
+        void ProcessFileIfIndexed(Uri uri, IDocument doc);
+    }
+}

src/Analysis/Ast/Impl/Indexing/IIndexParser.cs

@@ -0,0 +1,9 @@
+using System;
+using System.Threading;
+using System.Threading.Tasks;
+
+namespace Microsoft.Python.Analysis.Indexing {
+    internal interface IIndexParser : IDisposable {
+        Task ParseAsync(Uri uri, CancellationToken cancellationToken = default);
+    }
+}

src/Analysis/Ast/Impl/Indexing/ISymbolIndex.cs

@@ -0,0 +1,11 @@
+using System;
+using System.Collections.Generic;
+using Microsoft.Python.Parsing.Ast;
+
+namespace Microsoft.Python.Analysis.Indexing {
+    internal interface ISymbolIndex {
+        void UpdateIndex(Uri uri, PythonAst pythonAst);
+        IEnumerable<FlatSymbol> WorkspaceSymbols(string query);
+        void Delete(Uri uri);
+    }
+}

src/Analysis/Ast/Impl/Indexing/IndexManager.cs

@@ -0,0 +1,96 @@
+using System;
+using System.Collections.Concurrent;
+using System.Collections.Generic;
+using System.Threading;
+using System.Threading.Tasks;
+using Microsoft.Python.Analysis.Core.Interpreter;
+using Microsoft.Python.Analysis.Documents;
+using Microsoft.Python.Core.Diagnostics;
+using Microsoft.Python.Core.IO;
+using Microsoft.Python.Parsing;
+
+namespace Microsoft.Python.Analysis.Indexing {
+    internal class IndexManager : IIndexManager {
+        private readonly ISymbolIndex _symbolIndex;
+        private readonly IFileSystem _fileSystem;
+        private readonly IndexParser _indexParser;
+        private readonly string _workspaceRootPath;
+        private readonly string[] _includeFiles;
+        private readonly string[] _excludeFiles;
+        private readonly ConcurrentDictionary<Uri, bool> _indexedFiles = new ConcurrentDictionary<Uri, bool>();
+        private readonly CancellationTokenSource _allIndexCts = new CancellationTokenSource();
+
+        public IndexManager(ISymbolIndex symbolIndex, IFileSystem fileSystem, PythonLanguageVersion version, string rootPath, string[] includeFiles,
+            string[] excludeFiles) {
+            Check.ArgumentNotNull(nameof(fileSystem), fileSystem);
+            Check.ArgumentNotNull(nameof(rootPath), rootPath);
+            Check.ArgumentNotNull(nameof(includeFiles), includeFiles);
+            Check.ArgumentNotNull(nameof(excludeFiles), excludeFiles);
+
+            _symbolIndex = symbolIndex;
+            _fileSystem = fileSystem;
+            _indexParser = new IndexParser(symbolIndex, fileSystem, version);
+            _workspaceRootPath = rootPath;
+            _includeFiles = includeFiles;
+            _excludeFiles = excludeFiles;
+        }
+
+        public Task AddRootDirectory(CancellationToken workspaceCancellationToken = default) {
+            var linkedCts = CancellationTokenSource.CreateLinkedTokenSource(workspaceCancellationToken, _allIndexCts.Token);
+            var parseTasks = new List<Task>();
+            foreach (var fileInfo in WorkspaceFiles()) {
+                if (ModulePath.IsPythonSourceFile(fileInfo.FullName)) {
+                    Uri uri = new Uri(fileInfo.FullName);
+                    parseTasks.Add(_indexParser.ParseAsync(uri, linkedCts.Token).ContinueWith((task) => {
+                        linkedCts.Token.ThrowIfCancellationRequested();
+                        _indexedFiles[uri] = true;
+                    }));
+                }
+            }
+            return Task.WhenAll(parseTasks.ToArray());
+        }
+
+        private IEnumerable<IFileSystemInfo> WorkspaceFiles() {
+            return _fileSystem.GetDirectoryInfo(_workspaceRootPath).EnumerateFileSystemInfos(_includeFiles, _excludeFiles);
+        }
+
+        private bool IsFileIndexed(Uri uri) {
+            _indexedFiles.TryGetValue(uri, out var val);
+            return val;
+        }
+
+        public Task ProcessClosedFile(Uri uri, CancellationToken fileCancellationToken = default) {
+            var linkedCts = CancellationTokenSource.CreateLinkedTokenSource(fileCancellationToken, _allIndexCts.Token);
+            // If path is on workspace
+            if (IsFileOnWorkspace(uri)) {
+                // updates index and ignores previous AST
+                return _indexParser.ParseAsync(uri, linkedCts.Token);
+            } else {
+                // remove file from index
+                _indexedFiles.TryRemove(uri, out _);
+                _symbolIndex.Delete(uri);
+                return Task.CompletedTask;
+            }
+        }
+
+        private bool IsFileOnWorkspace(Uri uri) {
+            return _fileSystem.IsPathUnderRoot(_workspaceRootPath, uri.AbsolutePath);
+        }
+
+        public void ProcessFile(Uri uri, IDocument doc) {
+            _indexedFiles[uri] = true;
+            _symbolIndex.UpdateIndex(uri, doc.GetAnyAst());
+        }
+
+        public void ProcessFileIfIndexed(Uri uri, IDocument doc) {
+            if (IsFileIndexed(uri)) {
+                ProcessFile(uri, doc);
+            }
+        }
+
+        public void Dispose() {
+            _allIndexCts.Cancel();
+            _allIndexCts.Dispose();
+        }
+    }
+}

src/Analysis/Ast/Impl/Indexing/IndexParser.cs

@@ -0,0 +1,45 @@
+using System;
+using System.IO;
+using System.Threading;
+using System.Threading.Tasks;
+using Microsoft.Python.Core.Diagnostics;
+using Microsoft.Python.Core.IO;
+using Microsoft.Python.Parsing;
+
+namespace Microsoft.Python.Analysis.Indexing {
+    internal sealed class IndexParser : IIndexParser {
+        private readonly ISymbolIndex _symbolIndex;
+        private readonly IFileSystem _fileSystem;
+        private readonly PythonLanguageVersion _version;
+        private readonly CancellationTokenSource _allProcessingCts = new CancellationTokenSource();
+
+        public IndexParser(ISymbolIndex symbolIndex, IFileSystem fileSystem, PythonLanguageVersion version) {
+            Check.ArgumentNotNull(nameof(symbolIndex), symbolIndex);
+            Check.ArgumentNotNull(nameof(fileSystem), fileSystem);
+
+            _symbolIndex = symbolIndex;
+            _fileSystem = fileSystem;
+            _version = version;
+        }
+
+        public void Dispose() {
+            _allProcessingCts.Cancel();
+            _allProcessingCts.Dispose();
+        }
+
+        public Task ParseAsync(Uri uri, CancellationToken parseCancellationToken = default) {
+            var linkedParseCts = CancellationTokenSource.CreateLinkedTokenSource(_allProcessingCts.Token, parseCancellationToken);
+            var linkedParseToken = linkedParseCts.Token;
+            return Task.Run(() => {
+                if (!_fileSystem.FileExists(uri.AbsolutePath)) {
+                    throw new FileNotFoundException($"{uri.AbsolutePath} does not exist", uri.AbsolutePath);
+                }
+                using (var stream = _fileSystem.FileOpen(uri.AbsolutePath, FileMode.Open)) {
+                    var parser = Parser.CreateParser(stream, _version);
+                    linkedParseToken.ThrowIfCancellationRequested();
+                    _symbolIndex.UpdateIndex(uri, parser.ParseFile());
+                }
+            }, linkedParseToken);
+        }
+    }
+}

src/Analysis/Ast/Impl/Indexing/SymbolIndex.cs

@@ -0,0 +1,48 @@
+using System;
+using System.Collections.Concurrent;
+using System.Collections.Generic;
+using System.Linq;
+using Microsoft.Python.Core;
+using Microsoft.Python.Parsing.Ast;
+
+namespace Microsoft.Python.Analysis.Indexing {
+    internal sealed class SymbolIndex : ISymbolIndex {
+        private readonly ConcurrentDictionary<Uri, IReadOnlyList<HierarchicalSymbol>> _index = new ConcurrentDictionary<Uri, IReadOnlyList<HierarchicalSymbol>>();
+
+        public IEnumerable<HierarchicalSymbol> HierarchicalDocumentSymbols(Uri uri)
+            => _index.TryGetValue(uri, out var list) ? list : Enumerable.Empty<HierarchicalSymbol>();
+
+        public IEnumerable<FlatSymbol> WorkspaceSymbols(string query) {
+            foreach (var kvp in _index) {
+                foreach (var found in WorkspaceSymbolsQuery(query, kvp.Key, kvp.Value)) {
+                    yield return found;
+                }
+            }
+        }
+
+        private IEnumerable<FlatSymbol> WorkspaceSymbolsQuery(string query, Uri uri, IEnumerable<HierarchicalSymbol> symbols) {
+            // Some semblance of a BFS.
+            var queue = new Queue<(HierarchicalSymbol, string)>(symbols.Select(s => (s, (string)null)));
+
+            while (queue.Count > 0) {
+                var (sym, parent) = queue.Dequeue();
+
+                if (sym.Name.ContainsOrdinal(query, ignoreCase: true)) {
+                    yield return new FlatSymbol(sym.Name, sym.Kind, uri, sym.SelectionRange, parent);
+                }
+
+                foreach (var child in sym.Children.MaybeEnumerate()) {
+                    queue.Enqueue((child, sym.Name));
+                }
+            }
+        }
+
+        public void UpdateIndex(Uri uri, PythonAst ast) {
+            var walker = new SymbolIndexWalker(ast);
+            ast.Walk(walker);
+            _index[uri] = walker.Symbols;
+        }
+
+        public void Delete(Uri uri) => _index.TryRemove(uri, out var _);
+    }
+}