Index brace and pipeline lookups during formatting

Author: PrzemyslawKlys

Engine/TokenOperations.cs

@@ -122,28 +122,37 @@ private IEnumerable<Token> GetBraceInCommandElement(TokenKind tokenKind)
                 yield break;
             }
 
-            Func<Token, Ast, bool> predicate;
-
             switch (tokenKind)
             {
                 case TokenKind.LCurly:
-                    predicate = (x, cmdElemAst) =>
-                        x.Kind == TokenKind.LCurly && x.Extent.StartOffset == cmdElemAst.Extent.StartOffset;
-                    break;
-
                 case TokenKind.RCurly:
-                    predicate = (x, cmdElemAst) =>
-                        x.Kind == TokenKind.RCurly && x.Extent.EndOffset == cmdElemAst.Extent.EndOffset;
                     break;
 
                 default:
                     throw new ArgumentException("", nameof(tokenKind));
             }
 
+            // Index once rather than scanning every token for each script block.
+            // Preserve the first matching token if callers supply duplicate offsets.
+            var braceTokens = new Dictionary<int, Token>();
+            foreach (var token in tokens)
+            {
+                if (token.Kind != tokenKind)
+                {
+                    continue;
+                }
+
+                int offset = tokenKind == TokenKind.LCurly ? token.Extent.StartOffset : token.Extent.EndOffset;
+                if (!braceTokens.ContainsKey(offset))
+                {
+                    braceTokens.Add(offset, token);
+                }
+            }
+
             foreach (var cmdElemAst in cmdElemAsts)
             {
-                var tokenFound = tokens.FirstOrDefault(token => predicate(token, cmdElemAst));
-                if (tokenFound != null)
+                int offset = tokenKind == TokenKind.LCurly ? cmdElemAst.Extent.StartOffset : cmdElemAst.Extent.EndOffset;
+                if (braceTokens.TryGetValue(offset, out var tokenFound))
                 {
                     yield return tokenFound;
                 }

Rules/UseConsistentIndentation.cs

@@ -133,13 +133,23 @@ public override IEnumerable<DiagnosticRecord> AnalyzeScript(Ast ast, string file
             var onNewLine = true;
             var pipelineAsts = ast.FindAll(testAst => testAst is PipelineAst && (testAst as PipelineAst).PipelineElements.Count > 1, true).ToList();
             // Sort by end position so that inner (nested) pipelines appear before outer ones.
-            // This is required by MatchingPipelineAstEnd, whose early-break optimization
-            // would otherwise skip nested pipelines that end before their outer pipeline.
+            // Preserve this order when selecting the first matching pipeline on each line.
             pipelineAsts.Sort((a, b) =>
             {
                 int lineCmp = a.Extent.EndScriptPosition.LineNumber.CompareTo(b.Extent.EndScriptPosition.LineNumber);
                 return lineCmp != 0 ? lineCmp : a.Extent.EndScriptPosition.ColumnNumber.CompareTo(b.Extent.EndScriptPosition.ColumnNumber);
             });
+            var pipelineEndsByLine = new Dictionary<int, List<Ast>>();
+            foreach (var pipeline in pipelineAsts)
+            {
+                int endLine = pipeline.Extent.EndScriptPosition.LineNumber;
+                if (!pipelineEndsByLine.TryGetValue(endLine, out var pipelinesOnLine))
+                {
+                    pipelinesOnLine = new List<Ast>();
+                    pipelineEndsByLine.Add(endLine, pipelinesOnLine);
+                }
+                pipelinesOnLine.Add(pipeline);
+            }
             // Track pipeline indentation increases per PipelineAst instead of as a single
             // flat counter. A flat counter caused all accumulated pipeline indentation to be
             // subtracted when *any* pipeline ended, instead of only the contribution from
@@ -306,7 +316,7 @@ public override IEnumerable<DiagnosticRecord> AnalyzeScript(Ast ast, string file
                 if (pipelineIndentationStyle == PipelineIndentationStyle.None) { continue; }
 
                 // Check if the current token matches the end of a PipelineAst
-                PipelineAst matchingPipeLineAstEnd = MatchingPipelineAstEnd(pipelineAsts, token);
+                PipelineAst matchingPipeLineAstEnd = MatchingPipelineAstEnd(pipelineEndsByLine, token);
                 if (matchingPipeLineAstEnd == null)
                 {
                     continue;
@@ -493,24 +503,22 @@ private static CommandBaseAst LastPipeOnFirstLineWithPipeUsage(PipelineAst pipel
             return lastPipeOnFirstLineWithPipeUsage;
         }
 
-        private static PipelineAst MatchingPipelineAstEnd(List<Ast> pipelineAsts, Token token)
+        private static PipelineAst MatchingPipelineAstEnd(Dictionary<int, List<Ast>> pipelineEndsByLine, Token token)
         {
-            PipelineAst matchingPipeLineAstEnd = null;
-            for (int i = 0; i < pipelineAsts.Count; i++)
+            if (!pipelineEndsByLine.TryGetValue(token.Extent.EndScriptPosition.LineNumber, out var pipelineAsts))
             {
-                if (pipelineAsts[i].Extent.EndScriptPosition.LineNumber > token.Extent.EndScriptPosition.LineNumber)
-                {
-                    break;
-                }
+                return null;
+            }
 
+            for (int i = 0; i < pipelineAsts.Count; i++)
+            {
                 if (PositionIsEqual(pipelineAsts[i].Extent.EndScriptPosition, token.Extent.EndScriptPosition))
                 {
-                    matchingPipeLineAstEnd = pipelineAsts[i] as PipelineAst;
-                    break;
+                    return pipelineAsts[i] as PipelineAst;
                 }
             }
 
-            return matchingPipeLineAstEnd;
+            return null;
         }
 
         /// <summary>

Tests/Engine/TokenOperations.tests.ps1

@@ -2,6 +2,36 @@
 # Licensed under the MIT License.
 
 Describe "TokenOperations" {
+    It 'Should return only script-block expression braces in AST order' {
+        $scriptText = @'
+function Test-Example { 'function body' }
+$value = @{ Handler = { 'hashtable value' } }
+Get-Item | ForEach-Object { Where-Object { $_.Name -ne '{literal}' } }
+& { 'invoked block' }
+'@
+        $tokens = $null
+        $parseErrors = $null
+        $scriptAst = [System.Management.Automation.Language.Parser]::ParseInput($scriptText, [ref] $tokens, [ref] $parseErrors)
+        $parseErrors | Should -BeNullOrEmpty
+        $tokenOperations = [Microsoft.Windows.PowerShell.ScriptAnalyzer.TokenOperations]::new($tokens, $scriptAst)
+        $openBraces = @($tokenOperations.GetOpenBracesInCommandElements())
+        $closeBraces = @($tokenOperations.GetCloseBracesInCommandElements())
+        $openBraces.Count | Should -Be 4
+        $closeBraces.Count | Should -Be 4
+        ($openBraces.Extent.StartOffset -join ',') | Should -Be (@(
+            $scriptText.IndexOf("{ 'hashtable value'")
+            $scriptText.IndexOf('{ Where-Object')
+            $scriptText.IndexOf('{ $_.Name')
+            $scriptText.IndexOf("{ 'invoked block'")
+        ) -join ',')
+        ($closeBraces.Extent.EndOffset -join ',') | Should -Be (@(
+            $scriptText.IndexOf("'hashtable value' }") + "'hashtable value' }".Length
+            $scriptText.IndexOf("'{literal}' } }") + "'{literal}' } }".Length
+            $scriptText.IndexOf("'{literal}' }") + "'{literal}' }".Length
+            $scriptText.IndexOf("'invoked block' }") + "'invoked block' }".Length
+        ) -join ',')
+    }
+
     It "Should return correct AST position for assignment operator in hash table" {
         $scriptText = @'
 $h = @{
@@ -195,4 +225,4 @@ $object.{MemberName
 
     }
 
-}
+}