Skip to content

Commit

Permalink
Merge pull request #64322 from CyrusNajmabadi/fawmnOverflow
Browse files Browse the repository at this point in the history
Change ForAttributeWithMetadataName to use an explicit stack when recursing
  • Loading branch information
CyrusNajmabadi committed Sep 29, 2022
2 parents 6452bca + 573df10 commit 0158ff5
Showing 1 changed file with 85 additions and 49 deletions.
Original file line number Diff line number Diff line change
Expand Up @@ -28,7 +28,8 @@ internal enum SourceGeneratorSyntaxTreeInfo

public partial struct SyntaxValueProvider
{
private static readonly ObjectPool<Stack<string>> s_stackPool = new ObjectPool<Stack<string>>(static () => new Stack<string>());
private static readonly ObjectPool<Stack<string>> s_stringStackPool = new ObjectPool<Stack<string>>(static () => new Stack<string>());
private static readonly ObjectPool<Stack<SyntaxNode>> s_nodeStackPool = new ObjectPool<Stack<SyntaxNode>>(static () => new Stack<SyntaxNode>());

/// <summary>
/// Returns all syntax nodes of that match <paramref name="predicate"/> if that node has an attribute on it that
Expand Down Expand Up @@ -170,87 +171,122 @@ public partial struct SyntaxValueProvider

// Used to ensure that as we recurse through alias names to see if they could bind to attributeName that we
// don't get into cycles.
var seenNames = s_stackPool.Allocate();
var seenNames = s_stringStackPool.Allocate();
var results = ArrayBuilder<SyntaxNode>.GetInstance();
var attributeTargets = ArrayBuilder<SyntaxNode>.GetInstance();

try
{
recurse(compilationUnit);
processCompilationUnit(compilationUnit);
}
finally
{
localAliases.Free();
seenNames.Clear();
s_stackPool.Free(seenNames);
s_stringStackPool.Free(seenNames);
attributeTargets.Free();
}

results.RemoveDuplicates();
return results.ToImmutableAndFree();

void recurse(SyntaxNode node)
void processCompilationUnit(SyntaxNode compilationUnit)
{
cancellationToken.ThrowIfCancellationRequested();

if (node is ICompilationUnitSyntax)
{
syntaxHelper.AddAliases(node.Green, localAliases, global: false);
if (compilationUnit is ICompilationUnitSyntax)
syntaxHelper.AddAliases(compilationUnit.Green, localAliases, global: false);

recurseChildren(node);
}
else if (syntaxHelper.IsAnyNamespaceBlock(node))
{
var localAliasCount = localAliases.Count;
syntaxHelper.AddAliases(node.Green, localAliases, global: false);
processCompilationOrNamespaceMembers(compilationUnit);
}

recurseChildren(node);
void processCompilationOrNamespaceMembers(SyntaxNode node)
{
cancellationToken.ThrowIfCancellationRequested();

// after recursing into this namespace, dump any local aliases we added from this namespace decl itself.
localAliases.Count = localAliasCount;
foreach (var child in node.ChildNodesAndTokens())
{
if (child.IsNode)
{
var childNode = child.AsNode()!;
if (syntaxHelper.IsAnyNamespaceBlock(childNode))
processNamespaceBlock(childNode);
else
processMember(childNode);
}
}
else if (syntaxHelper.IsAttributeList(node))
}

void processNamespaceBlock(SyntaxNode namespaceBlock)
{
cancellationToken.ThrowIfCancellationRequested();

var localAliasCount = localAliases.Count;
syntaxHelper.AddAliases(namespaceBlock.Green, localAliases, global: false);

processCompilationOrNamespaceMembers(namespaceBlock);

// after recursing into this namespace, dump any local aliases we added from this namespace decl itself.
localAliases.Count = localAliasCount;
}

void processMember(SyntaxNode member)
{
cancellationToken.ThrowIfCancellationRequested();

// nodes can be arbitrarily deep. Use an explicit stack over recursion to prevent a stack-overflow.
var nodeStack = s_nodeStackPool.Allocate();
nodeStack.Push(member);

try
{
foreach (var attribute in syntaxHelper.GetAttributesOfAttributeList(node))
while (nodeStack.Count > 0)
{
// Have to lookup both with the name in the attribute, as well as adding the 'Attribute' suffix.
// e.g. if there is [X] then we have to lookup with X and with XAttribute.
var simpleAttributeName = syntaxHelper.GetUnqualifiedIdentifierOfName(syntaxHelper.GetNameOfAttribute(attribute));
if (matchesAttributeName(simpleAttributeName, withAttributeSuffix: false) ||
matchesAttributeName(simpleAttributeName, withAttributeSuffix: true))
{
attributeTargets.Clear();
syntaxHelper.AddAttributeTargets(node, attributeTargets);
var node = nodeStack.Pop();

foreach (var target in attributeTargets)
if (syntaxHelper.IsAttributeList(node))
{
foreach (var attribute in syntaxHelper.GetAttributesOfAttributeList(node))
{
if (predicate(target, cancellationToken))
results.Add(target);
// Have to lookup both with the name in the attribute, as well as adding the 'Attribute' suffix.
// e.g. if there is [X] then we have to lookup with X and with XAttribute.
var simpleAttributeName = syntaxHelper.GetUnqualifiedIdentifierOfName(syntaxHelper.GetNameOfAttribute(attribute));
if (matchesAttributeName(simpleAttributeName, withAttributeSuffix: false) ||
matchesAttributeName(simpleAttributeName, withAttributeSuffix: true))
{
attributeTargets.Clear();
syntaxHelper.AddAttributeTargets(node, attributeTargets);

foreach (var target in attributeTargets)
{
if (predicate(target, cancellationToken))
results.Add(target);
}

break;
}
}

return;
// attributes can't have attributes inside of them. so no need to recurse when we're done.
}
else
{
// For any other node, just keep recursing deeper to see if we can find an attribute. Note: we cannot
// terminate the search anywhere as attributes may be found on things like local functions, and that
// means having to dive deep into statements and expressions.
foreach (var child in node.ChildNodesAndTokens().Reverse())
{
if (child.IsNode)
nodeStack.Push(child.AsNode()!);
}
}
}

// attributes can't have attributes inside of them. so no need to recurse when we're done.
}
}
else
finally
{
// For any other node, just keep recursing deeper to see if we can find an attribute. Note: we cannot
// terminate the search anywhere as attributes may be found on things like local functions, and that
// means having to dive deep into statements and expressions.
recurseChildren(node);
}

return;

void recurseChildren(SyntaxNode node)
{
foreach (var child in node.ChildNodesAndTokens())
{
if (child.IsNode)
recurse(child.AsNode()!);
}
nodeStack.Clear();
s_nodeStackPool.Free(nodeStack);
}
}

Expand Down

0 comments on commit 0158ff5

Please sign in to comment.