{WIP} (GH-813) Use AST for Syntax Folder

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,16 @@ private FoldingRange[] Fold(
             // TODO Should be using dynamic registrations
             if (!this.currentSettings.CodeFolding.Enable) { return null; }
             var result = new List<FoldingRange>();
-            foreach (FoldingReference fold in TokenOperations.FoldableRegions(
-                editorSession.Workspace.GetFile(documentUri).ScriptTokens,
-                this.currentSettings.CodeFolding.ShowLastLine))
+            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 FoldingOperations.FoldableRegions(
+                script.ScriptTokens,
+                script.ScriptAst).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 +1737,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/AstOperations.cs

@@ -330,5 +330,18 @@ 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,152 @@
+//
+// 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 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.FoldableRegions(tokens, ref foldableRegions);
+
+            return foldableRegions;
+        }
+    }
+}

src/PowerShellEditorServices/Language/FoldingReference.cs

@@ -4,6 +4,7 @@
 //
 
 using System;
+using System.Collections.Generic;
 
 namespace Microsoft.PowerShell.EditorServices
 {
@@ -60,4 +61,33 @@ 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; }
+            FoldingReference currentItem;
+            TryGetValue(item.StartLine, out currentItem);
+            // Only add the item if it hasn't been seen before or it's the largest range
+            if ((currentItem == null) || (currentItem.CompareTo(item) == 1)) { 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;
+        }
+    }
 }