[ASTImporter] Prevent the ASTImporter from creating multiple main FileIDs.

Summary:
When importing the main FileID the ASTImporter currently gives it no include location. This means
that any SourceLocations produced for this FileID look to Clang as if they are coming from the
main FileID (as the main FileID has no include location).

Clang seems to expect that there is only one main FileID in one translation unit (which makes sense
during normal compilation), so this behavior leads to several problems when producing diagnostics,
one being that when calling `SourceManager::isBeforeInTranslationUnit` on two SourceLocations
that come from two different ASTContext instances, Clang fails to sort the SourceLocations as
the include chains of the FileIDs don't end up in a single FileID. This causes that Clang crashes
with "Unsortable locations found" in this function.

This patch gives any imported main FileIDs the main FileID of the To ASTContext as its include
location. This allows Clang to sort all imported SourceLocations as now all include chains point
to the main FileID of the To ASTContext. The exact include location is currently set to the start
of the To main file (just because that should always be a valid SourceLocation).

Reviewers: martong, a_sidorin, a.sidorin, shafik, balazske

Reviewed By: martong, a_sidorin, shafik

Subscribers: balazske, rnkovacs, cfe-commits

Tags: #clang

Differential Revision: https://reviews.llvm.org/D74542

clang/lib/AST/ASTImporter.cpp

@@ -8479,6 +8479,15 @@ Expected<FileID> ASTImporter::Import(FileID FromID, bool IsBuiltin) {
       if (!ToIncludeLoc)
         return ToIncludeLoc.takeError();
 
+      // Every FileID that is not the main FileID needs to have a valid include
+      // location so that the include chain points to the main FileID. When
+      // importing the main FileID (which has no include location), we need to
+      // create a fake include location in the main file to keep this property
+      // intact.
+      SourceLocation ToIncludeLocOrFakeLoc = *ToIncludeLoc;
+      if (FromID == FromSM.getMainFileID())
+        ToIncludeLocOrFakeLoc = ToSM.getLocForStartOfFile(ToSM.getMainFileID());
+
       if (Cache->OrigEntry && Cache->OrigEntry->getDir()) {
         // FIXME: We probably want to use getVirtualFile(), so we don't hit the
         // disk again
@@ -8490,7 +8499,7 @@ Expected<FileID> ASTImporter::Import(FileID FromID, bool IsBuiltin) {
         // point to a valid file and we get no Entry here. In this case try with
         // the memory buffer below.
         if (Entry)
-          ToID = ToSM.createFileID(*Entry, *ToIncludeLoc,
+          ToID = ToSM.createFileID(*Entry, ToIncludeLocOrFakeLoc,
                                    FromSLoc.getFile().getFileCharacteristic());
       }
     }

clang/unittests/AST/ASTImporterTest.cpp

@@ -5845,6 +5845,48 @@ TEST_P(ImportAutoFunctions, ReturnWithTypeInSwitch) {
   EXPECT_TRUE(isa<AutoType>(To->getReturnType()));
 }
 
+struct ImportSourceLocations : ASTImporterOptionSpecificTestBase {};
+
+TEST_P(ImportSourceLocations, PreserveFileIDTreeStructure) {
+  // Tests that the FileID tree structure (with the links being the include
+  // chains) is preserved while importing other files (which need to be
+  // added to this structure with fake include locations.
+
+  SourceLocation Location1;
+  {
+    auto Pattern = varDecl(hasName("X"));
+    Decl *FromTU = getTuDecl("int X;", Lang_C, "input0.c");
+    auto *FromD = FirstDeclMatcher<VarDecl>().match(FromTU, Pattern);
+
+    Location1 = Import(FromD, Lang_C)->getLocation();
+  }
+  SourceLocation Location2;
+  {
+    auto Pattern = varDecl(hasName("Y"));
+    Decl *FromTU = getTuDecl("int Y;", Lang_C, "input1.c");
+    auto *FromD = FirstDeclMatcher<VarDecl>().match(FromTU, Pattern);
+
+    Location2 = Import(FromD, Lang_C)->getLocation();
+  }
+
+  SourceManager &ToSM = ToAST->getSourceManager();
+  FileID FileID1 = ToSM.getFileID(Location1);
+  FileID FileID2 = ToSM.getFileID(Location2);
+
+  // Check that the imported files look like as if they were included from the
+  // start of the main file.
+  SourceLocation FileStart = ToSM.getLocForStartOfFile(ToSM.getMainFileID());
+  EXPECT_NE(FileID1, ToSM.getMainFileID());
+  EXPECT_NE(FileID2, ToSM.getMainFileID());
+  EXPECT_EQ(ToSM.getIncludeLoc(FileID1), FileStart);
+  EXPECT_EQ(ToSM.getIncludeLoc(FileID2), FileStart);
+
+  // Let the SourceManager check the order of the locations. The order should
+  // be the order in which the declarations are imported.
+  EXPECT_TRUE(ToSM.isBeforeInTranslationUnit(Location1, Location2));
+  EXPECT_FALSE(ToSM.isBeforeInTranslationUnit(Location2, Location1));
+}
+
 INSTANTIATE_TEST_CASE_P(ParameterizedTests, ASTImporterLookupTableTest,
                         DefaultTestValuesForRunOptions, );
 
@@ -5903,5 +5945,8 @@ INSTANTIATE_TEST_CASE_P(ParameterizedTests, ImportVariables,
 INSTANTIATE_TEST_CASE_P(ParameterizedTests, LLDBLookupTest,
                         DefaultTestValuesForRunOptions, );
 
+INSTANTIATE_TEST_CASE_P(ParameterizedTests, ImportSourceLocations,
+                        DefaultTestValuesForRunOptions, );
+
 } // end namespace ast_matchers
 } // end namespace clang