{WIP} (GH-812) Update folder for DSC style scripts

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,17 @@ 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);
+            int endLineOffset = 0;
+            // If we're showing the last line, decrement the Endline of all regions by one.
+            if (this.currentSettings.CodeFolding.ShowLastLine) { endLineOffset = -1; }
+            foreach (FoldingReference fold in FoldingOperations.FoldableRegions(
+                script.ScriptTokens,
+                script.ScriptAst))
             {
                 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 +1738,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,19 @@ 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>
+        /// <returns>A collection of FoldingReference objects</returns>
+        public static IEnumerable<FoldingReference> FindFoldsInDocument(Ast scriptAst)
+        {
+            FindFoldsVisitor findFoldsVisitor = new FindFoldsVisitor();
+            scriptAst.Visit(findFoldsVisitor);
+
+            return findFoldsVisitor.FoldableRegions;
+        }
+
     }
 }

src/PowerShellEditorServices/Language/FindFoldsVisitor.cs

@@ -0,0 +1,148 @@
+//
+// 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 : AstVisitor
+    {
+        private const string RegionKindNone = null;
+
+        public List<FoldingReference> FoldableRegions { get; }
+
+        public FindFoldsVisitor()
+        {
+            this.FoldableRegions = new List<FoldingReference>();
+        }
+
+        /// <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))
+            {
+                this.FoldableRegions.Add(CreateFoldingReference(objAst.Extent, RegionKindNone));
+            }
+            return AstVisitAction.Continue;
+        }
+
+        public override AstVisitAction VisitHashtable(HashtableAst objAst)
+        {
+            if (IsValidFoldingExtent(objAst.Extent))
+            {
+                this.FoldableRegions.Add(CreateFoldingReference(objAst.Extent, RegionKindNone));
+            }
+            return AstVisitAction.Continue;
+        }
+
+        public override AstVisitAction VisitParamBlock(ParamBlockAst objAst)
+        {
+            if (IsValidFoldingExtent(objAst.Extent)) { this.FoldableRegions.Add(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))
+            {
+                this.FoldableRegions.Add(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))
+            {
+                this.FoldableRegions.Add(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++;
+                }
+                this.FoldableRegions.Add(foldRef);
+            }
+            return AstVisitAction.Continue;
+        }
+
+        public override AstVisitAction VisitStringConstantExpression(StringConstantExpressionAst objAst)
+        {
+            if (IsValidFoldingExtent(objAst.Extent))
+            {
+                this.FoldableRegions.Add(CreateFoldingReference(objAst.Extent, RegionKindNone));
+            }
+
+            return AstVisitAction.Continue;
+        }
+
+        public override AstVisitAction VisitSubExpression(SubExpressionAst objAst)
+        {
+            if (IsValidFoldingExtent(objAst.Extent))
+            {
+                this.FoldableRegions.Add(CreateFoldingReference(objAst.Extent, RegionKindNone));
+            }
+            return AstVisitAction.Continue;
+        }
+
+        public override AstVisitAction VisitVariableExpression(VariableExpressionAst objAst)
+        {
+            if (IsValidFoldingExtent(objAst.Extent))
+            {
+                this.FoldableRegions.Add(CreateFoldingReference(objAst.Extent, RegionKindNone));
+            }
+            return AstVisitAction.Continue;
+        }
+    }
+}

src/PowerShellEditorServices/Language/FoldingOperations.cs

@@ -0,0 +1,55 @@
+//
+// 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 FoldingReference[] FoldableRegions(
+            Token[] tokens,
+            Ast scriptAst)
+        {
+            List<FoldingReference> foldableRegions = new List<FoldingReference>();
+
+            // Add regions from AST
+            foldableRegions.AddRange(AstOperations.FindFoldsInDocument(scriptAst));
+
+            // Add regions from Tokens
+            foldableRegions.AddRange(TokenOperations.FoldableRegions(tokens));
+
+            // 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);
+            });
+
+            return foldableRegions.ToArray();
+        }
+    }
+}