Skip to content

Commit c060a90

Browse files
DougGregorDougGregor
committed
Generalize the LCA implementation in ASTScope for arbitrary macro nesting
The prior Least Common Ancestor (LCA) implementation in ASTScope, which is used to search child scopes to find a particular location, assumed that the source range it was given was contained within a single buffer and was specific to the child-scope search task. Generalize this to a more fundamental "is before" operation on source locations that respects macro expansions, simplifying the code in the process.
1 parent e32289b commit c060a90

File tree

2 files changed

+103
-78
lines changed

2 files changed

+103
-78
lines changed

‎lib/AST/ASTScopeLookup.cpp‎

Lines changed: 94 additions & 78 deletions
Original file line numberDiff line numberDiff line change
@@ -47,7 +47,26 @@ void ASTScopeImpl::unqualifiedLookup(
4747

4848
const ASTScopeImpl *ASTScopeImpl::findStartingScopeForLookup(
4949
SourceFile *sourceFile, const SourceLoc loc) {
50-
auto *const fileScope = sourceFile->getScope().impl;
50+
auto *fileScope = sourceFile->getScope().impl;
51+
52+
// Workaround for bad locations; just return the file scope.
53+
if (loc.isInvalid())
54+
return fileScope;
55+
56+
// Some callers get the actual source file wrong. Look for the actual
57+
// source file containing this location.
58+
auto actualSF =
59+
sourceFile->getParentModule()->getSourceFileContainingLocation(loc);
60+
61+
// If there is no source file containing this location, just return the
62+
// scope we have.
63+
if (!actualSF)
64+
return fileScope;
65+
66+
// Grab the new file scope.
67+
if (actualSF != sourceFile)
68+
fileScope = actualSF->getScope().impl;
69+
5170
const auto *innermost = fileScope->findInnermostEnclosingScope(
5271
sourceFile->getParentModule(), loc, nullptr);
5372
ASTScopeAssert(innermost->getWasExpanded(),
@@ -79,104 +98,101 @@ ASTScopeImpl *ASTScopeImpl::findInnermostEnclosingScopeImpl(
7998
/// If the \p loc is in a new buffer but \p range is not, consider the location
8099
/// is at the start of replaced range. Otherwise, returns \p loc as is.
81100
static SourceLoc translateLocForReplacedRange(SourceManager &sourceMgr,
82-
CharSourceRange range,
101+
SourceLoc rangeStart,
83102
SourceLoc loc) {
84103
for (const auto &pair : sourceMgr.getReplacedRanges()) {
85104
if (sourceMgr.rangeContainsTokenLoc(pair.second, loc) &&
86-
!sourceMgr.rangeContainsTokenLoc(pair.second, range.getStart())) {
105+
!sourceMgr.rangeContainsTokenLoc(pair.second, rangeStart)) {
87106
return pair.first.Start;
88107
}
89108
}
90109
return loc;
91110
}
92111

112+
/// Populate the ancestors list for this buffer, with the root source buffer
113+
/// at the beginning and the given source buffer at the end.
114+
static void populateAncestors(
115+
SourceManager &sourceMgr, unsigned bufferID,
116+
SmallVectorImpl<unsigned> &ancestors) {
117+
if (auto info = sourceMgr.getGeneratedSourceInfo(bufferID)) {
118+
auto ancestorLoc = info->originalSourceRange.getStart();
119+
if (ancestorLoc.isValid()) {
120+
auto ancestorBufferID = sourceMgr.findBufferContainingLoc(ancestorLoc);
121+
populateAncestors(sourceMgr, ancestorBufferID, ancestors);
122+
}
123+
}
124+
125+
ancestors.push_back(bufferID);
126+
}
127+
128+
/// Determine whether the first source location precedes the second, accounting
129+
/// for macro expansions.
130+
static bool isBeforeInSource(
131+
SourceManager &sourceMgr, SourceLoc firstLoc, SourceLoc secondLoc,
132+
bool allowEqual) {
133+
// If the two locations are in the same source buffer, compare their pointers.
134+
unsigned firstBufferID = sourceMgr.findBufferContainingLoc(firstLoc);
135+
unsigned secondBufferID = sourceMgr.findBufferContainingLoc(secondLoc);
136+
if (firstBufferID == secondBufferID) {
137+
return sourceMgr.isBeforeInBuffer(firstLoc, secondLoc) ||
138+
(allowEqual && firstLoc == secondLoc);
139+
}
140+
141+
// If the two locations are in different source buffers, we need to compute
142+
// the least common ancestor.
143+
SmallVector<unsigned, 4> firstAncestors;
144+
populateAncestors(sourceMgr, firstBufferID, firstAncestors);
145+
SmallVector<unsigned, 4> secondAncestors;
146+
populateAncestors(sourceMgr, secondBufferID, secondAncestors);
147+
148+
// Find the first mismatch between the two ancestor lists; this is the
149+
// point of divergence.
150+
auto [firstMismatch, secondMismatch] = std::mismatch(
151+
firstAncestors.begin(), firstAncestors.end(),
152+
secondAncestors.begin(), secondAncestors.end());
153+
assert(firstMismatch != firstAncestors.begin() &&
154+
secondMismatch != secondAncestors.begin() &&
155+
"Ancestors don't have the same root source file");
156+
157+
SourceLoc firstLocInLCA = firstMismatch == firstAncestors.end()
158+
? firstLoc
159+
: sourceMgr.getGeneratedSourceInfo(*firstMismatch)
160+
->originalSourceRange.getStart();
161+
SourceLoc secondLocInLCA = secondMismatch == secondAncestors.end()
162+
? secondLoc
163+
: sourceMgr.getGeneratedSourceInfo(*secondMismatch)
164+
->originalSourceRange.getStart();
165+
return sourceMgr.isBeforeInBuffer(firstLocInLCA, secondLocInLCA) ||
166+
(allowEqual && firstLocInLCA == secondLocInLCA);
167+
}
168+
93169
NullablePtr<ASTScopeImpl>
94170
ASTScopeImpl::findChildContaining(ModuleDecl *parentModule,
95171
SourceLoc loc,
96172
SourceManager &sourceMgr) const {
97-
auto *locSourceFile = parentModule->getSourceFileContainingLocation(loc);
173+
if (loc.isInvalid())
174+
return nullptr;
98175

99176
// Use binary search to find the child that contains this location.
100177
auto *const *child = llvm::lower_bound(
101178
getChildren(), loc,
102179
[&](const ASTScopeImpl *scope, SourceLoc loc) {
103-
auto rangeOfScope = scope->getCharSourceRangeOfScope(sourceMgr);
104-
ASTScopeAssert(!sourceMgr.isBeforeInBuffer(rangeOfScope.getEnd(),
105-
rangeOfScope.getStart()),
106-
"Source range is backwards");
107-
108-
// If the scope source range and the loc are in two different source
109-
// files, one or both of them are in a macro expansion buffer.
110-
111-
// Note that `scope->getSourceFile()` returns the root of the source tree,
112-
// not the source file containing the location of the ASTScope.
113-
auto scopeStart = scope->getSourceRangeOfThisASTNode().Start;
114-
auto *scopeSourceFile = parentModule->getSourceFileContainingLocation(scopeStart);
115-
116-
if (scopeSourceFile != locSourceFile) {
117-
// To compare a source location that is possibly inside a macro expansion
118-
// with a source range that is also possibly in a macro expansion (not
119-
// necessarily the same one as before) we need to find the LCA in the
120-
// source file tree of macro expansions, and compare the original source
121-
// ranges within that common ancestor. We can't walk all the way up to the
122-
// source file containing the parent scope we're searching the children of,
123-
// because two independent (possibly nested) macro expansions can have the
124-
// same original source range in that file; freestanding and peer macros
125-
// mean that we can have arbitrarily nested macro expansions that all add
126-
// declarations to the same scope, that all originate from a single macro
127-
// invocation in the original source file.
128-
129-
// A map from enclosing source files to original source ranges of the macro
130-
// expansions within that file, recording the chain of macro expansions for
131-
// the given scope.
132-
llvm::SmallDenseMap<const SourceFile *, CharSourceRange>
133-
scopeExpansions;
134-
135-
// Walk up the chain of macro expansion buffers for the scope, recording the
136-
// original source range of the macro expansion along the way using generated
137-
// source info.
138-
auto *scopeExpansion = scopeSourceFile;
139-
scopeExpansions[scopeExpansion] =
140-
Lexer::getCharSourceRangeFromSourceRange(
141-
sourceMgr, scope->getSourceRangeOfThisASTNode());
142-
while (auto *ancestor = scopeExpansion->getEnclosingSourceFile()) {
143-
auto generatedInfo =
144-
sourceMgr.getGeneratedSourceInfo(*scopeExpansion->getBufferID());
145-
scopeExpansions[ancestor] = generatedInfo->originalSourceRange;
146-
scopeExpansion = ancestor;
147-
}
148-
149-
// Walk up the chain of macro expansion buffers for the source loc we're
150-
// searching for to find the LCA using `scopeExpansions`.
151-
auto *potentialLCA = locSourceFile;
152-
auto expansionLoc = loc;
153-
while (potentialLCA) {
154-
auto scopeExpansion = scopeExpansions.find(potentialLCA);
155-
if (scopeExpansion != scopeExpansions.end()) {
156-
// Take the original expansion range within the LCA of the loc and
157-
// the scope to compare.
158-
rangeOfScope = scopeExpansion->second;
159-
loc = expansionLoc;
160-
break;
161-
}
162-
163-
auto generatedInfo =
164-
sourceMgr.getGeneratedSourceInfo(*potentialLCA->getBufferID());
165-
if (generatedInfo)
166-
expansionLoc = generatedInfo->originalSourceRange.getStart();
167-
potentialLCA = potentialLCA->getEnclosingSourceFile();
168-
}
169-
}
170-
171-
loc = translateLocForReplacedRange(sourceMgr, rangeOfScope, loc);
172-
return (rangeOfScope.getEnd() == loc ||
173-
sourceMgr.isBeforeInBuffer(rangeOfScope.getEnd(), loc));
180+
auto rangeOfScope = scope->getSourceRangeOfThisASTNode();
181+
auto endOfScope =
182+
Lexer::getLocForEndOfToken(sourceMgr, rangeOfScope.End);
183+
184+
loc = translateLocForReplacedRange(sourceMgr, rangeOfScope.Start, loc);
185+
return isBeforeInSource(
186+
sourceMgr, endOfScope, loc, /*allowEqual=*/true);
174187
});
175188

176189
if (child != getChildren().end()) {
177-
auto rangeOfScope = (*child)->getCharSourceRangeOfScope(sourceMgr);
178-
loc = translateLocForReplacedRange(sourceMgr, rangeOfScope, loc);
179-
if (rangeOfScope.contains(loc))
190+
auto rangeOfScope = (*child)->getSourceRangeOfThisASTNode();
191+
auto endOfScope = Lexer::getLocForEndOfToken(sourceMgr, rangeOfScope.End);
192+
loc = translateLocForReplacedRange(sourceMgr, rangeOfScope.Start, loc);
193+
if (isBeforeInSource(sourceMgr, rangeOfScope.Start, loc,
194+
/*allowEqual=*/true) &&
195+
isBeforeInSource(sourceMgr, loc, endOfScope, /*allowEqual=*/false))
180196
return *child;
181197
}
182198

‎lib/AST/NameLookup.cpp‎

Lines changed: 9 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1665,6 +1665,15 @@ bool
16651665
namelookup::isInMacroArgument(SourceFile *sourceFile, SourceLoc loc) {
16661666
bool inMacroArgument = false;
16671667

1668+
// Make sure that the source location is actually within the given source
1669+
// file.
1670+
if (sourceFile && loc.isValid()) {
1671+
sourceFile =
1672+
sourceFile->getParentModule()->getSourceFileContainingLocation(loc);
1673+
if (!sourceFile)
1674+
return false;
1675+
}
1676+
16681677
ASTScope::lookupEnclosingMacroScope(
16691678
sourceFile, loc,
16701679
[&](auto potentialMacro) -> bool {

0 commit comments

Comments
 (0)