Track declarations on SymbolType so we can use ReferenceTable more

This lets us (nearly) replace `FindSymbolsInDocument` by using the references
(symbols) from the `ScriptFile`.

Author: andyleejordan

src/PowerShellEditorServices/Services/CodeLens/ReferencesCodeLensProvider.cs

@@ -84,19 +84,18 @@ SymbolType.Class or
         }
 
         /// <summary>
-        /// Take a codelens and create a new codelens object with updated references.
+        /// 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>
         /// <param name="cancellationToken"></param>
-        /// <returns>A new code lens object describing the same data as the old one but with updated references.</returns>
+        /// <returns>A new CodeLens object describing the same data as the old one but with updated references.</returns>
         public async Task<CodeLens> ResolveCodeLens(
             CodeLens codeLens,
             ScriptFile scriptFile,
             CancellationToken cancellationToken)
         {
-            ScriptFile[] references = _workspaceService.ExpandScriptReferences(
-                scriptFile);
+            ScriptFile[] references = _workspaceService.ExpandScriptReferences(scriptFile);
 
             SymbolReference foundSymbol = SymbolsService.FindSymbolDefinitionAtLocation(
                 scriptFile,
@@ -123,7 +122,8 @@ await _symbolsService.ScanForReferencesOfSymbol(
                     // so it's helpful to add some yields.
                     await Task.Yield();
                     cancellationToken.ThrowIfCancellationRequested();
-                    if (IsReferenceDefinition(foundSymbol, foundReference))
+
+                    if (foundReference.IsDeclaration)
                     {
                         continue;
                     }
@@ -165,27 +165,6 @@ await _symbolsService.ScanForReferencesOfSymbol(
             };
         }
 
-        /// <summary>
-        /// Check whether a SymbolReference is the actual definition of that symbol.
-        /// </summary>
-        /// <param name="definition">The symbol definition that may be referenced.</param>
-        /// <param name="reference">The reference symbol to check.</param>
-        /// <returns>True if the reference is not a reference to the definition, false otherwise.</returns>
-        private static bool IsReferenceDefinition(
-            SymbolReference definition,
-            SymbolReference reference)
-        {
-            // First check if we are in the same file as the definition. if we are...
-            // check if it's on the same line number.
-
-            // TODO: Do we care about two symbol definitions of the same name?
-            // if we do, how could we possibly know that a reference in one file is a reference
-            // of a particular symbol definition?
-            return
-                definition.FilePath == reference.FilePath &&
-                definition.ScriptRegion.StartLineNumber == reference.ScriptRegion.StartLineNumber;
-        }
-
         /// <summary>
         /// Get the code lens header for the number of references on a definition,
         /// given the number of references.

src/PowerShellEditorServices/Services/Symbols/ReferenceTable.cs

@@ -71,9 +71,9 @@ internal void EnsureInitialized()
         _parent.ScriptAst.Visit(new ReferenceVisitor(this));
     }
 
-    private void AddReference(SymbolType type, string name, IScriptExtent extent)
+    private void AddReference(SymbolType type, string name, IScriptExtent extent, bool isDeclaration = false)
     {
-        SymbolReference symbol = new(type, name, extent, _parent);
+        SymbolReference symbol = new(type, name, extent, _parent, isDeclaration);
         _symbolReferences.AddOrUpdate(
             name,
             _ => new ConcurrentBag<SymbolReference> { symbol },
@@ -123,7 +123,8 @@ public override AstVisitAction VisitFunctionDefinition(FunctionDefinitionAst fun
             _references.AddReference(
                 SymbolType.Function,
                 functionDefinitionAst.Name,
-                nameExtent);
+                nameExtent,
+                isDeclaration: true);
 
             return AstVisitAction.Continue;
         }

src/PowerShellEditorServices/Services/Symbols/ScriptDocumentSymbolProvider.cs

@@ -2,8 +2,6 @@
 // Licensed under the MIT License.
 
 using System.Collections.Generic;
-using System.Linq;
-using System.Management.Automation.Language;
 using Microsoft.PowerShell.EditorServices.Services.TextDocument;
 
 namespace Microsoft.PowerShell.EditorServices.Services.Symbols
@@ -17,25 +15,6 @@ internal class ScriptDocumentSymbolProvider : IDocumentSymbolProvider
         string IDocumentSymbolProvider.ProviderId => nameof(ScriptDocumentSymbolProvider);
 
         IEnumerable<SymbolReference> IDocumentSymbolProvider.ProvideDocumentSymbols(
-            ScriptFile scriptFile)
-        {
-            // If we have an AST, then we know it's a PowerShell file
-            // so lets try to find symbols in the document.
-            return scriptFile?.ScriptAst != null
-                ? FindSymbolsInDocument(scriptFile.ScriptAst)
-                : Enumerable.Empty<SymbolReference>();
-        }
-
-        /// <summary>
-        /// Finds all symbols in a script
-        /// </summary>
-        /// <param name="scriptAst">The abstract syntax tree of the given script</param>
-        /// <returns>A collection of SymbolReference objects</returns>
-        public static IEnumerable<SymbolReference> FindSymbolsInDocument(Ast scriptAst)
-        {
-            FindSymbolsVisitor findSymbolsVisitor = new();
-            scriptAst.Visit(findSymbolsVisitor);
-            return findSymbolsVisitor.SymbolReferences;
-        }
+            ScriptFile scriptFile) => scriptFile.References.GetAllReferences();
     }
 }

