Begin implementation of visitor to find token in AST

Author: daviesj

Engine/TokenOperations.cs

@@ -232,5 +232,85 @@ private bool OnSameLine(Token token1, Token token2)
         {
             return token1.Extent.StartLineNumber == token2.Extent.EndLineNumber;
         }
+
+        /// <summary>
+        /// Finds the position of a given token in the AST.
+        /// </summary>
+        /// <param name="token">The <see cref="Token"/> to search for.</param>
+        /// <returns>The Ast node directly containing the provided <see cref="Token"/>.</returns>
+        public Ast GetAstPosition(Token token)
+        {
+            FindAstPostitionVisitor findAstVisitor = new FindAstPostitionVisitor(token.Extent.StartScriptPosition);
+            ast.Visit(findAstVisitor);
+            return findAstVisitor.AstPosition;
+        }
+
+    }
+
+    /// <summary>
+    /// Provides an efficient way to find the position in the AST corresponding to a given script position.
+    /// </summary>
+    public class FindAstPostitionVisitor : AstVisitor2
+    {
+        private IScriptPosition searchPosition;
+
+        /// <summary>
+        /// Contains the position in the AST corresponding to the provided <see cref="IScriptPosition"/> upon completion of the <see cref="Ast.Visit(AstVisitor)"/> method.
+        /// </summary>
+        public Ast AstPosition { get; private set; }
+
+        /// <summary>
+        /// Initializes a new instance of the <see cref="FindAstPostitionVisitor"/> class with the postition to search for.
+        /// </summary>
+        /// <param name="position">The script position to search for.</param>
+        public FindAstPostitionVisitor(IScriptPosition position)
+        {
+            this.searchPosition = position;
+        }
+
+        /// <summary>
+        /// Traverses the AST based on offests to find the leaf node which contains the provided <see cref="IScriptPosition"/>.
+        /// This method implements the entire functionality of this visitor. All <see cref="AstVisitor2"/> methods are overridden to simply invoke this one.
+        /// </summary>
+        /// <param name="ast">Current AST node to process.</param>
+        /// <returns>An <see cref="AstVisitAction"/> indicating whether to visit children of the current node.</returns>
+        private AstVisitAction Visit(Ast ast)
+        {
+            if (ast.Extent.StartOffset > searchPosition.Offset || ast.Extent.EndOffset < searchPosition.Offset)
+            {
+                return AstVisitAction.SkipChildren;
+            }
+            else
+            {
+                AstPosition = ast;
+                return AstVisitAction.Continue;
+            }
+        }
+
+        public override AstVisitAction VisitScriptBlock(ScriptBlockAst scriptBlockAst)
+        {
+            return Visit(scriptBlockAst);
+        }
+
+        public override AstVisitAction VisitNamedBlock(NamedBlockAst namedBlockAst)
+        {
+            return Visit(namedBlockAst);
+        }
+
+        public override AstVisitAction VisitAssignmentStatement(AssignmentStatementAst assignmentStatementAst)
+        {
+            return Visit(assignmentStatementAst);
+        }
+
+        public override AstVisitAction VisitCommandExpression(CommandExpressionAst commandExpressionAst)
+        {
+            return Visit(commandExpressionAst);
+        }
+
+        public override AstVisitAction VisitHashtable(HashtableAst hashtableAst)
+        {
+            return Visit(hashtableAst);
+        }
+
     }
 }

Rules/UseConsistentWhitespace.cs

@@ -536,23 +536,10 @@ private IEnumerable<DiagnosticRecord> FindOperatorViolations(TokenOperations tok
                 // exclude assignment operator inside of multi-line hash tables if requested
                 if (IgnoreAssignmentOperatorInsideHashTable && tokenNode.Value.Kind == TokenKind.Equals)
                 {
-                    var enclosingHashTables = tokenOperations.Ast.FindAll(a => a.Extent.StartOffset <= tokenNode.Value.Extent.StartOffset && a.Extent.EndOffset >= tokenNode.Value.Extent.EndOffset && a is HashtableAst, true);
-                    if (enclosingHashTables.Count() > 0)
+                    Ast containingAst = tokenOperations.GetAstPosition(tokenNode.Value);
+                    if (containingAst is HashtableAst && containingAst.Extent.EndLineNumber != containingAst.Extent.StartLineNumber)
                     {
-                        Ast innermostEnclosingHashTable = enclosingHashTables.First();
-                        int smallestSizeSoFar = int.MaxValue;
-                        foreach (var hashTable in enclosingHashTables){
-                            int currentHashTableSize = hashTable.Extent.EndOffset - hashTable.Extent.StartOffset;
-                            if (currentHashTableSize < smallestSizeSoFar)
-                            {
-                                innermostEnclosingHashTable = hashTable;
-                                smallestSizeSoFar = currentHashTableSize;
-                            }
-                        }
-                        if (innermostEnclosingHashTable.Extent.StartLineNumber != innermostEnclosingHashTable.Extent.EndLineNumber)
-                        {
-                            continue;
-                        }
+                        continue;
                     }
                 }
 

Tests/Engine/TokenOperations.tests.ps1

@@ -0,0 +1,21 @@
+# Copyright (c) Microsoft Corporation. All rights reserved.
+# Licensed under the MIT License.
+
+Describe "TokenOperations" {
+    It "Should return correct AST position for assignment operator in hash table" {
+        $scriptText = @'
+$h = @{
+    a = 72
+    b = @{ z = "hi" }
+}
+'@
+        $tokens = $null
+        $parseErrors = $null
+        $scriptAst = [System.Management.Automation.Language.Parser]::ParseInput($scriptText, [ref] $tokens, [ref] $parseErrors)
+        $tokenOperations = New-Object Microsoft.Windows.PowerShell.ScriptAnalyzer.TokenOperations -ArgumentList @($tokens, $scriptAst)
+        $operatorToken = $tokens | Where-Object { $_.Extent.StartOffset -eq 32 }
+        $hashTableAst = $tokenOperations.GetAstPosition($operatorToken)
+        $hashTableAst | Should -BeOfType [System.Management.Automation.Language.HashTableAst]
+        $hashTableAst.Extent.Text | Should -Be '@{ z = "hi" }'
+    }
+}