[SourceManager] Extract 'findFileIDsForFile' from 'translateFile' to allow mapping from one file entry to multiple FileIDs

[arphaman]

Summary

This patch extracts the code that searches for a file id from translateFile to findFileIDsForFile to allow the users to map from one file entry to multiple FileIDs.
Will be used as a basis for a clang-rename fix in https://reviews.llvm.org/D50801.

[ioeric]

include/clang/Basic/SourceManager.h

1551

nit: put this closer to the closely related translateFile?

1557

Callback pattern seems uncommon in LLVM/Clang. I'd suggest making this return a set of FileIDs and put the callback-based function as a helper (shared by this and translateFile)in the implementation.

[arphaman]

Address review comments.

include/clang/Basic/SourceManager.h

1557

I created a helper, but unfortunately it needs to be a member of SourceManager as FileID::get is private.

[ioeric]

lib/Basic/SourceManager.cpp

1705

Do we also need this in findFileIDsForFile?

unittests/Basic/SourceManagerTest.cpp

380

Could you add a test case for getting file ID for main file, just to make sure we also covered cases handled by if (MainFileID.isValid()) {...} code path in translateFile()?

[arphaman]

Address review comments

lib/Basic/SourceManager.cpp

1705

I don't really need this for my use case.

[ioeric]

lib/Basic/SourceManager.cpp

1705

But it's not clear from the interface AFAICT. We should either handle this case (maybe controlled with a flag), or make it clear in the API (with a different name or documentation).

[arphaman]

Address review comments.

[arphaman]

lib/Basic/SourceManager.cpp

1705

Hmm, I think that would be better. I pulled that code into the findFileIDsForFile helper function, so we can call it in two modes now. It's probably good to do that check just in case in the new API too, so it does that check as well.

[ioeric]

lib/Basic/SourceManager.cpp

1626

Should we check FileID::get(I) is valid?

1628

The conditions here are pretty hard to follow and reason about. Could we maybe split them (some documentation would also help)?

1628

In the original version, file system updates are checked last (after modules). Any reason to move it before modules?

Also, it seems that this code path could also be run when FileIDabove is invalid? So I wonder whether else if is correct here.

[arphaman]

Remove dead code for filesystem update fileID matching.

lib/Basic/SourceManager.cpp

1626

That's not really necessary. The FileID we get should be valid as a local SLoc entry should have a corresponding FileID. The SLoc 0 is not a File loc entry so we can never get FileID::get(0) here.

1628

I just realized that the original file system code has never been executed and was just dead code! If we take a look at this logic:

