Improve algorithm to avoid excess allocations

Author: andyleejordan

src/PowerShellEditorServices/Services/TextDocument/Handlers/DocumentSymbolHandler.cs

@@ -40,81 +40,106 @@ public PsesDocumentSymbolHandler(ILoggerFactory factory, WorkspaceService worksp
             DocumentSelector = LspUtils.PowerShellDocumentSelector
         };
 
-        // This turns a flat list of symbols into a hierarchical list. It's ugly because we're
-        // dealing with records and so sadly must slowly copy and replace things whenever need to do
-        // a modification, but it seems to work.
-        private static async Task<List<DocumentSymbol>> SortDocumentSymbols(List<DocumentSymbol> symbols, CancellationToken cancellationToken)
+        // This turns a flat list of symbols into a hierarchical list.
+        private static async Task<List<HierarchicalSymbol>> SortHierarchicalSymbols(List<HierarchicalSymbol> symbols, CancellationToken cancellationToken)
         {
-            // Sort by the start of the symbol definition.
+            // Sort by the start of the symbol definition (they're probably sorted but we need to be
+            // certain otherwise this algorithm won't work).
             symbols.Sort((x1, x2) => x1.Range.Start.CompareTo(x2.Range.Start));
 
-            List<DocumentSymbol> parents = new();
+            List<HierarchicalSymbol> parents = new();
 
-            foreach (DocumentSymbol symbol in symbols)
+            foreach (HierarchicalSymbol symbol in symbols)
             {
                 // This async method is pretty dense with synchronous code
                 // so it's helpful to add some yields.
                 await Task.Yield();
                 if (cancellationToken.IsCancellationRequested)
                 {
-                    return symbols;
+                    return parents;
                 }
-
-                // Base case.
+                // Base case where we haven't found any parents yet.
                 if (parents.Count == 0)
                 {
                     parents.Add(symbol);
                 }
-                // Symbol starts after end of last symbol parsed.
+                // If the symbol starts after end of last symbol parsed then it's a new parent.
                 else if (symbol.Range.Start > parents[parents.Count - 1].Range.End)
                 {
                     parents.Add(symbol);
                 }
-                // Find where it fits.
+                // Otherwise it's a child, we just need to figure out whose child it is.
                 else
                 {
-                    for (int i = 0; i < parents.Count; i++)
+                    foreach (HierarchicalSymbol parent in parents)
                     {
-                        DocumentSymbol parent = parents[i];
-                        if (parent.Range.Start <= symbol.Range.Start && symbol.Range.End <= parent.Range.End)
+                        // If the symbol starts after the parent starts and ends before the parent
+                        // ends then its a child.
+                        if (symbol.Range.Start > parent.Range.Start && symbol.Range.End < parent.Range.End)
                         {
-                            List<DocumentSymbol> children = new();
-                            if (parent.Children is not null)
-                            {
-                                children.AddRange(parent.Children);
-                            }
-                            children.Add(symbol);
-                            parents[i] = parent with { Children = children };
+                            parent.Children.Add(symbol);
                             break;
                         }
                     }
+                    // TODO: If we somehow exist the list of parents and didn't find a place for the
+                    // child...we have a problem.
                 }
             }
 
-            // Recursively sort the children.
-            for (int i = 0; i < parents.Count; i++)
+            // Now recursively sort the children into nested buckets of children too.
+            foreach (HierarchicalSymbol parent in parents)
             {
-                DocumentSymbol parent = parents[i];
-                if (parent.Children is not null)
-                {
-                    List<DocumentSymbol> children = new(parent.Children);
-                    children = await SortDocumentSymbols(children, cancellationToken).ConfigureAwait(false);
-                    parents[i] = parent with { Children = children };
-                }
+                List<HierarchicalSymbol> sortedChildren = await SortHierarchicalSymbols(parent.Children, cancellationToken).ConfigureAwait(false);
+                // Since this is a foreach we can't just assign to parent.Children and have to do
+                // this instead.
+                parent.Children.Clear();
+                parent.Children.AddRange(sortedChildren);
             }
 
             return parents;
         }
 
