Migrate FindReferencesVisitor to ReferenceTable.ReferenceVisitor

This required using a `ScriptFile` for the associated test in
`AstOperationsTests`. The coupling of the visitor to the file seems fine
since that's the _point_ of this improvement: a cached sets of
references for each file. We also have to carefully sort our list of
symbol references for things that expect them in order.

Author: andyleejordan

src/PowerShellEditorServices/Services/Symbols/IDocumentSymbolProvider.cs

@@ -20,6 +20,6 @@ internal interface IDocumentSymbolProvider
         /// The document for which SymbolReferences should be provided.
         /// </param>
         /// <returns>An IEnumerable collection of SymbolReferences.</returns>
-        IEnumerable<ISymbolReference> ProvideDocumentSymbols(ScriptFile scriptFile);
+        IEnumerable<SymbolReference> ProvideDocumentSymbols(ScriptFile scriptFile);
     }
 }

src/PowerShellEditorServices/Services/Symbols/PesterDocumentSymbolProvider.cs

@@ -17,7 +17,7 @@ internal class PesterDocumentSymbolProvider : IDocumentSymbolProvider
     {
         string IDocumentSymbolProvider.ProviderId => nameof(PesterDocumentSymbolProvider);
 
-        IEnumerable<ISymbolReference> IDocumentSymbolProvider.ProvideDocumentSymbols(
+        IEnumerable<SymbolReference> IDocumentSymbolProvider.ProvideDocumentSymbols(
             ScriptFile scriptFile)
         {
             if (!scriptFile.FilePath.EndsWith(

src/PowerShellEditorServices/Services/Symbols/PsdDocumentSymbolProvider.cs

@@ -17,7 +17,7 @@ internal class PsdDocumentSymbolProvider : IDocumentSymbolProvider
     {
         string IDocumentSymbolProvider.ProviderId => nameof(PsdDocumentSymbolProvider);
 
-        IEnumerable<ISymbolReference> IDocumentSymbolProvider.ProvideDocumentSymbols(
+        IEnumerable<SymbolReference> IDocumentSymbolProvider.ProvideDocumentSymbols(
             ScriptFile scriptFile)
         {
             if ((scriptFile.FilePath?.EndsWith(".psd1", StringComparison.OrdinalIgnoreCase) == true) ||

src/PowerShellEditorServices/Services/Symbols/ReferenceTable.cs

@@ -9,6 +9,8 @@
 using Microsoft.PowerShell.EditorServices.Services.TextDocument;
 using Microsoft.PowerShell.EditorServices.Services.PowerShell.Utility;
 using Microsoft.PowerShell.EditorServices.Services.Symbols;
+using System.Collections.Generic;
+using Microsoft.PowerShell.EditorServices.Utility;
 
 namespace Microsoft.PowerShell.EditorServices.Services;
 
@@ -19,7 +21,7 @@ internal sealed class ReferenceTable
 {
     private readonly ScriptFile _parent;
 
-    private readonly ConcurrentDictionary<string, ConcurrentBag<IScriptExtent>> _symbolReferences = new(StringComparer.OrdinalIgnoreCase);
+    private readonly ConcurrentDictionary<string, ConcurrentBag<SymbolReference>> _symbolReferences = new(StringComparer.OrdinalIgnoreCase);
 
     private bool _isInited;
 
@@ -41,12 +43,24 @@ public void TagAsChanged()
     /// </summary>
     private bool IsInitialized => !_symbolReferences.IsEmpty || _isInited;
 
-    internal bool TryGetReferences(string command, out ConcurrentBag<IScriptExtent>? references)
+    internal bool TryGetReferences(string command, out ConcurrentBag<SymbolReference>? references)
     {
         EnsureInitialized();
         return _symbolReferences.TryGetValue(command, out references);
     }
 
+    // TODO: Should this be improved, or pre-sorted?
+    internal IReadOnlyList<SymbolReference> GetAllReferences()
+    {
+        EnsureInitialized();
+        List<SymbolReference> allReferences = new();
+        foreach (ConcurrentBag<SymbolReference> bag in _symbolReferences.Values)
+        {
+            allReferences.AddRange(bag);
+        }
+        return allReferences;
+    }
+
     internal void EnsureInitialized()
     {
         if (IsInitialized)
@@ -57,51 +71,73 @@ internal void EnsureInitialized()
         _parent.ScriptAst.Visit(new ReferenceVisitor(this));
     }
 
-    private void AddReference(string symbol, IScriptExtent extent)
+    private void AddReference(SymbolType type, string name, IScriptExtent extent)
     {
+        SymbolReference symbol = new(type, name, extent, _parent);
         _symbolReferences.AddOrUpdate(
-            symbol,
-            _ => new ConcurrentBag<IScriptExtent> { extent },
+            name,
+            _ => new ConcurrentBag<SymbolReference> { symbol },
             (_, existing) =>
             {
-                existing.Add(extent);
+                existing.Add(symbol);
                 return existing;
             });
     }
 
-    private sealed class ReferenceVisitor : AstVisitor
+    // TODO: Should we move this to AstOperations.cs? It is highly coupled to `ReferenceTable`,
+    // perhaps it doesn't have to be.
+    private sealed class ReferenceVisitor : AstVisitor2
     {
         private readonly ReferenceTable _references;
 
         public ReferenceVisitor(ReferenceTable references) => _references = references;
 
-        private static string? GetCommandName(CommandAst commandAst)
+        public override AstVisitAction VisitCommand(CommandAst commandAst)
         {
-            string commandName = commandAst.GetCommandName();
-            if (!string.IsNullOrEmpty(commandName))
+            string? commandName = VisitorUtils.GetCommandName(commandAst);
+            if (string.IsNullOrEmpty(commandName))
             {
-                return commandName;
+                return AstVisitAction.Continue;
             }
 
-            if (commandAst.CommandElements[0] is not ExpandableStringExpressionAst expandableStringExpressionAst)
-            {
-                return null;
-            }
+            _references.AddReference(
+                SymbolType.Function,
+                CommandHelpers.StripModuleQualification(commandName, out _),
+                commandAst.CommandElements[0].Extent
+            );
 
-            return AstOperations.TryGetInferredValue(expandableStringExpressionAst, out string value) ? value : null;
+            return AstVisitAction.Continue;
         }
 
-        public override AstVisitAction VisitCommand(CommandAst commandAst)
+        // TODO: We should examine if we really want to constrain the extents to the name only. This
+        // means that highlighting only highlights the symbol name, but providing the whole extend
+        // means the whole function (or class etc.) gets highlighted, which seems to be a personal
+        // preference.
+        public override AstVisitAction VisitFunctionDefinition(FunctionDefinitionAst functionDefinitionAst)
         {
-            string? commandName = GetCommandName(commandAst);
-            if (string.IsNullOrEmpty(commandName))
+            // Extent for constructors and method trigger both this and VisitFunctionMember(). Covered in the latter.
+            // This will not exclude nested functions as they have ScriptBlockAst as parent
+            if (functionDefinitionAst.Parent is FunctionMemberAst)
             {
                 return AstVisitAction.Continue;
             }
 
+            // We only want the function name
+            IScriptExtent nameExtent = VisitorUtils.GetNameExtent(functionDefinitionAst);
             _references.AddReference(
-                CommandHelpers.StripModuleQualification(commandName, out _),
-                commandAst.CommandElements[0].Extent);
+                SymbolType.Function,
+                functionDefinitionAst.Name,
+                nameExtent);
+
+            return AstVisitAction.Continue;
+        }
+
+        public override AstVisitAction VisitCommandParameter(CommandParameterAst commandParameterAst)
+        {
+            _references.AddReference(
+                SymbolType.Parameter,
+                commandParameterAst.Extent.Text,
+                commandParameterAst.Extent);
 
             return AstVisitAction.Continue;
         }
@@ -111,10 +147,70 @@ public override AstVisitAction VisitVariableExpression(VariableExpressionAst var
             // TODO: Consider tracking unscoped variable references only when they declared within
             // the same function definition.
             _references.AddReference(
+                SymbolType.Variable,
                 $"${variableExpressionAst.VariablePath.UserPath}",
-                variableExpressionAst.Extent);
+                variableExpressionAst.Extent
+            );
+
+            return AstVisitAction.Continue;
+        }
+
+        public override AstVisitAction VisitTypeDefinition(TypeDefinitionAst typeDefinitionAst)
+        {
+            SymbolType symbolType = typeDefinitionAst.IsEnum ? SymbolType.Enum : SymbolType.Class;
+
+            // We only want the type name. Get start-location for name
+            IScriptExtent nameExtent = VisitorUtils.GetNameExtent(typeDefinitionAst);
+            _references.AddReference(symbolType, typeDefinitionAst.Name, nameExtent);
+
+            return AstVisitAction.Continue;
+        }
+
+        public override AstVisitAction VisitTypeExpression(TypeExpressionAst typeExpressionAst)
+        {
+            _references.AddReference(
+                SymbolType.Type,
+                typeExpressionAst.TypeName.Name,
+                typeExpressionAst.Extent);
 
             return AstVisitAction.Continue;
         }
+
+        public override AstVisitAction VisitTypeConstraint(TypeConstraintAst typeConstraintAst)
+        {
+            _references.AddReference(SymbolType.Type, typeConstraintAst.TypeName.Name, typeConstraintAst.Extent);
+
+            return AstVisitAction.Continue;
+        }
+
+        public override AstVisitAction VisitFunctionMember(FunctionMemberAst functionMemberAst)
+        {
+            SymbolType symbolType = functionMemberAst.IsConstructor
+                ? SymbolType.Constructor
+                : SymbolType.Method;
+
+            // We only want the method/ctor name. Get start-location for name
+            IScriptExtent nameExtent = VisitorUtils.GetNameExtent(functionMemberAst, true, false);
+            _references.AddReference(
+                symbolType,
+                VisitorUtils.GetMemberOverloadName(functionMemberAst, true, false),
+                nameExtent);
+
+            return AstVisitAction.Continue;
+        }
+
+        public override AstVisitAction VisitPropertyMember(PropertyMemberAst propertyMemberAst)
+        {
+            // We only want the property name. Get start-location for name
+            IScriptExtent nameExtent = VisitorUtils.GetNameExtent(propertyMemberAst, false);
+            _references.AddReference(
+                SymbolType.Property,
+                VisitorUtils.GetMemberOverloadName(propertyMemberAst, false),
+                nameExtent);
+
+            return AstVisitAction.Continue;
+        }
+
+        // TODO: What else can we implement?
     }
 }

src/PowerShellEditorServices/Services/Symbols/ScriptDocumentSymbolProvider.cs

@@ -16,7 +16,7 @@ internal class ScriptDocumentSymbolProvider : IDocumentSymbolProvider
     {
         string IDocumentSymbolProvider.ProviderId => nameof(ScriptDocumentSymbolProvider);
 
-        IEnumerable<ISymbolReference> IDocumentSymbolProvider.ProvideDocumentSymbols(
+        IEnumerable<SymbolReference> IDocumentSymbolProvider.ProvideDocumentSymbols(
             ScriptFile scriptFile)
         {
             // If we have an AST, then we know it's a PowerShell file

src/PowerShellEditorServices/Services/Symbols/SymbolReference.cs

@@ -100,6 +100,26 @@ public SymbolReference(
             //        "{0} {1}")
         }
 
+        public SymbolReference(
+            SymbolType symbolType,
+            string symbolName,
+            IScriptExtent scriptExtent,
+            ScriptFile file)
+        {
+            SymbolType = symbolType;
+            SymbolName = symbolName;
+            ScriptRegion = ScriptRegion.Create(scriptExtent);
+            FilePath = file.FilePath;
+            try
+            {
+                SourceLine = file.GetLine(ScriptRegion.StartLineNumber);
+            }
+            catch (System.ArgumentOutOfRangeException)
+            {
+                SourceLine = string.Empty;
+            }
+        }
+
         /// <summary>
         /// Constructs an instance of a SymbolReference
         /// </summary>

src/PowerShellEditorServices/Services/Symbols/SymbolsService.cs

@@ -110,23 +110,18 @@ public SymbolsService(
         /// Finds all the symbols in a file.
         /// </summary>
         /// <param name="scriptFile">The ScriptFile in which the symbol can be located.</param>
-        /// <returns></returns>
         public List<SymbolReference> FindSymbolsInFile(ScriptFile scriptFile)
         {
             Validate.IsNotNull(nameof(scriptFile), scriptFile);
 
-            List<SymbolReference> foundOccurrences = new();
+            List<SymbolReference> symbols = new();
             foreach (IDocumentSymbolProvider symbolProvider in GetDocumentSymbolProviders())
             {
-                foreach (SymbolReference reference in symbolProvider.ProvideDocumentSymbols(scriptFile))
-                {
-                    reference.SourceLine = scriptFile.GetLine(reference.ScriptRegion.StartLineNumber);
-                    reference.FilePath = scriptFile.FilePath;
-                    foundOccurrences.Add(reference);
-                }
+                // TODO: Each provider needs to set the source line and filepath.
+                symbols.AddRange(symbolProvider.ProvideDocumentSymbols(scriptFile));
             }
 
-            return foundOccurrences;
+            return symbols;
         }
 
         /// <summary>
@@ -174,6 +169,7 @@ public async Task<List<SymbolReference>> ScanForReferencesOfSymbol(
                 return null;
             }
 
+            // TODO: Should we handle aliases at a lower level?
             CommandHelpers.AliasMap aliases = await CommandHelpers.GetAliasesAsync(
                 _executionService,
                 cancellationToken).ConfigureAwait(false);
@@ -226,36 +222,21 @@ static string[] GetIdentifiers(string symbolName, SymbolType symbolType, Command
 
             foreach (ScriptFile file in _workspaceService.GetOpenedFiles())
             {
+                List<SymbolReference> fileReferences = new();
                 foreach (string targetIdentifier in allIdentifiers)
                 {
-                    if (!file.References.TryGetReferences(targetIdentifier, out ConcurrentBag<IScriptExtent> references))
+                    if (!file.References.TryGetReferences(targetIdentifier, out ConcurrentBag<SymbolReference> references))
                     {
                         continue;
                     }
 
-                    foreach (IScriptExtent extent in references.OrderBy(e => e.StartOffset))
-                    {
-                        SymbolReference reference = new(
-                            SymbolType.Function,
-                            foundSymbol.SymbolName,
-                            extent);
-
-                        try
-                        {
-                            reference.SourceLine = file.GetLine(reference.ScriptRegion.StartLineNumber);
-                        }
-                        catch (ArgumentOutOfRangeException e)
-                        {
-                            reference.SourceLine = string.Empty;
-                            _logger.LogException("Found reference is out of range in script file", e);
-                        }
-                        reference.FilePath = file.FilePath;
-                        symbolReferences.Add(reference);
-                    }
+                    fileReferences.AddRange(references);
 
                     await Task.Yield();
                     cancellationToken.ThrowIfCancellationRequested();
                 }
+
+                symbolReferences.AddRange(fileReferences.OrderBy(symbol => symbol.ScriptRegion.StartOffset));
             }
 
             return symbolReferences;
@@ -283,7 +264,12 @@ public static IReadOnlyList<SymbolReference> FindOccurrencesInFile(
                 return null;
             }
 
-            return AstOperations.FindReferencesOfSymbol(file.ScriptAst, foundSymbol).ToArray();
+            if (file.References.TryGetReferences(foundSymbol.SymbolName, out ConcurrentBag<SymbolReference> references))
+            {
+                return references.OrderBy(symbol => symbol.ScriptRegion.StartOffset).ToArray();
+            }
+
+            return null;
         }
 
         /// <summary>
@@ -425,19 +411,17 @@ public async Task<SymbolReference> GetDefinitionOfSymbolAsync(
             (Dictionary<string, List<string>> _, Dictionary<string, string> aliasToCmdlets) =
                 await CommandHelpers.GetAliasesAsync(_executionService).ConfigureAwait(false);
 
-            if (aliasToCmdlets.ContainsKey(foundSymbol.SymbolName))
+            if (aliasToCmdlets.TryGetValue(foundSymbol.SymbolName, out string value))
             {
                 foundSymbol = new SymbolReference(
                     foundSymbol.SymbolType,
-                    aliasToCmdlets[foundSymbol.SymbolName],
+value,
                     foundSymbol.ScriptRegion,
                     foundSymbol.FilePath,
                     foundSymbol.SourceLine);
             }
 
-            ScriptFile[] referencedFiles =
-                _workspaceService.ExpandScriptReferences(
-                    sourceFile);
+            ScriptFile[] referencedFiles = _workspaceService.ExpandScriptReferences(sourceFile);
 
             HashSet<string> filesSearched = new(StringComparer.OrdinalIgnoreCase);
 
@@ -461,7 +445,11 @@ public async Task<SymbolReference> GetDefinitionOfSymbolAsync(
                     string dotSourcedPath = GetDotSourcedPath(foundSymbol, scriptFile);
                     if (scriptFile.FilePath == dotSourcedPath)
                     {
-                        foundDefinition = new SymbolReference(SymbolType.Function, foundSymbol.SymbolName, scriptFile.ScriptAst.Extent, scriptFile.FilePath);
+                        foundDefinition = new SymbolReference(
+                            SymbolType.Function,
+                            foundSymbol.SymbolName,
+                            scriptFile.ScriptAst.Extent,
+                            scriptFile.FilePath);
                         break;
                     }
                 }