Overhaul workspace search for symbol references (#1917)

Significantly reduce performance overhead of reference finding in large workspaces.

* Dependent on PowerShell/vscode-powershell#4170
* Adds a reference cache to every ScriptFile
* Workspace scan is performed only once on first request
* An LSP file system watcher provides updates for files not changed via Did*TextDocument notifications
* Adds a setting to only search "open documents" for references. This disables both the initial workspace scan and the file system watcher, relying only on Did*TextDocument notifications.

As a stress test I opened up my profile directory (which has ~3k script files in total in the Modules directory) and created a file with a function definition on every line. I then tabbed to a different file, and then tabbed back to the new file. Before the changes, the references code lens took ~10 seconds to populate and my CPU spiked to ~50% usage. After the changes, they populated instantly and CPU spiked to ~2% usage.

Author: SeeminglyScience

src/PowerShellEditorServices/Server/PsesLanguageServer.cs

@@ -119,6 +119,7 @@ public async Task StartAsync()
                     .WithHandler<ShowHelpHandler>()
                     .WithHandler<ExpandAliasHandler>()
                     .WithHandler<PsesSemanticTokensHandler>()
+                    .WithHandler<DidChangeWatchedFilesHandler>()
                     // NOTE: The OnInitialize delegate gets run when we first receive the
                     // _Initialize_ request:
                     // https://microsoft.github.io/language-server-protocol/specifications/specification-current/#initialize

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

@@ -1,6 +1,7 @@
 // Copyright (c) Microsoft Corporation.
 // Licensed under the MIT License.
 
+using System;
 using System.Collections.Concurrent;
 using System.Collections.Generic;
 using System.Management.Automation;
@@ -57,6 +58,39 @@ public record struct AliasMap(
         internal static readonly ConcurrentDictionary<string, List<string>> s_cmdletToAliasCache = new(System.StringComparer.OrdinalIgnoreCase);
         internal static readonly ConcurrentDictionary<string, string> s_aliasToCmdletCache = new(System.StringComparer.OrdinalIgnoreCase);
 
+        /// <summary>
+        /// Gets the actual command behind a fully module qualified command invocation, e.g.
+        /// <c>Microsoft.PowerShell.Management\Get-ChildItem</c> will return <c>Get-ChildItem</c>
+        /// </summary>
+        /// <param name="invocationName">
+        /// The potentially module qualified command name at the site of invocation.
+        /// </param>
+        /// <param name="moduleName">
+        /// A reference that will contain the module name if the invocation is module qualified.
+        /// </param>
+        /// <returns>The actual command name.</returns>
+        public static string StripModuleQualification(string invocationName, out ReadOnlyMemory<char> moduleName)
+        {
+            int slashIndex = invocationName.LastIndexOfAny(new[] { '\\', '/' });
+            if (slashIndex is -1)
+            {
+                moduleName = default;
+                return invocationName;
+            }
+
+            // If '\' is the last character then it's probably not a module qualified command.
+            if (slashIndex == invocationName.Length - 1)
+            {
+                moduleName = default;
+                return invocationName;
+            }
+
+            // Storing moduleName as ROMemory saves a string allocation in the common case where it
+            // is not needed.
+            moduleName = invocationName.AsMemory().Slice(0, slashIndex);
+            return invocationName.Substring(slashIndex + 1);
+        }
+
         /// <summary>
         /// Gets the CommandInfo instance for a command with a particular name.
         /// </summary>

src/PowerShellEditorServices/Services/Symbols/ReferenceTable.cs

@@ -0,0 +1,118 @@
+// Copyright (c) Microsoft Corporation.
+// Licensed under the MIT License.
+
+#nullable enable
+
+using System;
+using System.Collections.Concurrent;
+using System.Management.Automation.Language;
+using Microsoft.PowerShell.EditorServices.Services.TextDocument;
+using Microsoft.PowerShell.EditorServices.Services.PowerShell.Utility;
+using Microsoft.PowerShell.EditorServices.Services.Symbols;
+
+namespace Microsoft.PowerShell.EditorServices.Services;
+
+/// <summary>
+/// Represents the symbols that are referenced and their locations within a single document.
+/// </summary>
+internal sealed class ReferenceTable
+{
+    private readonly ScriptFile _parent;
+
+    private readonly ConcurrentDictionary<string, ConcurrentBag<IScriptExtent>> _symbolReferences = new(StringComparer.OrdinalIgnoreCase);
+
+    private bool _isInited;
+
+    public ReferenceTable(ScriptFile parent) => _parent = parent;
+
+    /// <summary>
+    /// Clears the reference table causing it to rescan the source AST when queried.
+    /// </summary>
+    public void TagAsChanged()
+    {
+        _symbolReferences.Clear();
+        _isInited = false;
+    }
+
+    // Prefer checking if the dictionary has contents to determine if initialized. The field
+    // `_isInited` is to guard against rescanning files with no command references, but will
+    // generally be less reliable of a check.
+    private bool IsInitialized => !_symbolReferences.IsEmpty || _isInited;
+
+    internal bool TryGetReferences(string command, out ConcurrentBag<IScriptExtent>? references)
+    {
+        EnsureInitialized();
+        return _symbolReferences.TryGetValue(command, out references);
+    }
+
+    internal void EnsureInitialized()
+    {
+        if (IsInitialized)
+        {
+            return;
+        }
+
+        _parent.ScriptAst.Visit(new ReferenceVisitor(this));
+    }
+
+    private void AddReference(string symbol, IScriptExtent extent)
+    {
+        _symbolReferences.AddOrUpdate(
+            symbol,
+            _ => new ConcurrentBag<IScriptExtent> { extent },
+            (_, existing) =>
+            {
+                existing.Add(extent);
+                return existing;
+            });
+    }
+
+    private sealed class ReferenceVisitor : AstVisitor
+    {
+        private readonly ReferenceTable _references;
+
+        public ReferenceVisitor(ReferenceTable references) => _references = references;
+
+        public override AstVisitAction VisitCommand(CommandAst commandAst)
+        {
+            string? commandName = GetCommandName(commandAst);
+            if (string.IsNullOrEmpty(commandName))
+            {
+                return AstVisitAction.Continue;
+            }
+
+            _references.AddReference(
+                CommandHelpers.StripModuleQualification(commandName, out _),
+                commandAst.CommandElements[0].Extent);
+
+            return AstVisitAction.Continue;
+
+            static string? GetCommandName(CommandAst commandAst)
+            {
+                string commandName = commandAst.GetCommandName();
+                if (!string.IsNullOrEmpty(commandName))
+                {
+                    return commandName;
+                }
+
+                if (commandAst.CommandElements[0] is not ExpandableStringExpressionAst expandableStringExpressionAst)
+                {
+                    return null;
+                }
+
+                return AstOperations.TryGetInferredValue(expandableStringExpressionAst, out string value) ? value : null;
+            }
+        }
+
+        public override AstVisitAction VisitVariableExpression(VariableExpressionAst variableExpressionAst)
+        {
+            // TODO: Consider tracking unscoped variable references only when they declared within
+            // the same function definition.
+            _references.AddReference(
+                $"${variableExpressionAst.VariablePath.UserPath}",
+                variableExpressionAst.Extent);
+
+            return AstVisitAction.Continue;
+        }
+    }
+}