Clean up and fix FindReferencesVisitor

This consolidated the constructors so we don't have some unused (and
confusing) boolean differentiating them, and instead just provide one
constructor with default values. Also clean up `AstOperationsTests`.

Author: andyleejordan

src/PowerShellEditorServices/Services/Symbols/SymbolsService.cs

@@ -264,10 +264,7 @@ public static IReadOnlyList<SymbolReference> FindOccurrencesInFile(
                 return null;
             }
 
-            return AstOperations.FindReferencesOfSymbol(
-                file.ScriptAst,
-                foundSymbol,
-                needsAliases: false).ToArray();
+            return AstOperations.FindReferencesOfSymbol(file.ScriptAst, foundSymbol).ToArray();
         }
 
         /// <summary>

src/PowerShellEditorServices/Services/Symbols/Vistors/AstOperations.cs

@@ -2,6 +2,7 @@
 // Licensed under the MIT License.
 
 using System;
+using System.Collections.Concurrent;
 using System.Collections.Generic;
 using System.Diagnostics;
 using System.Linq;
@@ -153,42 +154,21 @@ public static SymbolReference FindCommandAtPosition(Ast scriptAst, int lineNumbe
         /// </summary>
         /// <param name="scriptAst">The abstract syntax tree of the given script</param>
         /// <param name="symbolReference">The symbol that we are looking for referneces of</param>
-        /// <param name="CmdletToAliasDictionary">Dictionary maping cmdlets to aliases for finding alias references</param>
-        /// <param name="AliasToCmdletDictionary">Dictionary maping aliases to cmdlets for finding alias references</param>
+        /// <param name="cmdletToAliasDictionary">Dictionary maping cmdlets to aliases for finding alias references</param>
+        /// <param name="aliasToCmdletDictionary">Dictionary maping aliases to cmdlets for finding alias references</param>
         /// <returns></returns>
         public static IEnumerable<SymbolReference> FindReferencesOfSymbol(
             Ast scriptAst,
             SymbolReference symbolReference,
-            Dictionary<String, List<String>> CmdletToAliasDictionary,
-            Dictionary<String, String> AliasToCmdletDictionary)
+            ConcurrentDictionary<string, List<string>> cmdletToAliasDictionary = default,
+            ConcurrentDictionary<string, string> aliasToCmdletDictionary = default)
         {
             // find the symbol evaluators for the node types we are handling
-            FindReferencesVisitor referencesVisitor =
-                new FindReferencesVisitor(
-                    symbolReference,
-                    CmdletToAliasDictionary,
-                    AliasToCmdletDictionary);
-            scriptAst.Visit(referencesVisitor);
+            FindReferencesVisitor referencesVisitor = new(
+                symbolReference,
+                cmdletToAliasDictionary,
+                aliasToCmdletDictionary);
 
-            return referencesVisitor.FoundReferences;
-        }
-
-        /// <summary>
-        /// Finds all references (not including aliases) in a script for the given symbol
-        /// </summary>
-        /// <param name="scriptAst">The abstract syntax tree of the given script</param>
-        /// <param name="foundSymbol">The symbol that we are looking for referneces of</param>
-        /// <param name="needsAliases">If this reference search needs aliases.
-        /// This should always be false and used for occurence requests</param>
-        /// <returns>A collection of SymbolReference objects that are refrences to the symbolRefrence
-        /// not including aliases</returns>
-        public static IEnumerable<SymbolReference> FindReferencesOfSymbol(
-            ScriptBlockAst scriptAst,
-            SymbolReference foundSymbol,
-            bool needsAliases)
-        {
-            FindReferencesVisitor referencesVisitor =
-                new FindReferencesVisitor(foundSymbol);
             scriptAst.Visit(referencesVisitor);
 
             return referencesVisitor.FoundReferences;

src/PowerShellEditorServices/Services/Symbols/Vistors/FindReferencesVisitor.cs

@@ -2,6 +2,7 @@
 // Licensed under the MIT License.
 
 using System;
+using System.Collections.Concurrent;
 using System.Collections.Generic;
 using System.Management.Automation.Language;
 
@@ -12,48 +13,47 @@ namespace Microsoft.PowerShell.EditorServices.Services.Symbols
     /// </summary>
     internal class FindReferencesVisitor : AstVisitor
     {
-        private SymbolReference symbolRef;
-        private Dictionary<String, List<String>> CmdletToAliasDictionary;
-        private Dictionary<String, String> AliasToCmdletDictionary;
-        private string symbolRefCommandName;
-        private bool needsAliases;
+        private readonly SymbolReference _symbolRef;
+        private readonly ConcurrentDictionary<string, List<string>> _cmdletToAliasDictionary;
+        private readonly ConcurrentDictionary<string, string> _aliasToCmdletDictionary;
+        private readonly string _symbolRefCommandName;
+        private readonly bool _needsAliases;
 
         public List<SymbolReference> FoundReferences { get; set; }
 
         /// <summary>
         /// Constructor used when searching for aliases is needed
         /// </summary>
         /// <param name="symbolReference">The found symbolReference that other symbols are being compared to</param>
-        /// <param name="CmdletToAliasDictionary">Dictionary maping cmdlets to aliases for finding alias references</param>
-        /// <param name="AliasToCmdletDictionary">Dictionary maping aliases to cmdlets for finding alias references</param>
+        /// <param name="cmdletToAliasDictionary">Dictionary maping cmdlets to aliases for finding alias references</param>
+        /// <param name="aliasToCmdletDictionary">Dictionary maping aliases to cmdlets for finding alias references</param>
         public FindReferencesVisitor(
             SymbolReference symbolReference,
-            Dictionary<String, List<String>> CmdletToAliasDictionary,
-            Dictionary<String, String> AliasToCmdletDictionary)
+            ConcurrentDictionary<string, List<string>> cmdletToAliasDictionary = default,
+            ConcurrentDictionary<string, string> aliasToCmdletDictionary = default)
         {
-            this.symbolRef = symbolReference;
-            this.FoundReferences = new List<SymbolReference>();
-            this.needsAliases = true;
-            this.CmdletToAliasDictionary = CmdletToAliasDictionary;
-            this.AliasToCmdletDictionary = AliasToCmdletDictionary;
-
-            // Try to get the symbolReference's command name of an alias,
-            // if a command name does not exists (if the symbol isn't an alias to a command)
-            // set symbolRefCommandName to and empty string value
-            AliasToCmdletDictionary.TryGetValue(symbolReference.ScriptRegion.Text, out symbolRefCommandName);
-            if (symbolRefCommandName == null) { symbolRefCommandName = string.Empty; }
+            _symbolRef = symbolReference;
+            FoundReferences = new List<SymbolReference>();
 
-        }
+            if (cmdletToAliasDictionary is null || aliasToCmdletDictionary is null)
+            {
+                _needsAliases = false;
+                return;
+            }
 
-        /// <summary>
-        /// Constructor used when searching for aliases is not needed
-        /// </summary>
-        /// <param name="foundSymbol">The found symbolReference that other symbols are being compared to</param>
-        public FindReferencesVisitor(SymbolReference foundSymbol)
-        {
-            this.symbolRef = foundSymbol;
-            this.FoundReferences = new List<SymbolReference>();
-            this.needsAliases = false;
+            _needsAliases = true;
+            _cmdletToAliasDictionary = cmdletToAliasDictionary;
+            _aliasToCmdletDictionary = aliasToCmdletDictionary;
+
+            // Try to get the symbolReference's command name of an alias. If a command name does not
+            // exists (if the symbol isn't an alias to a command) set symbolRefCommandName to an
+            // empty string.
+            aliasToCmdletDictionary.TryGetValue(symbolReference.ScriptRegion.Text, out _symbolRefCommandName);
+
+            if (_symbolRefCommandName == null)
+            {
+                _symbolRefCommandName = string.Empty;
+            }
         }
 
         /// <summary>
@@ -68,50 +68,44 @@ public override AstVisitAction VisitCommand(CommandAst commandAst)
             Ast commandNameAst = commandAst.CommandElements[0];
             string commandName = commandNameAst.Extent.Text;
 
-            if(symbolRef.SymbolType.Equals(SymbolType.Function))
+            if (_symbolRef.SymbolType.Equals(SymbolType.Function))
             {
-                if (needsAliases)
+                if (_needsAliases)
                 {
-                    // Try to get the commandAst's name and aliases,
-                    // if a command does not exists (if the symbol isn't an alias to a command)
-                    // set command to and empty string value string command
-                    // if the aliases do not exist (if the symvol isn't a command that has aliases)
+                    // Try to get the commandAst's name and aliases.
+                    //
+                    // If a command does not exist (if the symbol isn't an alias to a command) set
+                    // command to an empty string value string command.
+                    //
+                    // If the aliases do not exist (if the symbol isn't a command that has aliases)
                     // set aliases to an empty List<string>
-                    string command;
-                    List<string> alaises;
-                    CmdletToAliasDictionary.TryGetValue(commandName, out alaises);
-                    AliasToCmdletDictionary.TryGetValue(commandName, out command);
-                    if (alaises == null) { alaises = new List<string>(); }
+                    _cmdletToAliasDictionary.TryGetValue(commandName, out List<string> aliases);
+                    _aliasToCmdletDictionary.TryGetValue(commandName, out string command);
+                    if (aliases == null) { aliases = new List<string>(); }
                     if (command == null) { command = string.Empty; }
 
-                    if (symbolRef.SymbolType.Equals(SymbolType.Function))
+                    // Check if the found symbol's name is the same as the commandAst's name OR
+                    // if the symbol's name is an alias for this commandAst's name (commandAst is a cmdlet) OR
+                    // if the symbol's name is the same as the commandAst's cmdlet name (commandAst is a alias)
+                    if (commandName.Equals(_symbolRef.SymbolName, StringComparison.OrdinalIgnoreCase)
+                        // Note that PowerShell command names and aliases are case insensitive.
+                        || aliases.Exists((match) => string.Equals(match, _symbolRef.ScriptRegion.Text, StringComparison.OrdinalIgnoreCase))
+                        || command.Equals(_symbolRef.ScriptRegion.Text, StringComparison.OrdinalIgnoreCase)
+                        || (!string.IsNullOrEmpty(command)
+                            && command.Equals(_symbolRefCommandName, StringComparison.OrdinalIgnoreCase)))
                     {
-                        // Check if the found symbol's name is the same as the commandAst's name OR
-                        // if the symbol's name is an alias for this commandAst's name (commandAst is a cmdlet) OR
-                        // if the symbol's name is the same as the commandAst's cmdlet name (commandAst is a alias)
-                        if (commandName.Equals(symbolRef.SymbolName, StringComparison.CurrentCultureIgnoreCase) ||
-                        alaises.Contains(symbolRef.ScriptRegion.Text.ToLower()) ||
-                        command.Equals(symbolRef.ScriptRegion.Text, StringComparison.CurrentCultureIgnoreCase) ||
-                        (!string.IsNullOrEmpty(command) && command.Equals(symbolRefCommandName, StringComparison.CurrentCultureIgnoreCase)))
-                        {
-                            this.FoundReferences.Add(new SymbolReference(
-                                SymbolType.Function,
-                                commandNameAst.Extent));
-                        }
+                        FoundReferences.Add(new SymbolReference(SymbolType.Function, commandNameAst.Extent));
                     }
-
                 }
                 else // search does not include aliases
                 {
-                    if (commandName.Equals(symbolRef.SymbolName, StringComparison.CurrentCultureIgnoreCase))
+                    if (commandName.Equals(_symbolRef.SymbolName, StringComparison.OrdinalIgnoreCase))
                     {
-                        this.FoundReferences.Add(new SymbolReference(
-                            SymbolType.Function,
-                            commandNameAst.Extent));
+                        FoundReferences.Add(new SymbolReference(SymbolType.Function, commandNameAst.Extent));
                     }
                 }
-
             }
+
             return base.VisitCommand(commandAst);
         }
 
