Symbols limit

Author: MikhailArkhipov

src/Analysis/Engine/Impl/Infrastructure/Extensions/EnumerableExtensions.cs

@@ -59,5 +59,24 @@ public static bool SetEquals<T>(this IEnumerable<T> source, IEnumerable<T> other
 
         public static IEnumerable<T> Keys<T, U>(this IEnumerable<KeyValuePair<T, U>> source) => source.Select(GetKey);
         public static IEnumerable<T> ExcludeDefault<T>(this IEnumerable<T> source) => source.Where(i => !Equals(i, default(T)));
+
+
+        public static IEnumerable<T> TraverseBreadthFirst<T>(this T root, Func<T, IEnumerable<T>> selectChildren) {
+            var items = new Queue<T>();
+            items.Enqueue(root);
+            while (items.Count > 0) {
+                var item = items.Dequeue();
+                yield return item;
+
+                var childen = selectChildren(item);
+                if (childen == null) {
+                    continue;
+                }
+
+                foreach (var child in childen) {
+                    items.Enqueue(child);
+                }
+            }
+        }
     }
 }

src/LanguageServer/Impl/Implementation/Server.WorkspaceSymbols.cs

@@ -14,6 +14,7 @@
 // See the Apache Version 2.0 License for specific language governing
 // permissions and limitations under the License.
 
+using System;
 using System.Collections.Generic;
 using System.Linq;
 using System.Threading;
@@ -26,7 +27,7 @@
 
 namespace Microsoft.Python.LanguageServer.Implementation {
     public sealed partial class Server {
-        private static int _symbolHierarchyDepthLimit = 1;
+        private static int _symbolHierarchyMaxSymbols = 1000;
 
         public override async Task<SymbolInformation[]> WorkspaceSymbols(WorkspaceSymbolParams @params, CancellationToken cancellationToken) {
             await WaitForCompleteAnalysisAsync(cancellationToken);
@@ -82,16 +83,16 @@ private static IEnumerable<IMemberResult> GetModuleVariables(ProjectEntry entry,
                     }
                     return false;
                 })
-            .Concat(GetChildScopesVariables(analysis, analysis.Scope, opts, 0));
+                .Take(_symbolHierarchyMaxSymbols)
+                .Concat(GetChildScopesVariables(analysis, analysis.Scope, opts))
+                .Take(_symbolHierarchyMaxSymbols);
         }
 
-        private static IEnumerable<IMemberResult> GetChildScopesVariables(IModuleAnalysis analysis, IScope scope, GetMemberOptions opts, int currentDepth)
-            => currentDepth < _symbolHierarchyDepthLimit
-                ? scope.Children.SelectMany(c => GetScopeVariables(analysis, c, opts, currentDepth))
-                : Enumerable.Empty<IMemberResult>();
+        private static IEnumerable<IMemberResult> GetChildScopesVariables(IModuleAnalysis analysis, IScope scope, GetMemberOptions opts)
+            => scope.TraverseBreadthFirst(s => s.Children).SelectMany(c => GetScopeVariables(analysis, c, opts));
 
