Filter symbol references to type matches

Author: andyleejordan

src/PowerShellEditorServices/Services/Symbols/ReferenceTable.cs

@@ -42,10 +42,13 @@ public void TagAsChanged()
     /// </summary>
     private bool IsInitialized => !_symbolReferences.IsEmpty || _isInited;
 
-    internal bool TryGetReferences(string command, out ConcurrentBag<SymbolReference>? references)
+    internal IEnumerable<SymbolReference> TryGetReferences(SymbolReference? symbol)
     {
         EnsureInitialized();
-        return _symbolReferences.TryGetValue(command, out references);
+        return symbol is not null
+            && _symbolReferences.TryGetValue(symbol.SymbolName, out ConcurrentBag<SymbolReference>? bag)
+                ? bag.Where(i => SymbolTypeUtils.SymbolTypeMatches(symbol.SymbolType, i.SymbolType))
+                : Enumerable.Empty<SymbolReference>();
     }
 
     internal SymbolReference? TryGetSymbolAtPosition(int line, int column) => GetAllReferences()

src/PowerShellEditorServices/Services/Symbols/SymbolType.cs

@@ -99,5 +99,15 @@ internal static SymbolKind GetSymbolKind(SymbolType symbolType)
                 _ => SymbolKind.Variable,
             };
         }
+
+        // Provides a partial equivalence between type constraints and custom types.
+        internal static bool SymbolTypeMatches(SymbolType left, SymbolType right)
+        {
+            return left == right
+                || (left is SymbolType.Class or SymbolType.Enum or SymbolType.Type
+                && right is SymbolType.Class or SymbolType.Enum or SymbolType.Type)
+                || (left is SymbolType.EnumMember or SymbolType.Property
+                && right is SymbolType.EnumMember or SymbolType.Property);
+        }
     }
 }

src/PowerShellEditorServices/Services/Symbols/SymbolsService.cs

@@ -188,14 +188,7 @@ public async Task<IEnumerable<SymbolReference>> ScanForReferencesOfSymbolAsync(
                 {
                     await Task.Yield();
                     cancellationToken.ThrowIfCancellationRequested();
-
-                    _ = file.References.TryGetReferences(targetIdentifier, out ConcurrentBag<SymbolReference>? references);
-                    if (references is null)
-                    {
-                        continue;
-                    }
-
-                    symbols.AddRange(references);
+                    symbols.AddRange(file.References.TryGetReferences(symbol with { SymbolName = targetIdentifier }));
                 }
             }
 
@@ -205,19 +198,10 @@ public async Task<IEnumerable<SymbolReference>> ScanForReferencesOfSymbolAsync(
         /// <summary>
         /// Finds all the occurrences of a symbol in the script given a file location.
         /// </summary>
-        public static IEnumerable<SymbolReference>? FindOccurrencesInFile(
-            ScriptFile scriptFile, int line, int column)
-        {
-            SymbolReference? symbol = FindSymbolAtLocation(scriptFile, line, column);
-
-            if (symbol is null)
-            {
-                return null;
-            }
-
-            scriptFile.References.TryGetReferences(symbol.SymbolName, out ConcurrentBag<SymbolReference>? references);
-            return references;
-        }
+        public static IEnumerable<SymbolReference> FindOccurrencesInFile(
+            ScriptFile scriptFile, int line, int column) => scriptFile
+                .References
+                .TryGetReferences(FindSymbolAtLocation(scriptFile, line, column));
 
         /// <summary>
         /// Finds the symbol at the location and returns it if it's a declaration.
@@ -236,15 +220,9 @@ public async Task<IEnumerable<SymbolReference>> ScanForReferencesOfSymbolAsync(
             ScriptFile scriptFile, int line, int column)
         {
             SymbolReference? symbol = FindSymbolAtLocation(scriptFile, line, column);
-            if (symbol is null)
-            {
-                return Task.FromResult<SymbolDetails?>(null);
-            }
-
-            return SymbolDetails.CreateAsync(
-                symbol,
-                _runspaceContext.CurrentRunspace,
-                _executionService);
+            return symbol is null
+                ? Task.FromResult<SymbolDetails?>(null)
+                : SymbolDetails.CreateAsync(symbol, _runspaceContext.CurrentRunspace, _executionService);
         }
 
         /// <summary>
@@ -310,34 +288,18 @@ public async Task<IEnumerable<SymbolReference>> GetDefinitionOfSymbolAsync(
             CancellationToken cancellationToken = default)
         {
             List<SymbolReference> declarations = new();
-            _ = scriptFile.References.TryGetReferences(symbol.SymbolName, out ConcurrentBag<SymbolReference>? symbols);
-            if (symbols is not null)
-            {
-                foreach (SymbolReference foundReference in symbols)
-                {
-                    if (foundReference.IsDeclaration)
-                    {
-                        _logger.LogDebug($"Found possible declaration in same file ${foundReference}");
-                        declarations.Add(foundReference);
-                    }
-                }
-            }
-
+            declarations.AddRange(scriptFile.References.TryGetReferences(symbol).Where(i => i.IsDeclaration));
             if (declarations.Any())
             {
+                _logger.LogDebug($"Found possible declaration in same file ${declarations}");
                 return declarations;
             }
 
-            foreach (SymbolReference foundReference in await ScanForReferencesOfSymbolAsync(
-                    symbol, cancellationToken).ConfigureAwait(false))
-            {
-                if (foundReference.IsDeclaration)
-                {
-                    _logger.LogDebug($"Found possible declaration in workspace ${foundReference}");
-                    declarations.Add(foundReference);
-                }
-            }
+            IEnumerable<SymbolReference> references =
+                await ScanForReferencesOfSymbolAsync(symbol, cancellationToken).ConfigureAwait(false);
+            declarations.AddRange(references.Where(i => i.IsDeclaration));
 
+            _logger.LogDebug($"Found possible declaration in workspace ${declarations}");
             return declarations;
         }
 

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

@@ -1,7 +1,7 @@
 // Copyright (c) Microsoft Corporation.
 // Licensed under the MIT License.
 
-using System.Collections.Concurrent;
+using System.Collections.Generic;
 using System.Linq;
 using Microsoft.Extensions.Logging.Abstractions;
 using Microsoft.PowerShell.EditorServices.Services;
@@ -46,9 +46,8 @@ public void CanFindReferencesOfSymbolAtPosition(int line, int column, Range[] sy
         {
             SymbolReference symbol = scriptFile.References.TryGetSymbolAtPosition(line, column);
 
-            Assert.True(scriptFile.References.TryGetReferences(
-                symbol.SymbolName,
-                out ConcurrentBag<SymbolReference> references));
+            IEnumerable<SymbolReference> references = scriptFile.References.TryGetReferences(symbol);
+            Assert.NotEmpty(references);
 
             int positionsIndex = 0;
             foreach (SymbolReference reference in references.OrderBy((i) => i.ScriptRegion.ToRange().Start))