for (unsigned I = 0, N = local_sloc_entry_size(); I != N; ++I) {
    FileID IFileID;
    IFileID.ID = I;
    const SLocEntry &SLoc = getSLocEntry(IFileID, &Invalid);
    if (Invalid)
      return false;

You'll notice that the loop starts iterating at 0. Get SLocEntry is called with FileID 0, which sets the Invalid flag to true. Then we simply return, so the loop never reached the code below. Looks like it's a regression that happened years ago. I removed this code for now, but I'll reinstate it correctly in a follow-up patch.

[jkorous]

LGTM and it seems like all of Eric's comments were answered too.

include/clang/Basic/SourceManager.h

SourceLocation translateFileLineCol(const FileEntry *SourceFile,
                                       unsigned Line, unsigned Col) const;
 
   /// Get the FileID for the given file.
   ///
   /// If the source file is included multiple times, the FileID will be the
   /// first inclusion.
   FileID translateFile(const FileEntry *SourceFile) const;
 
+  /// Looks through all the local and imported source locations to find the set
+  /// of FileIDs that correspond to the given entry.
+  ///
+  /// \returns the set of FileIDs for the given file, in the order of inclusion.
+  llvm::SmallVector<FileID, 2>
+  findFileIDsForFile(const FileEntry *SourceFile) const;
+
+private:
+  /// Looks through all the local and imported source locations to find the set
+  /// of FileIDs that correspond to the given entry.
+  ///
+  /// \param Callback should return true to exit the search.
+  ///
+  /// \returns true if the callback returned true, false otherwise.
+  bool findFileIDsForFile(const FileEntry *SourceFile,
+                          llvm::function_ref<bool(FileID)> Callback) const;
+
+public:
   /// Get the source location in \p FID for the given line:col.
   /// Returns null location if \p FID is not a file SLocEntry.
   SourceLocation translateLineCol(FileID FID,
                                   unsigned Line, unsigned Col) const;
 
   /// If \p Loc points inside a function macro argument, the returned
   /// location will be the macro location in which the argument was expanded.
   /// If a macro argument is used multiple times, the expanded location will
   /// be at the first expansion of the argument.
   /// e.g.
   ///   MY_MACRO(foo);
   ///             ^
   /// Passing a file location pointing at 'foo', will yield a macro location
   /// where 'foo' was expanded into.
   SourceLocation getMacroArgExpandedLocation(SourceLocation Loc) const;
 
   /// Determines the order of 2 source locations in the translation unit.
   ///
   /// \returns true if LHS source location comes before RHS, false otherwise.
   bool isBeforeInTranslationUnit(SourceLocation LHS, SourceLocation RHS) const;

lib/Basic/SourceManager.cpp

SourceLocation SourceManager::translateFileLineCol(const FileEntry *SourceFile,
                                                   unsigned Col) const {
   assert(SourceFile && "Null source file!");
   assert(Line && Col && "Line and column should start from 1!");
 
   FileID FirstFID = translateFile(SourceFile);
   return translateLineCol(FirstFID, Line, Col);
 }
 
+bool SourceManager::findFileIDsForFile(
+    const FileEntry *SourceFile,
+    llvm::function_ref<bool(FileID)> Callback) const {
+  assert(SourceFile && "Null source file!");
+
+  // Look through all of the local source locations.
+  for (unsigned I = 0, N = local_sloc_entry_size(); I != N; ++I) {
+    bool Invalid = false;
+    const SLocEntry &SLoc = getLocalSLocEntry(I, &Invalid);
+    if (Invalid)
+      return false;
+
+    if (SLoc.isFile() && SLoc.getFile().getContentCache() &&
+        SLoc.getFile().getContentCache()->OrigEntry == SourceFile) {
+      if (Callback(FileID::get(I)))
+        return true;
+    }
+  }
+
+  // If that still didn't help, try the modules.
+  for (unsigned I = 0, N = loaded_sloc_entry_size(); I != N; ++I) {
+    const SLocEntry &SLoc = getLoadedSLocEntry(I);
+    if (SLoc.isFile() && SLoc.getFile().getContentCache() &&
+        SLoc.getFile().getContentCache()->OrigEntry == SourceFile) {
+      if (Callback(FileID::get(-int(I) - 2)))
+        return true;
+    }
+  }
+  return false;
+}
+
+llvm::SmallVector<FileID, 2>
+SourceManager::findFileIDsForFile(const FileEntry *SourceFile) const {
+  llvm::SmallVector<FileID, 2> Result;
+  findFileIDsForFile(SourceFile, [&](FileID F) {
+    Result.push_back(F);
+    return false;
+  });
+  return Result;
+}
+
 /// Get the FileID for the given file.
 ///
 /// If the source file is included multiple times, the FileID will be the
 /// first inclusion.
 FileID SourceManager::translateFile(const FileEntry *SourceFile) const {
   assert(SourceFile && "Null source file!");
 
   // Find the first file ID that corresponds to the given file.
if (MainSLoc.isFile()) {
               }
             }
           }
         }
       }
     }
   }
 