-        private static IEnumerable<IMemberResult> GetScopeVariables(IModuleAnalysis analysis, IScope scope, GetMemberOptions opts, int currentDepth)
-            => analysis.GetAllAvailableMembersFromScope(scope, opts).Concat(GetChildScopesVariables(analysis, scope, opts, currentDepth + 1));
+        private static IEnumerable<IMemberResult> GetScopeVariables(IModuleAnalysis analysis, IScope scope, GetMemberOptions opts)
+            => analysis.GetAllAvailableMembersFromScope(scope, opts).Concat(GetChildScopesVariables(analysis, scope, opts));
 
         private SymbolInformation ToSymbolInformation(IMemberResult m) {
             var res = new SymbolInformation {
@@ -115,8 +116,9 @@ private SymbolInformation ToSymbolInformation(IMemberResult m) {
         }
 
         private DocumentSymbol[] ToDocumentSymbols(List<IMemberResult> members) {
-            var childMap = new Dictionary<IMemberResult, List<IMemberResult>>();
             var topLevel = new List<IMemberResult>();
+            var childMap = new Dictionary<IMemberResult, List<IMemberResult>>();
+            var totalCount = 0;
 
             foreach (var m in members) {
                 var parent = members.FirstOrDefault(x => x.Scope?.Node == m.Scope?.OuterScope?.Node && x.Name == m.Scope?.Name);
@@ -128,18 +130,23 @@ private DocumentSymbol[] ToDocumentSymbols(List<IMemberResult> members) {
                 } else {
                     topLevel.Add(m);
                 }
+                if (++totalCount >= _symbolHierarchyMaxSymbols) {
+                    break;
+                }
             }
 
-            var symbols = topLevel
-                .GroupBy(mr => mr.Name)
-                .Select(g => g.First())
-                .Select(m => ToDocumentSymbol(m, childMap, 0))
-                .ToArray();
+            var symbols = topLevel.SelectMany(t => t
+                    .TraverseBreadthFirst(c => childMap.ContainsKey(c) ? childMap[c] : Enumerable.Empty<IMemberResult>())
+                    .Take(_symbolHierarchyMaxSymbols)
+                    .GroupBy(mr => mr.Name)
+                    .Select(g => g.First())
+                    .Select(m => ToDocumentSymbol(m, childMap)))
+                    .ToArray();
 
             return symbols;
         }
 
-        private DocumentSymbol ToDocumentSymbol(IMemberResult m, Dictionary<IMemberResult, List<IMemberResult>> childMap, int currentDepth) {
+        private DocumentSymbol ToDocumentSymbol(IMemberResult m, Dictionary<IMemberResult, List<IMemberResult>> childMap) {
             var res = new DocumentSymbol {
                 name = m.Name,
                 detail = m.Name,
@@ -148,10 +155,12 @@ private DocumentSymbol ToDocumentSymbol(IMemberResult m, Dictionary<IMemberResul
                 _functionKind = GetFunctionKind(m)
             };
 
-            if (childMap.TryGetValue(m, out var children) && currentDepth < _symbolHierarchyDepthLimit) {
-                res.children = children.Select(x => ToDocumentSymbol(x, childMap, currentDepth + 1)).ToArray();
+            if (childMap.TryGetValue(m, out var children)) {
+                res.children = children
+                    .Select(x => ToDocumentSymbol(x, childMap))
+                    .ToArray();
             } else {
-                res.children = new DocumentSymbol[0];
+                res.children = Array.Empty<DocumentSymbol>();
             }
 
             var loc = m.Locations.FirstOrDefault(l => !string.IsNullOrEmpty(l.FilePath));

src/LanguageServer/Impl/Implementation/Server.cs

@@ -684,7 +684,7 @@ private bool HandleConfigurationChanges(ServerSettings newSettings) {
             var oldSettings = Settings;
             Settings = newSettings;
 
-            _symbolHierarchyDepthLimit = Settings.analysis.symbolsHierarchyDepthLimit;
+            _symbolHierarchyMaxSymbols = Settings.analysis.symbolsHierarchyMaxSymbols;
 
             if (oldSettings == null) {
                 return true;

src/LanguageServer/Impl/LanguageServer.cs

@@ -142,7 +142,7 @@ public async Task DidChangeConfiguration(JToken token, CancellationToken cancell
                 var analysis = pythonSection["analysis"];
                 settings.analysis.openFilesOnly = GetSetting(analysis, "openFilesOnly", false);
                 settings.diagnosticPublishDelay = GetSetting(analysis, "diagnosticPublishDelay", 1000);
-                settings.symbolsHierarchyDepthLimit = GetSetting(analysis, "symbolsHierarchyDepthLimit", 10);
+                settings.symbolsHierarchyMaxSymbols = GetSetting(analysis, "symbolsHierarchyDepthLimit", 1000);
 
                 _ui.SetLogLevel(GetLogLevel(analysis));
 

src/LanguageServer/Impl/LanguageServerSettings.cs

@@ -17,6 +17,6 @@
 namespace Microsoft.Python.LanguageServer.Implementation {
     public sealed class LanguageServerSettings: ServerSettings {
         public int diagnosticPublishDelay = 1000;
-        public int symbolsHierarchyDepthLimit = 10;
+        public int symbolsHierarchyMaxSymbols = 1000;
     }
 }