-        // AKA the outline feature
+        // This struct and the mapping function below exist to allow us to skip a *bunch* of
+        // unnecessary allocations when sorting the symbols since DocumentSymbol (which this is
+        // pretty much a mirror of) is an immutable record...but we need to constantly modify the
+        // list of children when sorting.
+        private struct HierarchicalSymbol
+        {
+            public SymbolKind Kind;
+            public Range Range;
+            public Range SelectionRange;
+            public string Name;
+            public List<HierarchicalSymbol> Children;
+        }
+
+        // Recursively turn our HierarchicalSymbol struct into OmniSharp's DocumentSymbol record.
+        private static List<DocumentSymbol> GetDocumentSymbolsFromHierarchicalSymbols(IEnumerable<HierarchicalSymbol> hierarchicalSymbols)
+        {
+            List<DocumentSymbol> documentSymbols = new();
+            foreach (HierarchicalSymbol symbol in hierarchicalSymbols)
+            {
+                documentSymbols.Add(new DocumentSymbol
+                {
+                    Kind = symbol.Kind,
+                    Range = symbol.Range,
+                    SelectionRange = symbol.SelectionRange,
+                    Name = symbol.Name,
+                    Children = GetDocumentSymbolsFromHierarchicalSymbols(symbol.Children)
+                });
+            }
+            return documentSymbols;
+        }
+
+        // AKA the outline feature!
         public override async Task<SymbolInformationOrDocumentSymbolContainer> Handle(DocumentSymbolParams request, CancellationToken cancellationToken)
         {
             _logger.LogDebug($"Handling document symbols for {request.TextDocument.Uri}");
 
             ScriptFile scriptFile = _workspaceService.GetFile(request.TextDocument.Uri);
-            List<DocumentSymbol> symbols = new();
 
-            foreach (SymbolReference r in ProvideDocumentSymbols(scriptFile))
+            List<HierarchicalSymbol> hierarchicalSymbols = new();
+
+            foreach (SymbolReference symbolReference in ProvideDocumentSymbols(scriptFile))
             {
                 // This async method is pretty dense with synchronous code
                 // so it's helpful to add some yields.
@@ -128,36 +153,33 @@ public override async Task<SymbolInformationOrDocumentSymbolContainer> Handle(Do
                 //
                 // TODO: We should also include function invocations that are part of DSLs (like
                 // Invoke-Build etc.).
-                if (!r.IsDeclaration || r.Type is SymbolType.Parameter)
+                if (!symbolReference.IsDeclaration || symbolReference.Type is SymbolType.Parameter)
                 {
                     continue;
                 }
 
-                // TODO: This now needs the Children property filled out to support hierarchical
-                // symbols, and we don't have the information nor algorithm to do that currently.
-                // OmniSharp was previously doing this for us based on the range, perhaps we can
-                // find that logic and reuse it.
-                symbols.Add(new DocumentSymbol
+                hierarchicalSymbols.Add(new HierarchicalSymbol
                 {
-                    Kind = SymbolTypeUtils.GetSymbolKind(r.Type),
-                    Range = r.ScriptRegion.ToRange(),
-                    SelectionRange = r.NameRegion.ToRange(),
-                    Name = r.Name
+                    Kind = SymbolTypeUtils.GetSymbolKind(symbolReference.Type),
+                    Range = symbolReference.ScriptRegion.ToRange(),
+                    SelectionRange = symbolReference.NameRegion.ToRange(),
+                    Name = symbolReference.Name,
+                    Children = new List<HierarchicalSymbol>()
                 });
             }
 
             // Short-circuit if we have no symbols.
-            if (symbols.Count == 0)
+            if (hierarchicalSymbols.Count == 0)
             {
                 return s_emptySymbolInformationOrDocumentSymbolContainer;
             }
 
             // Otherwise slowly sort them into a hierarchy.
-            symbols = await SortDocumentSymbols(symbols, cancellationToken).ConfigureAwait(false);
+            hierarchicalSymbols = await SortHierarchicalSymbols(hierarchicalSymbols, cancellationToken).ConfigureAwait(false);
 
             // And finally convert them to the silly SymbolInformationOrDocumentSymbol wrapper.
             List<SymbolInformationOrDocumentSymbol> container = new();
-            foreach (DocumentSymbol symbol in symbols)
+            foreach (DocumentSymbol symbol in GetDocumentSymbolsFromHierarchicalSymbols(hierarchicalSymbols))
             {
                 container.Add(new SymbolInformationOrDocumentSymbol(symbol));
             }