(GH-1631) Fix folding with spanning regions

Author: glennsarti

src/PowerShellEditorServices/Language/TokenOperations.cs

@@ -3,6 +3,7 @@
 // Licensed under the MIT license. See LICENSE file in the project root for full license information.
 //
 
+using System;
 using System.Collections.Generic;
 using System.Management.Automation.Language;
 using System.Text.RegularExpressions;
@@ -28,24 +29,17 @@ internal static FoldingReference[] FoldableRegions(
         {
             List<FoldingReference> foldableRegions = new List<FoldingReference>();
 
-            // Find matching braces { -> }
-            foldableRegions.AddRange(
-                MatchTokenElements(tokens, TokenKind.LCurly, TokenKind.RCurly, RegionKindNone)
-            );
-
-            // Find matching braces ( -> )
-            foldableRegions.AddRange(
-                MatchTokenElements(tokens, TokenKind.LParen, TokenKind.RParen, RegionKindNone)
-            );
-
-            // Find matching arrays @( -> )
+            // Find matching braces  { -> }
+            // Find matching hashes @{ -> }
             foldableRegions.AddRange(
-                MatchTokenElements(tokens, TokenKind.AtParen, TokenKind.RParen, RegionKindNone)
+                MatchTokenElements(tokens, new TokenKind[] { TokenKind.LCurly, TokenKind.AtCurly }, TokenKind.RCurly, RegionKindNone)
             );
 
-            // Find matching hashes @{ -> }
+            // Find matching parentheses     ( -> )
+            // Find matching array literals @( -> )
+            // Find matching subexpressions $( -> )
             foldableRegions.AddRange(
-                MatchTokenElements(tokens, TokenKind.AtCurly, TokenKind.RParen, RegionKindNone)
+                MatchTokenElements(tokens, new TokenKind[] { TokenKind.LParen, TokenKind.AtParen, TokenKind.DollarParen }, TokenKind.RParen, RegionKindNone)
             );
 
             // Find contiguous here strings @' -> '@
@@ -146,19 +140,19 @@ static private FoldingReference CreateFoldingReference(
         }
 
         /// <summary>
-        /// Given an array of tokens, find matching regions which start and end with a different TokenKind
+        /// 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(
             Token[] tokens,
-            TokenKind startTokenKind,
+            TokenKind[] startTokenKind,
             TokenKind endTokenKind,
             string matchKind)
         {
             List<FoldingReference> result = new List<FoldingReference>();
             Stack<Token> tokenStack = new Stack<Token>();
             foreach (Token token in tokens)
             {
-                if (token.Kind == startTokenKind) {
+                if (Array.IndexOf(startTokenKind, token.Kind) != -1) {
                     tokenStack.Push(token);
                 }
                 if ((tokenStack.Count > 0) && (token.Kind == endTokenKind)) {

test/PowerShellEditorServices.Test/Language/TokenOperationsTests.cs

@@ -217,5 +217,33 @@ public void LaguageServiceFindsFoldablRegionsWithDuplicateRegions() {
             FoldingReference[] result = GetRegions(testString);
             AssertFoldingReferenceArrays(expectedFolds, result);
         }
+
+        // This tests that token matching { -> }, @{ -> } and
+        // ( -> ), @( -> ) and $( -> ) does not confuse the folder
+        [Fact]
+        public void LaguageServiceFindsFoldablRegionsWithSameEndToken() {
+            string testString =
+@"foreach ($1 in $2) {
+
+    $x = @{
+        'abc' = 'def'
+    }
+}
+
+$y = $(
+    $arr = @('1', '2'); Write-Host ($arr)
+)
+";
+            FoldingReference[] expectedFolds = {
+                CreateFoldingReference(0, 19, 4, 1, null),
+                CreateFoldingReference(2,  9, 3, 5, null),
+                CreateFoldingReference(7,  5, 8, 1, null)
+            };
+
+            FoldingReference[] result = GetRegions(testString);
+
+            AssertFoldingReferenceArrays(expectedFolds, result);
+        }
+
     }
 }