(GH-824)(GH-812) Improve code folding speed

src/PowerShellEditorServices.Protocol/Server/LanguageServer.cs

@@ -535,7 +535,7 @@ private async Task HandleGetCommandRequestAsync(
         {
             PSCommand psCommand = new PSCommand();
             if (!string.IsNullOrEmpty(param))
-            {    
+            {
                 psCommand.AddCommand("Microsoft.PowerShell.Core\\Get-Command").AddArgument(param);
             }
             else
@@ -1267,7 +1267,7 @@ protected async Task HandleCodeActionRequest(
                 }
             }
 
-            // Add "show documentation" commands last so they appear at the bottom of the client UI. 
+            // Add "show documentation" commands last so they appear at the bottom of the client UI.
             // These commands do not require code fixes. Sometimes we get a batch of diagnostics
             // to create commands for. No need to create multiple show doc commands for the same rule.
             var ruleNamesProcessed = new HashSet<string>();
@@ -1390,14 +1390,15 @@ private FoldingRange[] Fold(string documentUri)
             if (!editorSession.Workspace.TryGetFile(documentUri, out scriptFile)) { return null; }
 
             var result = new List<FoldingRange>();
-            FoldingReference[] foldableRegions = 
-                TokenOperations.FoldableRegions(scriptFile.ScriptTokens, this.currentSettings.CodeFolding.ShowLastLine);
 
-            foreach (FoldingReference fold in foldableRegions)
+            // If we're showing the last line, decrement the Endline of all regions by one.
+            int endLineOffset = this.currentSettings.CodeFolding.ShowLastLine ? -1 : 0;
+
+            foreach (FoldingReference fold in TokenOperations.FoldableReferences(scriptFile.ScriptTokens).References)
             {
                 result.Add(new FoldingRange {
                     EndCharacter   = fold.EndCharacter,
-                    EndLine        = fold.EndLine,
+                    EndLine        = fold.EndLine + endLineOffset,
                     Kind           = fold.Kind,
                     StartCharacter = fold.StartCharacter,
                     StartLine      = fold.StartLine
@@ -1744,7 +1745,7 @@ await eventSender(
                 });
         }
 
-        // Generate a unique id that is used as a key to look up the associated code action (code fix) when 
+        // Generate a unique id that is used as a key to look up the associated code action (code fix) when
         // we receive and process the textDocument/codeAction message.
         private static string GetUniqueIdFromDiagnostic(Diagnostic diagnostic)
         {

src/PowerShellEditorServices/Language/FoldingReference.cs

@@ -4,6 +4,7 @@
 //
 
 using System;
+using System.Collections.Generic;
 
 namespace Microsoft.PowerShell.EditorServices
 {
@@ -60,4 +61,52 @@ public int CompareTo(FoldingReference that) {
             return string.Compare(this.Kind, that.Kind);
         }
     }
+
+    /// <summary>
+    /// A class that holds a list of FoldingReferences and ensures that when adding a reference that the
+    /// folding rules are obeyed, e.g. Only one fold per start line
+    /// </summary>
+    public class FoldingReferenceList
+    {
+        private readonly Dictionary<int, FoldingReference> references = new Dictionary<int, FoldingReference>();
+
+        /// <summary>
+        /// Return all references in the list
+        /// </summary>
+        public IEnumerable<FoldingReference> References
+        {
+            get
+            {
+                return references.Values;
+            }
+        }
+
+        /// <summary>
+        /// Adds a FoldingReference to the list and enforces ordering rules e.g. Only one fold per start line
+        /// </summary>
+        public void SafeAdd(FoldingReference item)
+        {
+            if (item == null) { return; }
+
+            // Only add the item if it hasn't been seen before or it's the largest range
+            if (references.TryGetValue(item.StartLine, out FoldingReference currentItem))
+            {
+                if (currentItem.CompareTo(item) == 1) { references[item.StartLine] = item; }
+            }
+            else
+            {
+                references[item.StartLine] = item;
+            }
+        }
+
+        /// <summary>
+        /// Helper method to easily convert the Dictionary Values into an array
+        /// </summary>
+        public FoldingReference[] ToArray()
+        {
+            var result = new FoldingReference[references.Count];
+            references.Values.CopyTo(result, 0);
+            return result;
+        }
+    }
 }