@@ -213,19 +213,109 @@ TokenBuffer::spelledForExpandedToken(const syntax::Token *Expanded) const {
213213 // Our token could only be produced by the previous mapping.
214214 if (It == File.Mappings .begin ()) {
215215 // No previous mapping, no need to modify offsets.
216- return {&File.SpelledTokens [ExpandedIndex - File.BeginExpanded ], nullptr };
216+ return {&File.SpelledTokens [ExpandedIndex - File.BeginExpanded ],
217+ /* Mapping=*/ nullptr };
217218 }
218219 --It; // 'It' now points to last mapping that started before our token.
219220
220221 // Check if the token is part of the mapping.
221222 if (ExpandedIndex < It->EndExpanded )
222- return {&File.SpelledTokens [It->BeginSpelled ], /* Mapping*/ &*It};
223+ return {&File.SpelledTokens [It->BeginSpelled ], /* Mapping= */
223224
224225 // Not part of the mapping, use the index from previous mapping to compute the
225226 // corresponding spelled token.
226227 return {
227228 &File.SpelledTokens [It->EndSpelled + (ExpandedIndex - It->EndExpanded )],
228- /* Mapping*/ nullptr };
229+ /* Mapping=*/ nullptr };
230+ }
231+
232+ const TokenBuffer::Mapping *
233+ TokenBuffer::mappingStartingBeforeSpelled (const MarkedFile &F,
234+ const syntax::Token *Spelled) {
235+ assert (F.SpelledTokens .data () <= Spelled);
236+ unsigned SpelledI = Spelled - F.SpelledTokens .data ();
237+ assert (SpelledI < F.SpelledTokens .size ());
238+
239+ auto It = llvm::partition_point (F.Mappings , [SpelledI](const Mapping &M) {
240+ return M.BeginSpelled <= SpelledI;
241+ });
242+ if (It == F.Mappings .begin ())
243+ return nullptr ;
244+ --It;
245+ return &*It;
246+ }
247+
248+ llvm::SmallVector<llvm::ArrayRef<syntax::Token>, 1 >
249+ TokenBuffer::expandedForSpelled (llvm::ArrayRef<syntax::Token> Spelled) const {
250+ if (Spelled.empty ())
251+ return {};
252+ assert (Spelled.front ().location ().isFileID ());
253+
254+ auto FID = sourceManager ().getFileID (Spelled.front ().location ());
255+ auto It = Files.find (FID);
256+ assert (It != Files.end ());
257+
258+ const MarkedFile &File = It->second ;
259+ // `Spelled` must be a subrange of `File.SpelledTokens`.
260+ assert (File.SpelledTokens .data () <= Spelled.data ());
261+ assert (&Spelled.back () <=
262+ File.SpelledTokens .data () + File.SpelledTokens .size ());
263+ #ifndef NDEBUG
264+ auto T1 = Spelled.back ().location ();
265+ auto T2 = File.SpelledTokens .back ().location ();
266+ assert (T1 == T2 || sourceManager ().isBeforeInTranslationUnit (T1, T2));
267+ #endif
268+
269+ auto *FrontMapping = mappingStartingBeforeSpelled (File, &Spelled.front ());
270+ unsigned SpelledFrontI = &Spelled.front () - File.SpelledTokens .data ();
271+ assert (SpelledFrontI < File.SpelledTokens .size ());
272+ unsigned ExpandedBegin;
273+ if (!FrontMapping) {
274+ // No mapping that starts before the first token of Spelled, we don't have
275+ // to modify offsets.
276+ ExpandedBegin = File.BeginExpanded + SpelledFrontI;
277+ } else if (SpelledFrontI < FrontMapping->EndSpelled ) {
278+ // This mapping applies to Spelled tokens.
279+ if (SpelledFrontI != FrontMapping->BeginSpelled ) {
280+ // Spelled tokens don't cover the entire mapping, returning empty result.
281+ return {}; // FIXME: support macro arguments.
282+ }
283+ // Spelled tokens start at the beginning of this mapping.
284+ ExpandedBegin = FrontMapping->BeginExpanded ;
285+ } else {
286+ // Spelled tokens start after the mapping ends (they start in the hole
287+ // between 2 mappings, or between a mapping and end of the file).
288+ ExpandedBegin =
289+ FrontMapping->EndExpanded + (SpelledFrontI - FrontMapping->EndSpelled );
290+ }
291+
292+ auto *BackMapping = mappingStartingBeforeSpelled (File, &Spelled.back ());
293+ unsigned SpelledBackI = &Spelled.back () - File.SpelledTokens .data ();
294+ unsigned ExpandedEnd;
295+ if (!BackMapping) {
296+ // No mapping that starts before the last token of Spelled, we don't have to
297+ // modify offsets.
298+ ExpandedEnd = File.BeginExpanded + SpelledBackI + 1 ;
299+ } else if (SpelledBackI < BackMapping->EndSpelled ) {
300+ // This mapping applies to Spelled tokens.
301+ if (SpelledBackI + 1 != BackMapping->EndSpelled ) {
302+ // Spelled tokens don't cover the entire mapping, returning empty result.
303+ return {}; // FIXME: support macro arguments.
304+ }
305+ ExpandedEnd = BackMapping->EndExpanded ;
306+ } else {
307+ // Spelled tokens end after the mapping ends.
308+ ExpandedEnd =
309+ BackMapping->EndExpanded + (SpelledBackI - BackMapping->EndSpelled ) + 1 ;
310+ }
311+
312+ assert (ExpandedBegin < ExpandedTokens.size ());
313+ assert (ExpandedEnd < ExpandedTokens.size ());
314+ // Avoid returning empty ranges.
315+ if (ExpandedBegin == ExpandedEnd)
316+ return {};
317+ return {llvm::makeArrayRef (ExpandedTokens.data () + ExpandedBegin,
318+ ExpandedTokens.data () + ExpandedEnd)};
229319}
230320
231321llvm::ArrayRef<syntax::Token> TokenBuffer::spelledTokens (FileID FID) const {
@@ -436,14 +526,38 @@ class TokenCollector::CollectPPExpansions : public PPCallbacks {
436526 SourceRange Range, const MacroArgs *Args) override {
437527 if (!Collector)
438528 return ;
439- // Only record top-level expansions, not those where:
529+ const auto &SM = Collector->PP .getSourceManager ();
530+ // Only record top-level expansions that directly produce expanded tokens.
531+ // This excludes those where:
440532 // - the macro use is inside a macro body,
441533 // - the macro appears in an argument to another macro.
442- if (!MacroNameTok.getLocation ().isFileID () ||
443- (LastExpansionEnd.isValid () &&
444- Collector->PP .getSourceManager ().isBeforeInTranslationUnit (
445- Range.getBegin (), LastExpansionEnd)))
534+ // However macro expansion isn't really a tree, it's token rewrite rules,
535+ // so there are other cases, e.g.
536+ // #define B(X) X
537+ // #define A 1 + B
538+ // A(2)
539+ // Both A and B produce expanded tokens, though the macro name 'B' comes
540+ // from an expansion. The best we can do is merge the mappings for both.
541+
542+ // The *last* token of any top-level macro expansion must be in a file.
543+ // (In the example above, see the closing paren of the expansion of B).
544+ if (!Range.getEnd ().isFileID ())
545+ return ;
546+ // If there's a current expansion that encloses this one, this one can't be
547+ // top-level.
548+ if (LastExpansionEnd.isValid () &&
549+ !SM.isBeforeInTranslationUnit (LastExpansionEnd, Range.getEnd ()))
446550 return ;
551+
552+ // If the macro invocation (B) starts in a macro (A) but ends in a file,
553+ // we'll create a merged mapping for A + B by overwriting the endpoint for
554+ // A's startpoint.
555+ if (!Range.getBegin ().isFileID ()) {
556+ Range.setBegin (SM.getExpansionLoc (Range.getBegin ()));
557+ assert (Collector->Expansions .count (Range.getBegin ().getRawEncoding ()) &&
558+ " Overlapping macros should have same expansion location"
559+ }
560+
447561 Collector->Expansions [Range.getBegin ().getRawEncoding ()] = Range.getEnd ();
448562 LastExpansionEnd = Range.getEnd ();
449563 }
@@ -526,6 +640,20 @@ class TokenCollector::Builder {
526640 for (const auto &File : Result.Files )
527641 discard (File.first );
528642
643+ #ifndef NDEBUG
644+ for (auto &pair : Result.Files ) {
645+ auto &mappings = pair.second .Mappings ;
646+ assert (std::is_sorted (
647+ mappings.begin (), mappings.end (),
648+ [](const TokenBuffer::Mapping &M1, const TokenBuffer::Mapping &M2) {
649+ return M1.BeginSpelled < M2.BeginSpelled &&
650+ M1.EndSpelled < M2.EndSpelled &&
651+ M1.BeginExpanded < M2.BeginExpanded &&
652+ M1.EndExpanded < M2.EndExpanded ;
653+ }));
654+ }
655+ #endif
656+
529657 return std::move (Result);
530658 }
531659
0 commit comments