src/PowerShellEditorServices/Services/Symbols/SymbolReference.cs

@@ -25,6 +25,8 @@ internal record SymbolReference
 
         public string FilePath { get; internal set; }
 
+        public bool IsDeclaration { get; }
+
         /// <summary>
         /// Constructs and instance of a SymbolReference
         /// </summary>
@@ -33,26 +35,30 @@ internal record SymbolReference
         /// <param name="scriptExtent">The script extent of the symbol</param>
         /// <param name="filePath">The file path of the symbol</param>
         /// <param name="sourceLine">The line contents of the given symbol (defaults to empty string)</param>
+        /// <param name="isDeclaration">True if this reference is the definition of the symbol</param>
         public SymbolReference(
             SymbolType symbolType,
             string symbolName,
             IScriptExtent scriptExtent,
             string filePath = "",
-            string sourceLine = "")
+            string sourceLine = "",
+            bool isDeclaration = false)
         {
             // TODO: Verify params
             SymbolType = symbolType;
             SymbolName = symbolName;
             ScriptRegion = new(scriptExtent);
             FilePath = filePath;
             SourceLine = sourceLine;
+            IsDeclaration = isDeclaration;
         }
 
         public SymbolReference(
             SymbolType symbolType,
             string symbolName,
             IScriptExtent scriptExtent,
-            ScriptFile file)
+            ScriptFile file,
+            bool isDeclaration)
         {
             SymbolType = symbolType;
             SymbolName = symbolName;
@@ -66,6 +72,7 @@ public SymbolReference(
             {
                 SourceLine = string.Empty;
             }
+            IsDeclaration = isDeclaration;
         }
 
         /// <summary>
@@ -74,7 +81,7 @@ public SymbolReference(
         /// <param name="symbolType">The higher level type of the symbol</param>
         /// <param name="scriptExtent">The script extent of the symbol</param>
         public SymbolReference(SymbolType symbolType, IScriptExtent scriptExtent)
-            : this(symbolType, scriptExtent.Text, scriptExtent, scriptExtent.File, "")
+            : this(symbolType, scriptExtent.Text, scriptExtent, scriptExtent.File)
         {
         }
     }

src/PowerShellEditorServices/Services/Symbols/SymbolType.cs

@@ -94,6 +94,7 @@ internal static SymbolKind GetSymbolKind(SymbolType symbolType)
                 SymbolType.Method => SymbolKind.Method,
                 SymbolType.Property => SymbolKind.Property,
                 SymbolType.EnumMember => SymbolKind.EnumMember,
+                SymbolType.Variable => SymbolKind.Variable,
                 // TODO: More delicately handle the other symbol types.
                 _ => SymbolKind.Variable,
             };

src/PowerShellEditorServices/Services/Symbols/SymbolsService.cs

@@ -153,7 +153,7 @@ public static SymbolReference FindSymbolAtLocation(
         }
 
         /// <summary>
-        /// Finds all the references of a symbol
+        /// Finds all the references of a symbol (excluding definitions)
         /// </summary>
         /// <param name="foundSymbol">The symbol to find all references for</param>
         /// <param name="referencedFiles">An array of scriptFiles to search for references in</param>

src/PowerShellEditorServices/Services/TextDocument/Handlers/ReferencesHandler.cs

@@ -47,18 +47,25 @@ await _symbolsService.ScanForReferencesOfSymbol(
                     _workspaceService.ExpandScriptReferences(scriptFile),
                     cancellationToken).ConfigureAwait(false);
 
-            List<Location> locations = new();
+            if (referencesResult is null)
+            {
+                return new LocationContainer();
+            }
 
-            if (referencesResult != null)
+            List<Location> locations = new();
+            foreach (SymbolReference foundReference in referencesResult)
             {
-                foreach (SymbolReference foundReference in referencesResult)
+                // Respect the request's setting to include declarations.
+                if (!request.Context.IncludeDeclaration && foundReference.IsDeclaration)
                 {
-                    locations.Add(new Location
-                    {
-                        Uri = DocumentUri.From(foundReference.FilePath),
-                        Range = foundReference.ScriptRegion.ToRange()
-                    });
+                    continue;
                 }
+
+                locations.Add(new Location
+                {
+                    Uri = DocumentUri.From(foundReference.FilePath),
+                    Range = foundReference.ScriptRegion.ToRange()
+                });
             }
 
             return new LocationContainer(locations);