-  if (FirstFID.isInvalid()) {
-    // The location we're looking for isn't in the main file; look
-    // through all of the local source locations.
-    for (unsigned I = 0, N = local_sloc_entry_size(); I != N; ++I) {
-      bool Invalid = false;
-      const SLocEntry &SLoc = getLocalSLocEntry(I, &Invalid);
-      if (Invalid)
-        return FileID();
+  if (FirstFID.isValid())
+    return FirstFID;
 
-      if (SLoc.isFile() &&
-          SLoc.getFile().getContentCache() &&
-          SLoc.getFile().getContentCache()->OrigEntry == SourceFile) {
-        FirstFID = FileID::get(I);
-        break;
-      }
-    }
-    // If that still didn't help, try the modules.
-    if (FirstFID.isInvalid()) {
-      for (unsigned I = 0, N = loaded_sloc_entry_size(); I != N; ++I) {
-        const SLocEntry &SLoc = getLoadedSLocEntry(I);
-        if (SLoc.isFile() &&
-            SLoc.getFile().getContentCache() &&
-            SLoc.getFile().getContentCache()->OrigEntry == SourceFile) {
-          FirstFID = FileID::get(-int(I) - 2);
-          break;
-        }
-      }
-    }
-  }
+  // The location we're looking for isn't in the main file; look
+  // through all of the local and the imported source locations.
+  if (findFileIDsForFile(SourceFile, [&](FileID F) {
+        FirstFID = F;
+        return true;
+      }))
+    return FirstFID;
 
   // If we haven't found what we want yet, try again, but this time stat()
   // each of the files in case the files have changed since we originally
   // parsed the file.
   if (FirstFID.isInvalid() &&
       (SourceFileName ||
        (SourceFileName = llvm::sys::path::filename(SourceFile->getName()))) &&
       (SourceFileUID || (SourceFileUID = getActualFileUID(SourceFile)))) {
     bool Invalid = false;
     for (unsigned I = 0, N = local_sloc_entry_size(); I != N; ++I) {

unittests/Basic/SourceManagerTest.cpp

TEST_F(SourceManagerTest, isBeforeInTranslationUnitWithMacroInInclude) {
   // MACRO_IN_INCLUDE definition of the included file.
   EXPECT_TRUE(SourceMgr.isBeforeInTranslationUnit(Macros[3].Loc, Macros[4].Loc));
 
   // The INC2 expansion in #include M(INC2) comes before the second
   // MACRO_IN_INCLUDE definition of the included file.
   EXPECT_TRUE(SourceMgr.isBeforeInTranslationUnit(Macros[10].Loc, Macros[11].Loc));
 }
 
+TEST_F(SourceManagerTest, findFileIDsForFile) {
+  const char *header = "int x;";
+
+  const char *main = "#include </test-header.h>\n"
+                     "#include </test-header.h>\n";
+
+  std::unique_ptr<llvm::MemoryBuffer> HeaderBuf =
+      llvm::MemoryBuffer::getMemBuffer(header);
+  std::unique_ptr<llvm::MemoryBuffer> MainBuf =
+      llvm::MemoryBuffer::getMemBuffer(main);
+
+  const FileEntry *HeaderFile =
+      FileMgr.getVirtualFile("/test-header.h", HeaderBuf->getBufferSize(), 0);
+  SourceMgr.overrideFileContents(HeaderFile, std::move(HeaderBuf));
+  const FileEntry *MainFile =
+      FileMgr.getVirtualFile("main.cpp", MainBuf->getBufferSize(), 0);
+  SourceMgr.overrideFileContents(MainFile, std::move(MainBuf));
+  SourceMgr.setMainFileID(
+      SourceMgr.createFileID(MainFile, SourceLocation(), SrcMgr::C_User));
+
+  TrivialModuleLoader ModLoader;
+  MemoryBufferCache PCMCache;
+  HeaderSearch HeaderInfo(std::make_shared<HeaderSearchOptions>(), SourceMgr,
+                          Diags, LangOpts, &*Target);
+  Preprocessor PP(std::make_shared<PreprocessorOptions>(), Diags, LangOpts,
+                  SourceMgr, PCMCache, HeaderInfo, ModLoader,
+                  /*IILookup =*/nullptr,
+                  /*OwnsHeaderSearch =*/false);
+  PP.Initialize(*Target);
+
+  PP.EnterMainSourceFile();
+
+  while (1) {
+    Token tok;
+    PP.Lex(tok);
+    if (tok.is(tok::eof))
+      break;
+  }
+
+  llvm::SmallVector<FileID, 2> MainFileID =
+      SourceMgr.findFileIDsForFile(MainFile);
+  ASSERT_EQ(1U, MainFileID.size());
+  ASSERT_EQ(MainFileID[0], SourceMgr.getMainFileID());
+
+  llvm::SmallVector<FileID, 2> Files = SourceMgr.findFileIDsForFile(HeaderFile);
+
+  ASSERT_EQ(2U, Files.size());
+  ASSERT_NE(Files[0], Files[1]);
+  SourceLocation Loc1 = SourceMgr.translateLineCol(Files[0], 1, 1);
+  SourceLocation Loc2 = SourceMgr.translateLineCol(Files[1], 1, 1);
+  ASSERT_NE(Loc1, Loc2);
+  ASSERT_TRUE(SourceMgr.isBeforeInTranslationUnit(Loc1, Loc2));
+}
+
 #endif
 
 } // anonymous namespace