(GH-813) Use AST for code folding

src/PowerShellEditorServices.Protocol/Server/LanguageServer.cs

@@ -1388,7 +1388,9 @@ private FoldingRange[] Fold(string 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.FoldableReferences(scriptFile.ScriptTokens).References)
+            foreach (FoldingReference fold in FoldingOperations.FoldableRegions(
+                scriptFile.ScriptTokens,
+                scriptFile.ScriptAst).References)
             {
                 result.Add(new FoldingRange {
                     EndCharacter   = fold.EndCharacter,

src/PowerShellEditorServices/Language/AstOperations.cs

@@ -339,5 +339,17 @@ static public string[] FindDotSourcedIncludes(Ast scriptAst, string psScriptRoot
 
             return dotSourcedVisitor.DotSourcedFiles.ToArray();
         }
+
+        /// <summary>
+        /// Finds all foldable regions in a script based on AST
+        /// </summary>
+        /// <param name="scriptAst">The abstract syntax tree of the given script</param>
+        /// <param name="refList">The FoldingReferenceList object to add the folds to</param>
+        /// <returns>A collection of FoldingReference objects</returns>
+        public static void FindFoldsInDocument(Ast scriptAst, ref FoldingReferenceList refList)
+        {
+            FindFoldsVisitor findFoldsVisitor = new FindFoldsVisitor(ref refList);
+            scriptAst.Visit(findFoldsVisitor);
+        }
     }
 }

src/PowerShellEditorServices/Language/FindFoldsVisitor.cs

@@ -0,0 +1,158 @@
+//
+// Copyright (c) Microsoft. All rights reserved.
+// 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;
+
+namespace Microsoft.PowerShell.EditorServices
+{
+    /// <summary>
+    /// The visitor used to find the all folding regions in an AST
+    /// </summary>
+    internal class FindFoldsVisitor : AstVisitor2
+    {
+        private const string RegionKindNone = null;
+
+        private FoldingReferenceList _refList;
+
+        public FindFoldsVisitor(ref FoldingReferenceList refList)
+        {
+            _refList = refList;
+        }
+
+        /// <summary>
+        /// Returns whether an Extent could be used as a valid folding region
+        /// </summary>
+        private bool IsValidFoldingExtent(
+            IScriptExtent extent)
+        {
+            // The extent must span at least one line
+            return extent.EndLineNumber > extent.StartLineNumber;
+        }
+
+        /// <summary>
+        /// Creates an instance of a FoldingReference object from a script extent
+        /// </summary>
+        private FoldingReference CreateFoldingReference(
+            IScriptExtent extent,
+            string matchKind)
+        {
+            // Extents are base 1, but LSP is base 0, so minus 1 off all lines and character positions
+            return new FoldingReference
+            {
+                StartLine      = extent.StartLineNumber - 1,
+                StartCharacter = extent.StartColumnNumber - 1,
+                EndLine        = extent.EndLineNumber - 1,
+                EndCharacter   = extent.EndColumnNumber - 1,
+                Kind           = matchKind
+            };
+        }
+
+        // AST object visitor methods
+        public override AstVisitAction VisitArrayExpression(ArrayExpressionAst objAst)
+        {
+            if (IsValidFoldingExtent(objAst.Extent))
+            {
+                _refList.SafeAdd(CreateFoldingReference(objAst.Extent, RegionKindNone));
+            }
+            return AstVisitAction.Continue;
+        }
+
+        public override AstVisitAction VisitHashtable(HashtableAst objAst)
+        {
+            if (IsValidFoldingExtent(objAst.Extent))
+            {
+                _refList.SafeAdd(CreateFoldingReference(objAst.Extent, RegionKindNone));
+            }
+            return AstVisitAction.Continue;
+        }
+
+        public override AstVisitAction VisitParamBlock(ParamBlockAst objAst)
+        {
+            if (IsValidFoldingExtent(objAst.Extent)) { _refList.SafeAdd(CreateFoldingReference(objAst.Extent, RegionKindNone)); }
+            return AstVisitAction.Continue;
+        }
+
+        public override AstVisitAction VisitStatementBlock(StatementBlockAst objAst)
+        {
+            // These parent visitors will get this AST Object.  No need to process it
+            if (objAst.Parent == null) { return AstVisitAction.Continue; }
+            if (objAst.Parent is ArrayExpressionAst) { return AstVisitAction.Continue; }
+            if (IsValidFoldingExtent(objAst.Extent))
+            {
+                _refList.SafeAdd(CreateFoldingReference(objAst.Extent, RegionKindNone));
+            }
+            return AstVisitAction.Continue;
+        }
+
+        public override AstVisitAction VisitScriptBlock(ScriptBlockAst objAst)
+        {
+            // If the Parent object is null then this represents the entire script.  We don't want to fold that
+            if (objAst.Parent == null) { return AstVisitAction.Continue; }
+            // The ScriptBlockExpressionAst visitor will get this AST Object.  No need to process it
+            if (objAst.Parent is ScriptBlockExpressionAst) { return AstVisitAction.Continue; }
+            if (IsValidFoldingExtent(objAst.Extent))
+            {
+                _refList.SafeAdd(CreateFoldingReference(objAst.Extent, RegionKindNone));
+            }
+            return AstVisitAction.Continue;
+        }
+
+        public override AstVisitAction VisitScriptBlockExpression(ScriptBlockExpressionAst objAst)
+        {
+            if (IsValidFoldingExtent(objAst.Extent)) {
+                FoldingReference foldRef = CreateFoldingReference(objAst.ScriptBlock.Extent, RegionKindNone);
+                if (objAst.Parent == null) { return AstVisitAction.Continue; }
+                if (objAst.Parent is InvokeMemberExpressionAst) {
+                    // This is a bit naive.  The ScriptBlockExpressionAst Extent does not include the actual { and }
+                    // characters so the StartCharacter and EndCharacter indexes are out by one.  This could be a bug in
+                    // PowerShell Parser. This is just a workaround
+                    foldRef.StartCharacter--;
+                    foldRef.EndCharacter++;
+                }
+                _refList.SafeAdd(foldRef);
+            }
+            return AstVisitAction.Continue;
+        }
+
+        public override AstVisitAction VisitStringConstantExpression(StringConstantExpressionAst objAst)
+        {
+            if (IsValidFoldingExtent(objAst.Extent))
+            {
+                _refList.SafeAdd(CreateFoldingReference(objAst.Extent, RegionKindNone));
+            }
+
+            return AstVisitAction.Continue;
+        }
+
+        public override AstVisitAction VisitSubExpression(SubExpressionAst objAst)
+        {
+            if (IsValidFoldingExtent(objAst.Extent))
+            {
+                _refList.SafeAdd(CreateFoldingReference(objAst.Extent, RegionKindNone));
+            }
+            return AstVisitAction.Continue;
+        }
+
+        public override AstVisitAction VisitTypeDefinition(TypeDefinitionAst objAst)
+        {
+            if (IsValidFoldingExtent(objAst.Extent))
+            {
+                _refList.SafeAdd(CreateFoldingReference(objAst.Extent, RegionKindNone));
+            }
+            return AstVisitAction.Continue;
+        }
+
+        public override AstVisitAction VisitVariableExpression(VariableExpressionAst objAst)
+        {
+            if (IsValidFoldingExtent(objAst.Extent))
+            {
+                _refList.SafeAdd(CreateFoldingReference(objAst.Extent, RegionKindNone));
+            }
+            return AstVisitAction.Continue;
+        }
+    }
+}

src/PowerShellEditorServices/Language/FoldingOperations.cs

@@ -0,0 +1,36 @@
+//
+// Copyright (c) Microsoft. All rights reserved.
+// 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;
+
+namespace Microsoft.PowerShell.EditorServices
+{
+    /// <summary>
+    /// Provides common operations for code folding in a script
+    /// </summary>
+    internal static class FoldingOperations
+    {
+        /// <summary>
+        /// Extracts all of the unique foldable regions in a script given a script AST and the list tokens
+        /// used to generate the AST
+        /// </summary>
+        internal static FoldingReferenceList FoldableRegions(
+            Token[] tokens,
+            Ast scriptAst)
+        {
+            var foldableRegions = new FoldingReferenceList();
+
+            // Add regions from AST
+            AstOperations.FindFoldsInDocument(scriptAst, ref foldableRegions);
+
+            // Add regions from Tokens
+            TokenOperations.FoldableReferences(tokens, ref foldableRegions);
+
+            return foldableRegions;
+        }
+    }
+}

src/PowerShellEditorServices/Language/TokenOperations.cs

@@ -33,61 +33,10 @@ internal static class TokenOperations
         /// <summary>
         /// Extracts all of the unique foldable regions in a script given the list tokens
         /// </summary>
-        internal static FoldingReferenceList FoldableReferences(
-            Token[] tokens)
+        internal static void FoldableReferences(
+            Token[] tokens,
+            ref FoldingReferenceList refList)
         {
-            var refList = new FoldingReferenceList();
-
-            Stack<Token> tokenCurlyStack = new Stack<Token>();
-            Stack<Token> tokenParenStack = new Stack<Token>();
-            foreach (Token token in tokens)
-            {
-                switch (token.Kind)
-                {
-                    // Find matching braces  { -> }
-                    // Find matching hashes @{ -> }
-                    case TokenKind.LCurly:
-                    case TokenKind.AtCurly:
-                        tokenCurlyStack.Push(token);
-                        break;
-
-                    case TokenKind.RCurly:
-                        if (tokenCurlyStack.Count > 0)
-                        {
-                            refList.SafeAdd(CreateFoldingReference(tokenCurlyStack.Pop(), token, RegionKindNone));
-                        }
-                        break;
-
-                    // Find matching parentheses     ( -> )
-                    // Find matching array literals @( -> )
-                    // Find matching subexpressions $( -> )
-                    case TokenKind.LParen:
-                    case TokenKind.AtParen:
-                    case TokenKind.DollarParen:
-                        tokenParenStack.Push(token);
-                        break;
-
-                    case TokenKind.RParen:
-                        if (tokenParenStack.Count > 0)
-                        {
-                            refList.SafeAdd(CreateFoldingReference(tokenParenStack.Pop(), token, RegionKindNone));
-                        }
-                        break;
-
-                    // Find contiguous here strings @' -> '@
-                    // Find unopinionated variable names ${ \n \n }
-                    // Find contiguous expandable here strings @" -> "@
-                    case TokenKind.HereStringLiteral:
-                    case TokenKind.Variable:
-                    case TokenKind.HereStringExpandable:
-                        if (token.Extent.StartLineNumber != token.Extent.EndLineNumber)
-                        {
-                            refList.SafeAdd(CreateFoldingReference(token, token, RegionKindNone));
-                        }
-                        break;
-                }
-            }
-
             // Find matching comment regions   #region -> #endregion
             // 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
@@ -153,8 +102,6 @@ internal static FoldingReferenceList FoldableReferences(
             {
                 refList.SafeAdd(CreateFoldingReference(blockStartToken, blockNextLine - 1, RegionKindComment));
             }
-
-            return refList;
         }
 
         /// <summary>

test/PowerShellEditorServices.Test/Language/TokenOperationsTests.cs

@@ -20,7 +20,9 @@ private FoldingReference[] GetRegions(string text) {
                 text,
                 Version.Parse("5.0"));
 
-            var result = Microsoft.PowerShell.EditorServices.TokenOperations.FoldableReferences(scriptFile.ScriptTokens).ToArray();
+            var result = Microsoft.PowerShell.EditorServices.FoldingOperations.FoldableRegions(
+                scriptFile.ScriptTokens,
+                scriptFile.ScriptAst).ToArray();
             // The foldable regions need to be deterministic for testing so sort the array.
             Array.Sort(result);
             return result;
@@ -256,6 +258,7 @@ public void LaguageServiceFindsFoldablRegionsWithSameEndToken() {
         }
 
         // A simple PowerShell Classes test
+        [Trait("Category", "Folding")]
         [Fact]
         public void LaguageServiceFindsFoldablRegionsWithClasses() {
             string testString =
@@ -271,7 +274,7 @@ [string] TestMethod() {
 }
 ";
             FoldingReference[] expectedFolds = {
-                CreateFoldingReference(0, 16, 9,  1, null),
+                CreateFoldingReference(0, 0, 9,  1, null),
                 CreateFoldingReference(1, 31, 4, 16, null),
                 CreateFoldingReference(6, 26, 8,  5, null)
             };
@@ -282,6 +285,7 @@ [string] TestMethod() {
         }
 
         // This tests DSC style keywords and param blocks
+        [Trait("Category", "Folding")]
         [Fact]
         public void LaguageServiceFindsFoldablRegionsWithDSC() {
             string testString =
@@ -322,7 +326,7 @@ AdcsCertificationAuthority CertificateAuthority
 ";
             FoldingReference[] expectedFolds = {
                 CreateFoldingReference(1,  0, 33, 1, null),
-                CreateFoldingReference(3,  4, 12, 5, null),
+                CreateFoldingReference(2,  4, 12, 5, null),
                 CreateFoldingReference(17, 4, 32, 5, null),
                 CreateFoldingReference(19, 8, 22, 9, null),
                 CreateFoldingReference(25, 8, 31, 9, null)