diff --git a/src/PowerShellEditorServices.Protocol/Server/LanguageServer.cs b/src/PowerShellEditorServices.Protocol/Server/LanguageServer.cs
index d325b2d65..8a8c414b8 100644
--- a/src/PowerShellEditorServices.Protocol/Server/LanguageServer.cs
+++ b/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,
diff --git a/src/PowerShellEditorServices/Language/AstOperations.cs b/src/PowerShellEditorServices/Language/AstOperations.cs
index 354e5f188..a2bd603c4 100644
--- a/src/PowerShellEditorServices/Language/AstOperations.cs
+++ b/src/PowerShellEditorServices/Language/AstOperations.cs
@@ -339,5 +339,17 @@ static public string[] FindDotSourcedIncludes(Ast scriptAst, string psScriptRoot
return dotSourcedVisitor.DotSourcedFiles.ToArray();
}
+
+ ///
+ /// Finds all foldable regions in a script based on AST
+ ///
+ /// The abstract syntax tree of the given script
+ /// The FoldingReferenceList object to add the folds to
+ /// A collection of FoldingReference objects
+ public static void FindFoldsInDocument(Ast scriptAst, ref FoldingReferenceList refList)
+ {
+ FindFoldsVisitor findFoldsVisitor = new FindFoldsVisitor(ref refList);
+ scriptAst.Visit(findFoldsVisitor);
+ }
}
}
diff --git a/src/PowerShellEditorServices/Language/FindFoldsVisitor.cs b/src/PowerShellEditorServices/Language/FindFoldsVisitor.cs
new file mode 100644
index 000000000..35a784435
--- /dev/null
+++ b/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
+{
+ ///
+ /// The visitor used to find the all folding regions in an AST
+ ///
+ internal class FindFoldsVisitor : AstVisitor2
+ {
+ private const string RegionKindNone = null;
+
+ private FoldingReferenceList _refList;
+
+ public FindFoldsVisitor(ref FoldingReferenceList refList)
+ {
+ _refList = refList;
+ }
+
+ ///
+ /// Returns whether an Extent could be used as a valid folding region
+ ///
+ private bool IsValidFoldingExtent(
+ IScriptExtent extent)
+ {
+ // The extent must span at least one line
+ return extent.EndLineNumber > extent.StartLineNumber;
+ }
+
+ ///
+ /// Creates an instance of a FoldingReference object from a script extent
+ ///
+ 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;
+ }
+ }
+}
diff --git a/src/PowerShellEditorServices/Language/FoldingOperations.cs b/src/PowerShellEditorServices/Language/FoldingOperations.cs
new file mode 100644
index 000000000..9f832235e
--- /dev/null
+++ b/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
+{
+ ///
+ /// Provides common operations for code folding in a script
+ ///
+ internal static class FoldingOperations
+ {
+ ///
+ /// Extracts all of the unique foldable regions in a script given a script AST and the list tokens
+ /// used to generate the AST
+ ///
+ 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;
+ }
+ }
+}
diff --git a/src/PowerShellEditorServices/Language/TokenOperations.cs b/src/PowerShellEditorServices/Language/TokenOperations.cs
index e62dca68c..60f2a2593 100644
--- a/src/PowerShellEditorServices/Language/TokenOperations.cs
+++ b/src/PowerShellEditorServices/Language/TokenOperations.cs
@@ -33,61 +33,10 @@ internal static class TokenOperations
///
/// Extracts all of the unique foldable regions in a script given the list tokens
///
- internal static FoldingReferenceList FoldableReferences(
- Token[] tokens)
+ internal static void FoldableReferences(
+ Token[] tokens,
+ ref FoldingReferenceList refList)
{
- var refList = new FoldingReferenceList();
-
- Stack tokenCurlyStack = new Stack();
- Stack tokenParenStack = new Stack();
- 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;
}
///
diff --git a/test/PowerShellEditorServices.Test/Language/TokenOperationsTests.cs b/test/PowerShellEditorServices.Test/Language/TokenOperationsTests.cs
index 38a62993b..0e81a9682 100644
--- a/test/PowerShellEditorServices.Test/Language/TokenOperationsTests.cs
+++ b/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)