(GH-824) Refactor the FoldingReference arrays and lists into it's own class

Previously the folding provider created many intermediate arrays and lists and
required post-processing.  This commit changes the behaviour to use an
accumlator patter with an extended Dictionary class.  This new class adds a
`SafeAdd` method to add FoldingRanges, which then has the logic to determine if
the range should indeed be added, for example, passing nulls or pre-existing
larger ranges.

By passing around this list using ByReference we can avoid creating many objects
which are just then thrown away.

This commit also moves the ShowLastLine code from the FoldingProvider into the
Language Server.  This reduces the number of array enumerations to one.

Author: glennsarti

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>();
@@ -1382,13 +1382,18 @@ private FoldingRange[] Fold(
             // TODO Should be using dynamic registrations
             if (!this.currentSettings.CodeFolding.Enable) { return null; }
             var result = new List<FoldingRange>();
+
+            ScriptFile script = editorSession.Workspace.GetFile(documentUri);
+
+            // 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.FoldableRegions(
-                editorSession.Workspace.GetFile(documentUri).ScriptTokens,
-                this.currentSettings.CodeFolding.ShowLastLine))
+                script.ScriptTokens).ToArray())
             {
                 result.Add(new FoldingRange {
                     EndCharacter   = fold.EndCharacter,
-                    EndLine        = fold.EndLine,
+                    EndLine        = fold.EndLine + endLineOffset,
                     Kind           = fold.Kind,
                     StartCharacter = fold.StartCharacter,
                     StartLine      = fold.StartLine
@@ -1734,7 +1739,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,39 @@ 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 : Dictionary<int, FoldingReference>
+    {
+        /// <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 (TryGetValue(item.StartLine, out FoldingReference currentItem))
+            {
+                if (currentItem.CompareTo(item) == 1) { this[item.StartLine] = item; }
+            }
+            else
+            {
+                this[item.StartLine] = item;
+            }
+        }
+
+        /// <summary>
+        /// Helper method to easily convert the Dictionary Values into an array
+        /// </summary>
+        public FoldingReference[] ToArray()
+        {
+            var result = new FoldingReference[Count];
+            Values.CopyTo(result, 0);
+            return result;
+        }
+    }
 }

src/PowerShellEditorServices/Language/TokenOperations.cs

@@ -39,85 +39,39 @@ internal static class TokenOperations
         /// <summary>
         /// Extracts all of the unique foldable regions in a script given the list tokens
         /// </summary>
-        internal static FoldingReference[] FoldableRegions(
-            Token[] tokens,
-            bool ShowLastLine)
+        internal static FoldingReferenceList FoldableRegions(
+            Token[] tokens)
         {
-            List<FoldingReference> foldableRegions = new List<FoldingReference>();
+            var refList = new FoldingReferenceList();
 
             // Find matching braces  { -> }
             // Find matching hashes @{ -> }
-            foldableRegions.AddRange(
-                MatchTokenElements(tokens, s_openingBraces, TokenKind.RCurly, RegionKindNone)
-            );
+            MatchTokenElements(tokens, s_openingBraces, TokenKind.RCurly, RegionKindNone, ref refList);
 
             // Find matching parentheses     ( -> )
             // Find matching array literals @( -> )
             // Find matching subexpressions $( -> )
-            foldableRegions.AddRange(
-                MatchTokenElements(tokens, s_openingParens, TokenKind.RParen, RegionKindNone)
-            );
+            MatchTokenElements(tokens, s_openingParens, TokenKind.RParen, RegionKindNone, ref refList);
 
             // Find contiguous here strings @' -> '@
-            foldableRegions.AddRange(
-                MatchTokenElement(tokens, TokenKind.HereStringLiteral, RegionKindNone)
-            );
+            MatchTokenElement(tokens, TokenKind.HereStringLiteral, RegionKindNone, ref refList);
 
             // Find unopinionated variable names ${ \n \n }
-            foldableRegions.AddRange(
-                MatchTokenElement(tokens, TokenKind.Variable, RegionKindNone)
-            );
+            MatchTokenElement(tokens, TokenKind.Variable, RegionKindNone, ref refList);
 
             // Find contiguous here strings @" -> "@
-            foldableRegions.AddRange(
-                MatchTokenElement(tokens, TokenKind.HereStringExpandable, RegionKindNone)
-            );
+            MatchTokenElement(tokens, TokenKind.HereStringExpandable, RegionKindNone, ref refList);
 
             // Find matching comment regions   #region -> #endregion
-            foldableRegions.AddRange(
-                MatchCustomCommentRegionTokenElements(tokens, RegionKindRegion)
-            );
+            MatchCustomCommentRegionTokenElements(tokens, RegionKindRegion, ref refList);
 
             // Find blocks of line comments # comment1\n# comment2\n...
-            foldableRegions.AddRange(
-                MatchBlockCommentTokenElement(tokens, RegionKindComment)
-            );
+            MatchBlockCommentTokenElement(tokens, RegionKindComment, ref refList);
 
             // Find comments regions <# -> #>
-            foldableRegions.AddRange(
-                MatchTokenElement(tokens, TokenKind.Comment, RegionKindComment)
-            );
-
-            // Remove any null entries. Nulls appear if the folding reference is invalid
-            // or missing
-            foldableRegions.RemoveAll(item => item == null);
-
-            // Sort the FoldingReferences, starting at the top of the document,
-            // and ensure that, in the case of multiple ranges starting the same line,
-            // that the largest range (i.e. most number of lines spanned) is sorted
-            // first. This is needed to detect duplicate regions. The first in the list
-            // will be used and subsequent duplicates ignored.
-            foldableRegions.Sort();
-
-            // It's possible to have duplicate or overlapping ranges, that is, regions which have the same starting
-            // line number as the previous region. Therefore only emit ranges which have a different starting line
-            // than the previous range.
-            foldableRegions.RemoveAll( (FoldingReference item) => {
-                // Note - I'm not happy with searching here, but as the RemoveAll
-                // doesn't expose the index in the List, we need to calculate it. Fortunately the
-                // list is sorted at this point, so we can use BinarySearch.
-                int index = foldableRegions.BinarySearch(item);
-                if (index == 0) { return false; }
-                return (item.StartLine == foldableRegions[index - 1].StartLine);
-            });
-
-            // Some editors have different folding UI, sometimes the lastline should be displayed
-            // If we do want to show the last line, just change the region to be one line less
-            if (ShowLastLine) {
-                foldableRegions.ForEach( item => { item.EndLine--; });
-            }
+            MatchTokenElement(tokens, TokenKind.Comment, RegionKindComment, ref refList);
 
-            return foldableRegions.ToArray();
+            return refList;
         }
 
         /// <summary>
@@ -163,42 +117,40 @@ static private FoldingReference CreateFoldingReference(
         /// <summary>
         /// Given an array of tokens, find matching regions which start (array of tokens) and end with a different TokenKind
         /// </summary>
-        static private List<FoldingReference> MatchTokenElements(
+        static private void MatchTokenElements(
             Token[] tokens,
             TokenKind[] startTokenKind,
             TokenKind endTokenKind,
-            string matchKind)
+            string matchKind,
+            ref FoldingReferenceList refList)
         {
-            List<FoldingReference> result = new List<FoldingReference>();
             Stack<Token> tokenStack = new Stack<Token>();
             foreach (Token token in tokens)
             {
                 if (Array.IndexOf(startTokenKind, token.Kind) != -1) {
                     tokenStack.Push(token);
                 }
                 if ((tokenStack.Count > 0) && (token.Kind == endTokenKind)) {
-                    result.Add(CreateFoldingReference(tokenStack.Pop(), token, matchKind));
+                    refList.SafeAdd(CreateFoldingReference(tokenStack.Pop(), token, matchKind));
                 }
             }
-            return result;
         }
 
         /// <summary>
         /// Given an array of token, finds a specific token
         /// </summary>
-        static private List<FoldingReference> MatchTokenElement(
+        static private void MatchTokenElement(
             Token[] tokens,
             TokenKind tokenKind,
-            string matchKind)
+            string matchKind,
+            ref FoldingReferenceList refList)
         {
-            List<FoldingReference> result = new List<FoldingReference>();
             foreach (Token token in tokens)
             {
                 if ((token.Kind == tokenKind) && (token.Extent.StartLineNumber != token.Extent.EndLineNumber)) {
-                    result.Add(CreateFoldingReference(token, token, matchKind));
+                    refList.SafeAdd(CreateFoldingReference(token, token, matchKind));
                 }
             }
-            return result;
         }
 
         /// <summary>
@@ -230,11 +182,11 @@ static private bool IsBlockComment(int index, Token[] tokens) {
         /// classed as comments.  To workaround this we search for valid block comments (See IsBlockCmment)
         /// and then determine contiguous line numbers from there
         /// </summary>
-        static private List<FoldingReference> MatchBlockCommentTokenElement(
+        static private void MatchBlockCommentTokenElement(
             Token[] tokens,
-            string matchKind)
+            string matchKind,
+             ref FoldingReferenceList refList)
         {
-            List<FoldingReference> result = new List<FoldingReference>();
             Token startToken = null;
             int nextLine = -1;
             for (int index = 0; index < tokens.Length; index++)
@@ -243,7 +195,7 @@ static private List<FoldingReference> MatchBlockCommentTokenElement(
                 if (IsBlockComment(index, tokens) && s_nonRegionLineCommentRegex.IsMatch(thisToken.Text)) {
                     int thisLine = thisToken.Extent.StartLineNumber - 1;
                     if ((startToken != null) && (thisLine != nextLine)) {
-                        result.Add(CreateFoldingReference(startToken, nextLine - 1, matchKind));
+                        refList.SafeAdd(CreateFoldingReference(startToken, nextLine - 1, matchKind));
                         startToken = thisToken;
                     }
                     if (startToken == null) { startToken = thisToken; }
@@ -253,27 +205,26 @@ static private List<FoldingReference> MatchBlockCommentTokenElement(
             // If we exit the token array and we're still processing comment lines, then the
             // comment block simply ends at the end of document
             if (startToken != null) {
-                result.Add(CreateFoldingReference(startToken, nextLine - 1, matchKind));
+                refList.SafeAdd(CreateFoldingReference(startToken, nextLine - 1, matchKind));
             }
-            return result;
         }
 
         /// <summary>
         /// Given a list of tokens, find the tokens that are comments and
         /// the comment text is either `# region` or `# endregion`, and then use a stack to determine
         /// the ranges they span
         /// </summary>
-        static private List<FoldingReference> MatchCustomCommentRegionTokenElements(
+        static private void MatchCustomCommentRegionTokenElements(
             Token[] tokens,
-            string matchKind)
+            string matchKind,
+            ref FoldingReferenceList refList)
         {
             // These regular expressions are used to match lines which mark the start and end of region comment in a PowerShell
             // script. They are based on the defaults in the VS Code Language Configuration at;
             // https://github.com/Microsoft/vscode/blob/64186b0a26/extensions/powershell/language-configuration.json#L26-L31
             string startRegionText = @"^\s*#region\b";
             string endRegionText = @"^\s*#endregion\b";
 
-            List<FoldingReference> result = new List<FoldingReference>();
             Stack<Token> tokenStack = new Stack<Token>();
             for (int index = 0; index < tokens.Length; index++)
             {
@@ -283,11 +234,10 @@ static private List<FoldingReference> MatchCustomCommentRegionTokenElements(
                         tokenStack.Push(token);
                     }
                     if ((tokenStack.Count > 0) && (Regex.IsMatch(token.Text, endRegionText, RegexOptions.IgnoreCase))) {
-                        result.Add(CreateFoldingReference(tokenStack.Pop(), token, matchKind));
+                        refList.SafeAdd(CreateFoldingReference(tokenStack.Pop(), token, matchKind));
                     }
                 }
             }
-            return result;
         }
     }
 }