Skip to content

Commit 8b2d1b2

Browse files
andyleejordanandyleejordan
committed
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.
1 parent f5a0223 commit 8b2d1b2

File tree

16 files changed

+251
-410
lines changed

16 files changed

+251
-410
lines changed

‎src/PowerShellEditorServices/Services/Symbols/IDocumentSymbolProvider.cs‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -20,6 +20,6 @@ internal interface IDocumentSymbolProvider
2020
/// The document for which SymbolReferences should be provided.
2121
/// </param>
2222
/// <returns>An IEnumerable collection of SymbolReferences.</returns>
23-
IEnumerable<ISymbolReference> ProvideDocumentSymbols(ScriptFile scriptFile);
23+
IEnumerable<SymbolReference> ProvideDocumentSymbols(ScriptFile scriptFile);
2424
}
2525
}

‎src/PowerShellEditorServices/Services/Symbols/PesterDocumentSymbolProvider.cs‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -17,7 +17,7 @@ internal class PesterDocumentSymbolProvider : IDocumentSymbolProvider
1717
{
1818
string IDocumentSymbolProvider.ProviderId => nameof(PesterDocumentSymbolProvider);
1919

20-
IEnumerable<ISymbolReference> IDocumentSymbolProvider.ProvideDocumentSymbols(
20+
IEnumerable<SymbolReference> IDocumentSymbolProvider.ProvideDocumentSymbols(
2121
ScriptFile scriptFile)
2222
{
2323
if (!scriptFile.FilePath.EndsWith(

‎src/PowerShellEditorServices/Services/Symbols/PsdDocumentSymbolProvider.cs‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -17,7 +17,7 @@ internal class PsdDocumentSymbolProvider : IDocumentSymbolProvider
1717
{
1818
string IDocumentSymbolProvider.ProviderId => nameof(PsdDocumentSymbolProvider);
1919

20-
IEnumerable<ISymbolReference> IDocumentSymbolProvider.ProvideDocumentSymbols(
20+
IEnumerable<SymbolReference> IDocumentSymbolProvider.ProvideDocumentSymbols(
2121
ScriptFile scriptFile)
2222
{
2323
if ((scriptFile.FilePath?.EndsWith(".psd1", StringComparison.OrdinalIgnoreCase) == true) ||

‎src/PowerShellEditorServices/Services/Symbols/ReferenceTable.cs‎

Lines changed: 118 additions & 22 deletions
Original file line numberDiff line numberDiff line change
@@ -9,6 +9,8 @@
99
using Microsoft.PowerShell.EditorServices.Services.TextDocument;
1010
using Microsoft.PowerShell.EditorServices.Services.PowerShell.Utility;
1111
using Microsoft.PowerShell.EditorServices.Services.Symbols;
12+
using System.Collections.Generic;
13+
using Microsoft.PowerShell.EditorServices.Utility;
1214

1315
namespace Microsoft.PowerShell.EditorServices.Services;
1416

@@ -19,7 +21,7 @@ internal sealed class ReferenceTable
1921
{
2022
private readonly ScriptFile _parent;
2123

22-
private readonly ConcurrentDictionary<string, ConcurrentBag<IScriptExtent>> _symbolReferences = new(StringComparer.OrdinalIgnoreCase);
24+
private readonly ConcurrentDictionary<string, ConcurrentBag<SymbolReference>> _symbolReferences = new(StringComparer.OrdinalIgnoreCase);
2325

2426
private bool _isInited;
2527

@@ -41,12 +43,24 @@ public void TagAsChanged()
4143
/// </summary>
4244
private bool IsInitialized => !_symbolReferences.IsEmpty || _isInited;
4345

44-
internal bool TryGetReferences(string command, out ConcurrentBag<IScriptExtent>? references)
46+
internal bool TryGetReferences(string command, out ConcurrentBag<SymbolReference>? references)
4547
{
4648
EnsureInitialized();
4749
return _symbolReferences.TryGetValue(command, out references);
4850
}
4951

52+
// TODO: Should this be improved, or pre-sorted?
53+
internal IReadOnlyList<SymbolReference> GetAllReferences()
54+
{
55+
EnsureInitialized();
56+
List<SymbolReference> allReferences = new();
57+
foreach (ConcurrentBag<SymbolReference> bag in _symbolReferences.Values)
58+
{
59+
allReferences.AddRange(bag);
60+
}
61+
return allReferences;
62+
}
63+
5064
internal void EnsureInitialized()
5165
{
5266
if (IsInitialized)
@@ -57,51 +71,73 @@ internal void EnsureInitialized()
5771
_parent.ScriptAst.Visit(new ReferenceVisitor(this));
5872
}
5973

60-
private void AddReference(string symbol, IScriptExtent extent)
74+
private void AddReference(SymbolType type, string name, IScriptExtent extent)
6175
{
76+
SymbolReference symbol = new(type, name, extent, _parent);
6277
_symbolReferences.AddOrUpdate(
63-
symbol,
64-
_ => new ConcurrentBag<IScriptExtent> { extent },
78+
name,
79+
_ => new ConcurrentBag<SymbolReference> { symbol },
6580
(_, existing) =>
6681
{
67-
existing.Add(extent);
82+
existing.Add(symbol);
6883
return existing;
6984
});
7085
}
7186

72-
private sealed class ReferenceVisitor : AstVisitor
87+
// TODO: Should we move this to AstOperations.cs? It is highly coupled to `ReferenceTable`,
88+
// perhaps it doesn't have to be.
89+
private sealed class ReferenceVisitor : AstVisitor2
7390
{
7491
private readonly ReferenceTable _references;
7592

7693
public ReferenceVisitor(ReferenceTable references) => _references = references;
7794

78-
private static string? GetCommandName(CommandAst commandAst)
95+
public override AstVisitAction VisitCommand(CommandAst commandAst)
7996
{
80-
string commandName = commandAst.GetCommandName();
81-
if (!string.IsNullOrEmpty(commandName))
97+
string? commandName = VisitorUtils.GetCommandName(commandAst);
98+
if (string.IsNullOrEmpty(commandName))
8299
{
83-
return commandName;
100+
return AstVisitAction.Continue;
84101
}
85102

86-
if (commandAst.CommandElements[0] is not ExpandableStringExpressionAst expandableStringExpressionAst)
87-
{
88-
return null;
89-
}
103+
_references.AddReference(
104+
SymbolType.Function,
105+
CommandHelpers.StripModuleQualification(commandName, out _),
106+
commandAst.CommandElements[0].Extent
107+
);
90108

91-
return AstOperations.TryGetInferredValue(expandableStringExpressionAst, out string value) ? value : null;
109+
return AstVisitAction.Continue;
92110
}
93111

94-
public override AstVisitAction VisitCommand(CommandAst commandAst)
112+
// TODO: We should examine if we really want to constrain the extents to the name only. This
113+
// means that highlighting only highlights the symbol name, but providing the whole extend
114+
// means the whole function (or class etc.) gets highlighted, which seems to be a personal
115+
// preference.
116+
public override AstVisitAction VisitFunctionDefinition(FunctionDefinitionAst functionDefinitionAst)
95117
{
96-
string? commandName = GetCommandName(commandAst);
97-
if (string.IsNullOrEmpty(commandName))
118+
// Extent for constructors and method trigger both this and VisitFunctionMember(). Covered in the latter.
119+
// This will not exclude nested functions as they have ScriptBlockAst as parent
120+
if (functionDefinitionAst.Parent is FunctionMemberAst)
98121
{
99122
return AstVisitAction.Continue;
100123
}
101124

125+
// We only want the function name
126+
IScriptExtent nameExtent = VisitorUtils.GetNameExtent(functionDefinitionAst);
102127
_references.AddReference(
103-
CommandHelpers.StripModuleQualification(commandName, out _),
104-
commandAst.CommandElements[0].Extent);
128+
SymbolType.Function,
129+
functionDefinitionAst.Name,
130+
nameExtent);
131+
132+
return AstVisitAction.Continue;
133+
}
134+
135+
public override AstVisitAction VisitCommandParameter(CommandParameterAst commandParameterAst)
136+
{
137+
_references.AddReference(
138+
SymbolType.Parameter,
139+
commandParameterAst.Extent.Text,
140+
commandParameterAst.Extent);
105141

106142
return AstVisitAction.Continue;
107143
}
@@ -111,10 +147,70 @@ public override AstVisitAction VisitVariableExpression(VariableExpressionAst var
111147
// TODO: Consider tracking unscoped variable references only when they declared within
112148
// the same function definition.
113149
_references.AddReference(
150+
SymbolType.Variable,
114151
$"${variableExpressionAst.VariablePath.UserPath}",
115-
variableExpressionAst.Extent);
152+
variableExpressionAst.Extent
153+
);
154+
155+
return AstVisitAction.Continue;
156+
}
157+
158+
public override AstVisitAction VisitTypeDefinition(TypeDefinitionAst typeDefinitionAst)
159+
{
160+
SymbolType symbolType = typeDefinitionAst.IsEnum ? SymbolType.Enum : SymbolType.Class;
161+
162+
// We only want the type name. Get start-location for name
163+
IScriptExtent nameExtent = VisitorUtils.GetNameExtent(typeDefinitionAst);
164+
_references.AddReference(symbolType, typeDefinitionAst.Name, nameExtent);
165+
166+
return AstVisitAction.Continue;
167+
}
168+
169+
public override AstVisitAction VisitTypeExpression(TypeExpressionAst typeExpressionAst)
170+
{
171+
_references.AddReference(
172+
SymbolType.Type,
173+
typeExpressionAst.TypeName.Name,
174+
typeExpressionAst.Extent);
116175

117176
return AstVisitAction.Continue;
118177
}
178+
179+
public override AstVisitAction VisitTypeConstraint(TypeConstraintAst typeConstraintAst)
180+
{
181+
_references.AddReference(SymbolType.Type, typeConstraintAst.TypeName.Name, typeConstraintAst.Extent);
182+
183+
return AstVisitAction.Continue;
184+
}
185+
186+
public override AstVisitAction VisitFunctionMember(FunctionMemberAst functionMemberAst)
187+
{
188+
SymbolType symbolType = functionMemberAst.IsConstructor
189+
? SymbolType.Constructor
190+
: SymbolType.Method;
191+
192+
// We only want the method/ctor name. Get start-location for name
193+
IScriptExtent nameExtent = VisitorUtils.GetNameExtent(functionMemberAst, true, false);
194+
_references.AddReference(
195+
symbolType,
196+
VisitorUtils.GetMemberOverloadName(functionMemberAst, true, false),
197+
nameExtent);
198+
199+
return AstVisitAction.Continue;
200+
}
201+
202+
public override AstVisitAction VisitPropertyMember(PropertyMemberAst propertyMemberAst)
203+
{
204+
// We only want the property name. Get start-location for name
205+
IScriptExtent nameExtent = VisitorUtils.GetNameExtent(propertyMemberAst, false);
206+
_references.AddReference(
207+
SymbolType.Property,
208+
VisitorUtils.GetMemberOverloadName(propertyMemberAst, false),
209+
nameExtent);
210+
211+
return AstVisitAction.Continue;
212+
}
213+
214+
// TODO: What else can we implement?
119215
}
120216
}

‎src/PowerShellEditorServices/Services/Symbols/ScriptDocumentSymbolProvider.cs‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -16,7 +16,7 @@ internal class ScriptDocumentSymbolProvider : IDocumentSymbolProvider
1616
{
1717
string IDocumentSymbolProvider.ProviderId => nameof(ScriptDocumentSymbolProvider);
1818

19-
IEnumerable<ISymbolReference> IDocumentSymbolProvider.ProvideDocumentSymbols(
19+
IEnumerable<SymbolReference> IDocumentSymbolProvider.ProvideDocumentSymbols(
2020
ScriptFile scriptFile)
2121
{
2222
// If we have an AST, then we know it's a PowerShell file

‎src/PowerShellEditorServices/Services/Symbols/SymbolReference.cs‎

Lines changed: 20 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -100,6 +100,26 @@ public SymbolReference(
100100
// "{0} {1}")
101101
}
102102

103+
public SymbolReference(
104+
SymbolType symbolType,
105+
string symbolName,
106+
IScriptExtent scriptExtent,
107+
ScriptFile file)
108+
{
109+
SymbolType = symbolType;
110+
SymbolName = symbolName;
111+
ScriptRegion = ScriptRegion.Create(scriptExtent);
112+
FilePath = file.FilePath;
113+
try
114+
{
115+
SourceLine = file.GetLine(ScriptRegion.StartLineNumber);
116+
}
117+
catch (System.ArgumentOutOfRangeException)
118+
{
119+
SourceLine = string.Empty;
120+
}
121+
}
122+
103123
/// <summary>
104124
/// Constructs an instance of a SymbolReference
105125
/// </summary>

‎src/PowerShellEditorServices/Services/Symbols/SymbolsService.cs‎

Lines changed: 24 additions & 36 deletions
Original file line numberDiff line numberDiff line change
@@ -110,23 +110,18 @@ public SymbolsService(
110110
/// Finds all the symbols in a file.
111111
/// </summary>
112112
/// <param name="scriptFile">The ScriptFile in which the symbol can be located.</param>
113-
/// <returns></returns>
114113
public List<SymbolReference> FindSymbolsInFile(ScriptFile scriptFile)
115114
{
116115
Validate.IsNotNull(nameof(scriptFile), scriptFile);
117116

118-
List<SymbolReference> foundOccurrences = new();
117+
List<SymbolReference> symbols = new();
119118
foreach (IDocumentSymbolProvider symbolProvider in GetDocumentSymbolProviders())
120119
{
121-
foreach (SymbolReference reference in symbolProvider.ProvideDocumentSymbols(scriptFile))
122-
{
123-
reference.SourceLine = scriptFile.GetLine(reference.ScriptRegion.StartLineNumber);
124-
reference.FilePath = scriptFile.FilePath;
125-
foundOccurrences.Add(reference);
126-
}
120+
// TODO: Each provider needs to set the source line and filepath.
121+
symbols.AddRange(symbolProvider.ProvideDocumentSymbols(scriptFile));
127122
}
128123

129-
return foundOccurrences;
124+
return symbols;
130125
}
131126

132127
/// <summary>
@@ -174,6 +169,7 @@ public async Task<List<SymbolReference>> ScanForReferencesOfSymbol(
174169
return null;
175170
}
176171

172+
// TODO: Should we handle aliases at a lower level?
177173
CommandHelpers.AliasMap aliases = await CommandHelpers.GetAliasesAsync(
178174
_executionService,
179175
cancellationToken).ConfigureAwait(false);
@@ -226,36 +222,21 @@ static string[] GetIdentifiers(string symbolName, SymbolType symbolType, Command
226222

227223
foreach (ScriptFile file in _workspaceService.GetOpenedFiles())
228224
{
225+
List<SymbolReference> fileReferences = new();
229226
foreach (string targetIdentifier in allIdentifiers)
230227
{
231-
if (!file.References.TryGetReferences(targetIdentifier, out ConcurrentBag<IScriptExtent> references))
228+
if (!file.References.TryGetReferences(targetIdentifier, out ConcurrentBag<SymbolReference> references))
232229
{
233230
continue;
234231
}
235232

236-
foreach (IScriptExtent extent in references.OrderBy(e => e.StartOffset))
237-
{
238-
SymbolReference reference = new(
239-
SymbolType.Function,
240-
foundSymbol.SymbolName,
241-
extent);
242-
243-
try
244-
{
245-
reference.SourceLine = file.GetLine(reference.ScriptRegion.StartLineNumber);
246-
}
247-
catch (ArgumentOutOfRangeException e)
248-
{
249-
reference.SourceLine = string.Empty;
250-
_logger.LogException("Found reference is out of range in script file", e);
251-
}
252-
reference.FilePath = file.FilePath;
253-
symbolReferences.Add(reference);
254-
}
233+
fileReferences.AddRange(references);
255234

256235
await Task.Yield();
257236
cancellationToken.ThrowIfCancellationRequested();
258237
}
238+
239+
symbolReferences.AddRange(fileReferences.OrderBy(symbol => symbol.ScriptRegion.StartOffset));
259240
}
260241

261242
return symbolReferences;
@@ -283,7 +264,12 @@ public static IReadOnlyList<SymbolReference> FindOccurrencesInFile(
283264
return null;
284265
}
285266

286-
return AstOperations.FindReferencesOfSymbol(file.ScriptAst, foundSymbol).ToArray();
267+
if (file.References.TryGetReferences(foundSymbol.SymbolName, out ConcurrentBag<SymbolReference> references))
268+
{
269+
return references.OrderBy(symbol => symbol.ScriptRegion.StartOffset).ToArray();
270+
}
271+
272+
return null;
287273
}
288274

289275
/// <summary>
@@ -425,19 +411,17 @@ public async Task<SymbolReference> GetDefinitionOfSymbolAsync(
425411
(Dictionary<string, List<string>> _, Dictionary<string, string> aliasToCmdlets) =
426412
await CommandHelpers.GetAliasesAsync(_executionService).ConfigureAwait(false);
427413

428-
if (aliasToCmdlets.ContainsKey(foundSymbol.SymbolName))
414+
if (aliasToCmdlets.TryGetValue(foundSymbol.SymbolName, out string value))
429415
{
430416
foundSymbol = new SymbolReference(
431417
foundSymbol.SymbolType,
432-
aliasToCmdlets[foundSymbol.SymbolName],
418+
value,
433419
foundSymbol.ScriptRegion,
434420
foundSymbol.FilePath,
435421
foundSymbol.SourceLine);
436422
}
437423

438-
ScriptFile[] referencedFiles =
439-
_workspaceService.ExpandScriptReferences(
440-
sourceFile);
424+
ScriptFile[] referencedFiles = _workspaceService.ExpandScriptReferences(sourceFile);
441425

442426
HashSet<string> filesSearched = new(StringComparer.OrdinalIgnoreCase);
443427

@@ -461,7 +445,11 @@ public async Task<SymbolReference> GetDefinitionOfSymbolAsync(
461445
string dotSourcedPath = GetDotSourcedPath(foundSymbol, scriptFile);
462446
if (scriptFile.FilePath == dotSourcedPath)
463447
{
464-
foundDefinition = new SymbolReference(SymbolType.Function, foundSymbol.SymbolName, scriptFile.ScriptAst.Extent, scriptFile.FilePath);
448+
foundDefinition = new SymbolReference(
449+
SymbolType.Function,
450+
foundSymbol.SymbolName,
451+
scriptFile.ScriptAst.Extent,
452+
scriptFile.FilePath);
465453
break;
466454
}
467455
}

0 commit comments

Comments
 (0)