@@ -135,12 +129,10 @@ public override AstVisitAction VisitFunctionDefinition(FunctionDefinitionAst fun
                 File = functionDefinitionAst.Extent.File
             };
 
-            if (symbolRef.SymbolType.Equals(SymbolType.Function) &&
-                nameExtent.Text.Equals(symbolRef.SymbolName, StringComparison.CurrentCultureIgnoreCase))
+            if (_symbolRef.SymbolType.Equals(SymbolType.Function) &&
+                nameExtent.Text.Equals(_symbolRef.SymbolName, StringComparison.CurrentCultureIgnoreCase))
             {
-                this.FoundReferences.Add(new SymbolReference(
-                                          SymbolType.Function,
-                                          nameExtent));
+                FoundReferences.Add(new SymbolReference(SymbolType.Function, nameExtent));
             }
             return base.VisitFunctionDefinition(functionDefinitionAst);
         }
@@ -153,12 +145,10 @@ public override AstVisitAction VisitFunctionDefinition(FunctionDefinitionAst fun
         /// <returns>A visit action that continues the search for references</returns>
         public override AstVisitAction VisitCommandParameter(CommandParameterAst commandParameterAst)
         {
-            if (symbolRef.SymbolType.Equals(SymbolType.Parameter) &&
-                commandParameterAst.Extent.Text.Equals(symbolRef.SymbolName, StringComparison.CurrentCultureIgnoreCase))
+            if (_symbolRef.SymbolType.Equals(SymbolType.Parameter) &&
+                commandParameterAst.Extent.Text.Equals(_symbolRef.SymbolName, StringComparison.CurrentCultureIgnoreCase))
             {
-                this.FoundReferences.Add(new SymbolReference(
-                                         SymbolType.Parameter,
-                                         commandParameterAst.Extent));
+                FoundReferences.Add(new SymbolReference(SymbolType.Parameter, commandParameterAst.Extent));
             }
             return AstVisitAction.Continue;
         }
@@ -171,12 +161,10 @@ public override AstVisitAction VisitCommandParameter(CommandParameterAst command
         /// <returns>A visit action that continues the search for references</returns>
         public override AstVisitAction VisitVariableExpression(VariableExpressionAst variableExpressionAst)
         {
-            if(symbolRef.SymbolType.Equals(SymbolType.Variable) &&
-                variableExpressionAst.Extent.Text.Equals(symbolRef.SymbolName, StringComparison.CurrentCultureIgnoreCase))
+            if (_symbolRef.SymbolType.Equals(SymbolType.Variable)
+                && variableExpressionAst.Extent.Text.Equals(_symbolRef.SymbolName, StringComparison.CurrentCultureIgnoreCase))
             {
-                this.FoundReferences.Add(new SymbolReference(
-                                         SymbolType.Variable,
-                                         variableExpressionAst.Extent));
+                FoundReferences.Add(new SymbolReference(SymbolType.Variable, variableExpressionAst.Extent));
             }
             return AstVisitAction.Continue;
         }
@@ -186,7 +174,7 @@ private static (int, int) GetStartColumnAndLineNumbersFromAst(FunctionDefinition
         {
             int startColumnNumber = ast.Extent.StartColumnNumber;
             int startLineNumber = ast.Extent.StartLineNumber;
-            int astOffset = 0;
+            int astOffset;
 
             if (ast.IsFilter)
             {

test/PowerShellEditorServices.Test/Services/Symbols/AstOperationsTests.cs

@@ -7,13 +7,13 @@
 using Microsoft.PowerShell.EditorServices.Services.Symbols;
 using OmniSharp.Extensions.LanguageServer.Protocol.Models;
 using Xunit;
-using Xunit.Abstractions;
 
 namespace Microsoft.PowerShell.EditorServices.Test.Services.Symbols
 {
+    [Trait("Category", "AstOperations")]
     public class AstOperationsTests
     {
-        private static string s_scriptString = @"function BasicFunction {}
+        private const string s_scriptString = @"function BasicFunction {}
 BasicFunction
 
 function          FunctionWithExtraSpace
@@ -36,9 +36,8 @@ function          FunctionWithExtraSpace
 
     FunctionNameOnDifferentLine
 ";
-        private static ScriptBlockAst s_ast = (ScriptBlockAst) ScriptBlock.Create(s_scriptString).Ast;
+        private static readonly ScriptBlockAst s_ast = (ScriptBlockAst) ScriptBlock.Create(s_scriptString).Ast;
 
-        [Trait("Category", "AstOperations")]
         [Theory]
         [InlineData(2, 3, "BasicFunction")]
         [InlineData(7, 18, "FunctionWithExtraSpace")]
@@ -50,14 +49,13 @@ public void CanFindSymbolAtPostion(int lineNumber, int columnNumber, string expe
             Assert.Equal(expectedName, reference.SymbolName);
         }
 
-        [Trait("Category", "AstOperations")]
         [Theory]
-        [MemberData(nameof(FindReferencesOfSymbolAtPostionData), parameters: 3)]
+        [MemberData(nameof(FindReferencesOfSymbolAtPostionData))]
         public void CanFindReferencesOfSymbolAtPostion(int lineNumber, int columnNumber, Position[] positions)
         {
             SymbolReference symbol = AstOperations.FindSymbolAtPosition(s_ast, lineNumber, columnNumber);
 
-            IEnumerable<SymbolReference> references = AstOperations.FindReferencesOfSymbol(s_ast, symbol, needsAliases: false);
+            IEnumerable<SymbolReference> references = AstOperations.FindReferencesOfSymbol(s_ast, symbol);
 
             int positionsIndex = 0;
             foreach (SymbolReference reference in references)
@@ -69,9 +67,7 @@ public void CanFindReferencesOfSymbolAtPostion(int lineNumber, int columnNumber,
             }
         }
 
-        public static object[][] FindReferencesOfSymbolAtPostionData => s_findReferencesOfSymbolAtPostionData;
-
-        private static readonly object[][] s_findReferencesOfSymbolAtPostionData = new object[][]
+        public static object[][] FindReferencesOfSymbolAtPostionData { get; } = new object[][]
         {
             new object[] { 2, 3, new[] { new Position(1, 10), new Position(2, 1) } },
             new object[] { 7, 18, new[] { new Position(4, 19), new Position(7, 3) } },