Merge pull request #1656 from PowerShell/andschwa/resolve-aliases

Resurrect support to resolve aliased references

Author: andyleejordan

src/PowerShellEditorServices/Services/CodeLens/ICodeLensProvider.cs

@@ -1,6 +1,7 @@
 // Copyright (c) Microsoft Corporation.
 // Licensed under the MIT License.
 
+using System.Threading.Tasks;
 using Microsoft.PowerShell.EditorServices.Services.TextDocument;
 using OmniSharp.Extensions.LanguageServer.Protocol.Models;
 
@@ -34,13 +35,12 @@ internal interface ICodeLensProvider
         /// The CodeLens to resolve.
         /// </param>
         /// <param name="scriptFile">
-        /// A CancellationToken which can be used to cancel the
-        /// request.
+        /// The ScriptFile to resolve it in (sometimes unused).
         /// </param>
         /// <returns>
         /// A Task which returns the resolved CodeLens when completed.
         /// </returns>
-        CodeLens ResolveCodeLens(
+        Task<CodeLens> ResolveCodeLens(
             CodeLens codeLens,
             ScriptFile scriptFile);
     }

src/PowerShellEditorServices/Services/CodeLens/PesterCodeLensProvider.cs

@@ -2,6 +2,7 @@
 // Licensed under the MIT License.
 using System;
 using System.Collections.Generic;
+using System.Threading.Tasks;
 using Microsoft.PowerShell.EditorServices.Services;
 using Microsoft.PowerShell.EditorServices.Services.Symbols;
 using Microsoft.PowerShell.EditorServices.Services.TextDocument;
@@ -133,11 +134,11 @@ public CodeLens[] ProvideCodeLenses(ScriptFile scriptFile)
         /// <param name="codeLens">The code lens to resolve.</param>
         /// <param name="scriptFile">The script file.</param>
         /// <returns>The given CodeLens, wrapped in a task.</returns>
-        public CodeLens ResolveCodeLens(CodeLens codeLens, ScriptFile scriptFile)
+        public Task<CodeLens> ResolveCodeLens(CodeLens codeLens, ScriptFile scriptFile)
         {
             // This provider has no specific behavior for
             // resolving CodeLenses.
-            return codeLens;
+            return Task.FromResult(codeLens);
         }
     }
 }

src/PowerShellEditorServices/Services/CodeLens/ReferencesCodeLensProvider.cs

@@ -4,6 +4,7 @@
 using System;
 using System.Collections.Generic;
 using System.Text;
+using System.Threading.Tasks;
 using Microsoft.PowerShell.EditorServices.Services;
 using Microsoft.PowerShell.EditorServices.Services.Symbols;
 using Microsoft.PowerShell.EditorServices.Services.TextDocument;
@@ -79,10 +80,10 @@ public CodeLens[] ProvideCodeLenses(ScriptFile scriptFile)
         /// Take a codelens and create a new codelens object with updated references.
         /// </summary>
         /// <param name="codeLens">The old code lens to get updated references for.</param>
+        /// <param name="scriptFile"></param>
         /// <returns>A new code lens object describing the same data as the old one but with updated references.</returns>
-        public CodeLens ResolveCodeLens(CodeLens codeLens, ScriptFile scriptFile)
+        public async Task<CodeLens> ResolveCodeLens(CodeLens codeLens, ScriptFile scriptFile)
         {
-
             ScriptFile[] references = _workspaceService.ExpandScriptReferences(
                 scriptFile);
 
@@ -91,10 +92,10 @@ public CodeLens ResolveCodeLens(CodeLens codeLens, ScriptFile scriptFile)
                 codeLens.Range.Start.Line + 1,
                 codeLens.Range.Start.Character + 1);
 
-            List<SymbolReference> referencesResult = _symbolsService.FindReferencesOfSymbol(
+            List<SymbolReference> referencesResult = await _symbolsService.FindReferencesOfSymbol(
                 foundSymbol,
                 references,
-                _workspaceService);
+                _workspaceService).ConfigureAwait(false);
 
             Location[] referenceLocations;
             if (referencesResult == null)

