(GH-811) Fix folding with spanning regions

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;
@@ -15,10 +16,26 @@ namespace Microsoft.PowerShell.EditorServices
     /// </summary>
     internal static class TokenOperations
     {
+        // Region kinds to align with VSCode's region kinds
         private const string RegionKindComment = "comment";
         private const string RegionKindRegion = "region";
         private const string RegionKindNone = null;
 
+        // Opening tokens for { } and @{ }
+        private static readonly TokenKind[] s_openingBraces = new []
+        {
+            TokenKind.LCurly,
+            TokenKind.AtCurly
+        };
+
+        // Opening tokens for ( ), @( ), $( )
+        private static readonly TokenKind[] s_openingParens = new []
+        {
+            TokenKind.LParen,
+            TokenKind.AtParen,
+            TokenKind.DollarParen
+        };
+
         /// <summary>
         /// Extracts all of the unique foldable regions in a script given the list tokens
         /// </summary>
@@ -28,29 +45,27 @@ 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 ( -> )
+            // Find matching braces  { -> }
+            // Find matching hashes @{ -> }
             foldableRegions.AddRange(
-                MatchTokenElements(tokens, TokenKind.LParen, TokenKind.RParen, RegionKindNone)
+                MatchTokenElements(tokens, s_openingBraces, TokenKind.RCurly, RegionKindNone)
             );
 
-            // Find matching arrays @( -> )
+            // Find matching parentheses     ( -> )
+            // Find matching array literals @( -> )
+            // Find matching subexpressions $( -> )
             foldableRegions.AddRange(
-                MatchTokenElements(tokens, TokenKind.AtParen, TokenKind.RParen, RegionKindNone)
+                MatchTokenElements(tokens, s_openingParens, TokenKind.RParen, RegionKindNone)
             );
 
-            // Find matching hashes @{ -> }
+            // Find contiguous here strings @' -> '@
             foldableRegions.AddRange(
-                MatchTokenElements(tokens, TokenKind.AtCurly, TokenKind.RParen, RegionKindNone)
+                MatchTokenElement(tokens, TokenKind.HereStringLiteral, RegionKindNone)
             );
 
-            // Find contiguous here strings @' -> '@
+            // Find unopinionated variable names ${ \n \n }
             foldableRegions.AddRange(
-                MatchTokenElement(tokens, TokenKind.HereStringLiteral, RegionKindNone)
+                MatchTokenElement(tokens, TokenKind.Variable, RegionKindNone)
             );
 
             // Find contiguous here strings @" -> "@
@@ -146,19 +161,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

@@ -67,6 +67,9 @@ herestrings should fold
 
 '@
 
+# This won't confuse things
+Get-Command -Param @I
+
 $I = @""
 double quoted herestrings should also fold
 
@@ -111,23 +114,30 @@ double quoted herestrings should also fold
 # Comment Block 2
 # Comment Block 2
 $something = $true
-#endregion Comment Block 3";
+#endregion Comment Block 3
+
+# What about anonymous variable assignment
+${this
+is
+valid} = 5
+";
         private FoldingReference[] expectedAllInOneScriptFolds = {
             CreateFoldingReference(0,   0,  3, 10, "region"),
             CreateFoldingReference(1,   0,  2,  2, "comment"),
             CreateFoldingReference(10,  0, 14,  2, "comment"),
-            CreateFoldingReference(16, 30, 59,  1, null),
+            CreateFoldingReference(16, 30, 62,  1, null),
             CreateFoldingReference(17,  0, 21,  2, "comment"),
             CreateFoldingReference(23,  7, 25,  2, null),
-            CreateFoldingReference(28,  5, 30,  2, null),
-            CreateFoldingReference(35,  2, 36,  0, "comment"),
-            CreateFoldingReference(39,  2, 48, 14, "region"),
-            CreateFoldingReference(41,  4, 44, 14, "region"),
-            CreateFoldingReference(51,  7, 52,  3, null),
-            CreateFoldingReference(56,  7, 58,  3, null),
-            CreateFoldingReference(64,  0, 65,  0, "comment"),
-            CreateFoldingReference(67,  0, 71, 26, "region"),
-            CreateFoldingReference(68,  0, 69,  0, "comment")
+            CreateFoldingReference(31,  5, 33,  2, null),
+            CreateFoldingReference(38,  2, 39,  0, "comment"),
+            CreateFoldingReference(42,  2, 51, 14, "region"),
+            CreateFoldingReference(44,  4, 47, 14, "region"),
+            CreateFoldingReference(54,  7, 55,  3, null),
+            CreateFoldingReference(59,  7, 61,  3, null),
+            CreateFoldingReference(67,  0, 68,  0, "comment"),
+            CreateFoldingReference(70,  0, 74, 26, "region"),
+            CreateFoldingReference(71,  0, 72,  0, "comment"),
+            CreateFoldingReference(78,  0, 79,  6, null),
         };
 
         /// <summary>
@@ -217,5 +227,32 @@ 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);
+        }
     }
 }