[clang] Make sure argument expansion locations are correct in presence of predefined buffer

[kadircet]

Summary

Macro argument expansion logic relies on skipping file IDs that created
as a result of an include. Unfortunately it fails to do that for
predefined buffer since it doesn't have a valid insertion location.

As a result of that any file ID created for an include inside the
predefined buffers breaks the traversal logic in
SourceManager::computeMacroArgsCache.

To fix this issue we first record number of created FIDs for predefined
buffer, and then skip them explicitly in source manager.

Another solution would be to just give predefined buffers a valid source
location, but it is unclear where that should be..

[sammccall]

clang/lib/Basic/SourceManager.cpp

1805

nit: reverse the order of this check to avoid the string compare?
You could also consider hoisting bool IsMainFile out of the loop, maybe the optimizer can see this but I think it's probably a readablity win anyway.

1806

I don't love having this code duplicated.

What about:

SourceLocation IncludeLoc = ...
bool IncludedFromMain = isInFileID(IncludeLoc, FID) || (IsMainFile && Filename = "<built-in>");
if (IncludedFromMain) {
  // increment ID
} else if (IncludeLoc.isValid()) { // included, but not from this file
  return
}
continue;

clang/lib/Lex/PPLexerChange.cpp

420

Can we ever get spurious equality here because both sides are 0?

[kadircet]

clang/lib/Lex/PPLexerChange.cpp

420

the latter is set during Preprocessor::EnterMainSourceFile() which initializes the PP and all users seem to be calling it first, therefore I assumed it would alwyas be a valid file id.

putting a && getPredefinesFileID().isValid() to be on the safe side.

[kadircet]

• address comments

[sammccall]

clang/lib/Basic/SourceManager.cpp

1816

wasn't included from FID -> was included but not from FID

clang/lib/Basic/SourceManager.cpp

while (true) {
     }
 
     bool Invalid = false;
     const SrcMgr::SLocEntry &Entry = getSLocEntryByID(ID, &Invalid);
     if (Invalid)
       return;
     if (Entry.isFile()) {
       SourceLocation IncludeLoc = Entry.getFile().getIncludeLoc();
-      if (IncludeLoc.isInvalid())
-        continue;
-      if (!isInFileID(IncludeLoc, FID))
-        return; // No more files/macros that may be "contained" in this file.
-
-      // Skip the files/macros of the #include'd file, we only care about macros
-      // that lexed macro arguments from our file.
+      bool IncludedInFID =
+          (IncludeLoc.isValid() && isInFileID(IncludeLoc, FID)) ||
+          // Predefined header doesn't have a valid include location in main
+          // file, but any files created by it should still be skipped when
+          // computing macro args expanded in the main file.
+          (FID == MainFileID && Entry.getFile().Filename == "<built-in>");
+      if (IncludedInFID) {
+        // Skip the files/macros of the #include'd file, we only care about
+        // macros that lexed macro arguments from our file.
         if (Entry.getFile().NumCreatedFIDs)
-        ID += Entry.getFile().NumCreatedFIDs - 1/*because of next ++ID*/;
+          ID += Entry.getFile().NumCreatedFIDs - 1 /*because of next ++ID*/;
         continue;
+      } else if (IncludeLoc.isValid()) {
+        // If file was included but not from FID, there is no more files/macros
+        // that may be "contained" in this file.
+        return;
+      }
+      continue;
     }
 
     const ExpansionInfo &ExpInfo = Entry.getExpansion();
 
     if (ExpInfo.getExpansionLocStart().isFileID()) {
       if (!isInFileID(ExpInfo.getExpansionLocStart(), FID))
         return; // No more files/macros that may be "contained" in this file.
     }

clang/lib/Lex/PPLexerChange.cpp

if (isCodeCompletionEnabled() && CurPPLexer &&
       CurLexer.reset();
 
       CurPPLexer = nullptr;
       recomputeCurLexerKind();
       return true;
     }
 
     if (!isEndOfMacro && CurPPLexer &&
-        SourceMgr.getIncludeLoc(CurPPLexer->getFileID()).isValid()) {
+        (SourceMgr.getIncludeLoc(CurPPLexer->getFileID()).isValid() ||
+         // Predefines file doesn't have a valid include location.
+         (PredefinesFileID.isValid() &&
+          CurPPLexer->getFileID() == PredefinesFileID))) {
       // Notify SourceManager to record the number of FileIDs that were created
       // during lexing of the #include'd file.
       unsigned NumFIDs =
           SourceMgr.local_sloc_entry_size() -
           CurPPLexer->getInitialNumSLocEntries() + 1/*#include'd file*/;
       SourceMgr.setNumCreatedFIDsForFileID(CurPPLexer->getFileID(), NumFIDs);
     }
 

clang/unittests/Basic/SourceManagerTest.cpp

TEST_F(SourceManagerTest, getMacroArgExpandedLocation) {
 
   const FileEntry *headerFile = FileMgr.getVirtualFile("/test-header.h",
                                                  HeaderBuf->getBufferSize(), 0);
   SourceMgr.overrideFileContents(headerFile, std::move(HeaderBuf));
 
   TrivialModuleLoader ModLoader;
   HeaderSearch HeaderInfo(std::make_shared<HeaderSearchOptions>(), SourceMgr,
                           Diags, LangOpts, &*Target);
+
   Preprocessor PP(std::make_shared<PreprocessorOptions>(), Diags, LangOpts,
                   SourceMgr, HeaderInfo, ModLoader,
                   /*IILookup =*/nullptr,
                   /*OwnsHeaderSearch =*/false);
+  // Ensure we can get expanded locations in presence of implicit includes.
+  // These are different than normal includes since predefines buffer doesn't
+  // have a valid insertion location.
+  PP.setPredefines("#include \"/implicit-header.h\"");
+  FileMgr.getVirtualFile("/implicit-header.h", 0, 0);
   PP.Initialize(*Target);
   PP.EnterMainSourceFile();
 
   std::vector<Token> toks;
   while (1) {
     Token tok;
     PP.Lex(tok);
     if (tok.is(tok::eof))

clang/unittests/Lex/LexerTest.cpp

while (true) {
     if (T->is(tok::eof))
       break;
     GeneratedByNextToken.push_back(getSourceText(*T, *T));
     Loc = T->getLocation();
   }
   EXPECT_THAT(GeneratedByNextToken, ElementsAre("abcd", "=", "0", ";", "int",
                                                 "xyz", "=", "abcd", ";"));
 }
+
+TEST_F(LexerTest, CreatedFIDCountForPredefinedBuffer) {
+  TrivialModuleLoader ModLoader;
+  auto PP = CreatePP("", ModLoader);
+  while (1) {
+    Token tok;
+    PP->Lex(tok);
+    if (tok.is(tok::eof))
+      break;
+  }
+  EXPECT_EQ(SourceMgr.getNumCreatedFIDsForFileID(PP->getPredefinesFileID()),
+            1U);
+}
 } // anonymous namespace