src/PowerShellEditorServices/Services/PowerShell/Utility/CommandHelpers.cs

@@ -17,8 +17,8 @@ namespace Microsoft.PowerShell.EditorServices.Services.PowerShell.Utility
     /// </summary>
     internal static class CommandHelpers
     {
-        private static readonly HashSet<string> s_nounExclusionList = new HashSet<string>
-            {
+        private static readonly HashSet<string> s_nounExclusionList = new()
+        {
                 // PowerShellGet v2 nouns
                 "CredsFromCredentialProvider",
                 "DscResource",
@@ -36,8 +36,8 @@ internal static class CommandHelpers
             };
 
         // This is used when a noun exists in multiple modules (for example, "Command" is used in Microsoft.PowerShell.Core and also PowerShellGet)
-        private static readonly HashSet<string> s_cmdletExclusionList = new HashSet<string>
-            {
+        private static readonly HashSet<string> s_cmdletExclusionList = new()
+        {
                 // Commands in PowerShellGet with conflicting nouns
                 "Find-Command",
                 "Find-Module",
@@ -49,17 +49,17 @@ internal static class CommandHelpers
                 "Update-ModuleManifest",
             };
 
-        private static readonly ConcurrentDictionary<string, CommandInfo> s_commandInfoCache =
-            new ConcurrentDictionary<string, CommandInfo>();
-
-        private static readonly ConcurrentDictionary<string, string> s_synopsisCache =
-            new ConcurrentDictionary<string, string>();
+        private static readonly ConcurrentDictionary<string, CommandInfo> s_commandInfoCache = new();
+        private static readonly ConcurrentDictionary<string, string> s_synopsisCache = new();
+        private static readonly ConcurrentDictionary<string, List<string>> s_cmdletToAliasCache = new(System.StringComparer.OrdinalIgnoreCase);
+        private static readonly ConcurrentDictionary<string, string> s_aliasToCmdletCache = new(System.StringComparer.OrdinalIgnoreCase);
 
         /// <summary>
         /// Gets the CommandInfo instance for a command with a particular name.
         /// </summary>
         /// <param name="commandName">The name of the command.</param>
-        /// <param name="powerShellContext">The PowerShellContext to use for running Get-Command.</param>
+        /// <param name="currentRunspace">The current runspace.</param>
+        /// <param name="executionService">The execution service.</param>
         /// <returns>A CommandInfo object with details about the specified command.</returns>
         public static async Task<CommandInfo> GetCommandInfoAsync(
             string commandName,
@@ -97,7 +97,11 @@ public static async Task<CommandInfo> GetCommandInfoAsync(
                 .AddArgument(commandName)
                 .AddParameter("ErrorAction", "Ignore");
 
-            CommandInfo commandInfo = (await executionService.ExecutePSCommandAsync<CommandInfo>(command, CancellationToken.None).ConfigureAwait(false)).FirstOrDefault();
+            IReadOnlyList<CommandInfo> results = await executionService
+                .ExecutePSCommandAsync<CommandInfo>(command, CancellationToken.None)
+                .ConfigureAwait(false);
+
+            CommandInfo commandInfo = results.Count > 0 ? results[0] : null;
 
             // Only cache CmdletInfos since they're exposed in binaries they are likely to not change throughout the session.
             if (commandInfo?.CommandType == CommandTypes.Cmdlet)
@@ -112,8 +116,8 @@ public static async Task<CommandInfo> GetCommandInfoAsync(
         /// Gets the command's "Synopsis" documentation section.
         /// </summary>
         /// <param name="commandInfo">The CommandInfo instance for the command.</param>
-        /// <param name="executionService">The PowerShellContext to use for getting command documentation.</param>
-        /// <returns></returns>
+        /// <param name="executionService">The exeuction service to use for getting command documentation.</param>
+        /// <returns>The synopsis.</returns>
         public static async Task<string> GetCommandSynopsisAsync(
             CommandInfo commandInfo,
             IInternalPowerShellExecutionService executionService)
@@ -146,13 +150,13 @@ public static async Task<string> GetCommandSynopsisAsync(
                 .AddParameter("Online", false)
                 .AddParameter("ErrorAction", "Ignore");
 
-            IReadOnlyList<PSObject> results = await executionService.ExecutePSCommandAsync<PSObject>(command, CancellationToken.None).ConfigureAwait(false);
-            PSObject helpObject = results.FirstOrDefault();
+            IReadOnlyList<PSObject> results = await executionService
+                .ExecutePSCommandAsync<PSObject>(command, CancellationToken.None)
+                .ConfigureAwait(false);
 
             // Extract the synopsis string from the object
-            string synopsisString =
-                (string)helpObject?.Properties["synopsis"].Value ??
-                string.Empty;
+            PSObject helpObject = results.Count > 0 ? results[0] : null;
+            string synopsisString = (string)helpObject?.Properties["synopsis"].Value ?? string.Empty;
 
             // Only cache cmdlet infos because since they're exposed in binaries, the can never change throughout the session.
             if (commandInfo.CommandType == CommandTypes.Cmdlet)
@@ -168,5 +172,39 @@ public static async Task<string> GetCommandSynopsisAsync(
 
             return synopsisString;
         }
+
+        /// <summary>
+        /// Gets all aliases found in the runspace
+        /// </summary>
+        /// <param name="executionService"></param>
+        public static async Task<(Dictionary<string, List<string>>, Dictionary<string, string>)> GetAliasesAsync(IInternalPowerShellExecutionService executionService)
+        {
+            Validate.IsNotNull(nameof(executionService), executionService);
+
+            IEnumerable<CommandInfo> aliases = await executionService.ExecuteDelegateAsync<IEnumerable<CommandInfo>>(
+                nameof(GetAliasesAsync),
+                Execution.ExecutionOptions.Default,
+                (pwsh, _) =>
+                {
+                    CommandInvocationIntrinsics invokeCommand = pwsh.Runspace.SessionStateProxy.InvokeCommand;
+                    return invokeCommand.GetCommands("*", CommandTypes.Alias, nameIsPattern: true);
+                },
+                CancellationToken.None).ConfigureAwait(false);
+
+            foreach (AliasInfo aliasInfo in aliases)
+            {
+                // TODO: When we move to netstandard2.1, we can use another overload which generates
+                // static delegates and thus reduces allocations.
+                s_cmdletToAliasCache.AddOrUpdate(
+                    aliasInfo.Definition,
+                    (_) => new List<string> { aliasInfo.Name },
+                    (_, v) => { v.Add(aliasInfo.Name); return v; });
+
+                s_aliasToCmdletCache.TryAdd(aliasInfo.Name, aliasInfo.Definition);
+            }
+
+            return (new Dictionary<string, List<string>>(s_cmdletToAliasCache),
+                new Dictionary<string, string>(s_aliasToCmdletCache));
+        }
     }
 }

src/PowerShellEditorServices/Services/Symbols/SymbolsService.cs

@@ -39,7 +39,6 @@ internal class SymbolsService
 
         private readonly ConcurrentDictionary<string, ICodeLensProvider> _codeLensProviders;
         private readonly ConcurrentDictionary<string, IDocumentSymbolProvider> _documentSymbolProviders;
-
         #endregion
 
         #region Constructors
@@ -49,6 +48,10 @@ internal class SymbolsService
         /// the given Runspace to execute language service operations.
         /// </summary>
         /// <param name="factory">An ILoggerFactory implementation used for writing log messages.</param>
+        /// <param name="runspaceContext"></param>
+        /// <param name="executionService"></param>
+        /// <param name="workspaceService"></param>
+        /// <param name="configurationService"></param>
         public SymbolsService(
             ILoggerFactory factory,
             IRunspaceContext runspaceContext,
@@ -176,7 +179,7 @@ public static SymbolReference FindSymbolAtLocation(
         /// <param name="referencedFiles">An array of scriptFiles too search for references in</param>
         /// <param name="workspace">The workspace that will be searched for symbols</param>
         /// <returns>FindReferencesResult</returns>
-        public List<SymbolReference> FindReferencesOfSymbol(
+        public async Task<List<SymbolReference>> FindReferencesOfSymbol(
             SymbolReference foundSymbol,
             ScriptFile[] referencedFiles,
             WorkspaceService workspace)
@@ -186,7 +189,7 @@ public List<SymbolReference> FindReferencesOfSymbol(
                 return null;
             }
 
-            // NOTE: we use to make sure aliases were loaded but took it out because we needed the pipeline thread.
+            (Dictionary<string, List<string>> cmdletToAliases, Dictionary<string, string> aliasToCmdlets) = await CommandHelpers.GetAliasesAsync(_executionService).ConfigureAwait(false);
 
             // We want to look for references first in referenced files, hence we use ordered dictionary
             // TODO: File system case-sensitivity is based on filesystem not OS, but OS is a much cheaper heuristic
@@ -221,7 +224,8 @@ public List<SymbolReference> FindReferencesOfSymbol(
                 IEnumerable<SymbolReference> references = AstOperations.FindReferencesOfSymbol(
                     file.ScriptAst,
                     foundSymbol,
-                    needsAliases: false);
+                    cmdletToAliases,
+                    aliasToCmdlets);
 
                 foreach (SymbolReference reference in references)
                 {
@@ -264,10 +268,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>
@@ -306,7 +307,7 @@ public static SymbolReference FindFunctionDefinitionAtLocation(
         /// <param name="lineNumber">The line number at which the symbol can be located.</param>
         /// <param name="columnNumber">The column number at which the symbol can be located.</param>
         /// <returns></returns>
-        public async Task<SymbolDetails> FindSymbolDetailsAtLocationAsync(
+        public Task<SymbolDetails> FindSymbolDetailsAtLocationAsync(
             ScriptFile scriptFile,
             int lineNumber,
             int columnNumber)
@@ -319,16 +320,14 @@ public async Task<SymbolDetails> FindSymbolDetailsAtLocationAsync(
 
             if (symbolReference == null)
             {
-                return null;
+                return Task.FromResult<SymbolDetails>(null);
             }
 
             symbolReference.FilePath = scriptFile.FilePath;
-            SymbolDetails symbolDetails = await SymbolDetails.CreateAsync(
+            return SymbolDetails.CreateAsync(
                 symbolReference,
                 _runspaceContext.CurrentRunspace,
-                _executionService).ConfigureAwait(false);
-
-            return symbolDetails;
+                _executionService);
         }
 
         /// <summary>
@@ -446,8 +445,7 @@ public async Task<SymbolReference> GetDefinitionOfSymbolAsync(
             if (foundDefinition == null)
             {
                 // Get a list of all powershell files in the workspace path
-                IEnumerable<string> allFiles = _workspaceService.EnumeratePSFiles();
-                foreach (string file in allFiles)
+                foreach (string file in _workspaceService.EnumeratePSFiles())
                 {
                     if (filesSearched.Contains(file))
                     {
@@ -543,7 +541,7 @@ private ScriptFile[] GetBuiltinCommandScriptFiles(
             }
 
             string modPath = moduleInfo.Path;
-            List<ScriptFile> scriptFiles = new List<ScriptFile>();
+            List<ScriptFile> scriptFiles = new();
             ScriptFile newFile;
 
             // find any files where the moduleInfo's path ends with ps1 or psm1
@@ -598,7 +596,7 @@ public static FunctionDefinitionAst GetFunctionDefinitionForHelpComment(
             IEnumerable<Ast> foundAsts = scriptFile.ScriptAst.FindAll(
                 ast =>
                 {
-                    if (!(ast is FunctionDefinitionAst fdAst))
+                    if (ast is not FunctionDefinitionAst fdAst)
                     {
                         return false;
                     }
@@ -608,7 +606,7 @@ public static FunctionDefinitionAst GetFunctionDefinitionForHelpComment(
                 },
                 true);
 
-            if (foundAsts == null || !foundAsts.Any())
+            if (foundAsts?.Any() != true)
             {
                 helpLocation = null;
                 return null;

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

@@ -153,42 +153,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)
+            IDictionary<string, List<string>> cmdletToAliasDictionary = default,
+            IDictionary<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;