[CrossTU] Add a function to retrieve original source location.

[balazske]

Summary

A new function will be added to get the original SourceLocation
for a SourceLocation that was imported as result of getCrossTUDefinition.
The returned SourceLocation is in the context of the (original)
SourceManager for the original source file. Additionally the
ASTUnit object for that source file is returned. This is needed
to get a SourceManager to operate on with the returned source location.

The new function works if multiple different source files are loaded
with the same CrossTU context.

[martong]

LGTM! Thanks!

[martong]

clang/lib/CrossTU/CrossTranslationUnit.cpp

476

The Importer is not used, perhaps it is not needed in the lambda?

Or we might have the mapping of ToID->FromID in the ASTImporter, and the FromID->FromUnit mapping in the CTUContext?

[balazske]

Removed ASTImporter from "callback" function.

[balazske]

clang/lib/CrossTU/CrossTranslationUnit.cpp

476

Yes the Importer is not needed (if needed later it can be added to the lambda as a variable).
Additional mappings of FileID seems not needed now, the current solution is sufficient for this use case. If there is a ToID->FromID in the ASTImporter a "callback" lambda would be still needed to make the FromID->FromUnit mapping.

[balazske]

clang/include/clang/CrossTU/CrossTranslationUnit.h

175

It would be better to use std::tuple<SourceLocation, Preprocessor &, ASTContext &> here. (Or * if & does not work.)

[balazske]

Update diff to one commit with all changes.

clang/include/clang/AST/ASTImporter.h

class TypeSourceInfo;
   /// Imports selected nodes from one AST context into another context,
   /// merging AST nodes where appropriate.
   class ASTImporter {
     friend class ASTNodeImporter;
   public:
     using NonEquivalentDeclSet = llvm::DenseSet<std::pair<Decl *, Decl *>>;
     using ImportedCXXBaseSpecifierMap =
         llvm::DenseMap<const CXXBaseSpecifier *, CXXBaseSpecifier *>;
+    using FileIDImportHandlerType = std::function<void(
+        FileID /*ToID*/, FileID /*FromID*/, ASTImporter & /*Importer*/)>;
 
     // An ImportPath is the list of the AST nodes which we visit during an
     // Import call.
     // If node `A` depends on node `B` then the path contains an `A`->`B` edge.
     // From the call stack of the import functions we can read the very same
     // path.
     //
     // Now imagine the following AST, where the `->` represents dependency in
private:
       VecTy Nodes;
       // Auxiliary container to be able to answer "Do we have a cycle ending
       // at last element?" as fast as possible.
       // We count each Decl's occurrence over the path.
       llvm::SmallDenseMap<Decl *, int, 32> Aux;
     };
 
   private:
+    FileIDImportHandlerType FileIDImportHandler;
+
     std::shared_ptr<ASTImporterSharedState> SharedState = nullptr;
 
     /// The path which we go through during the import of a given AST node.
     ImportPathTy ImportPath;
     /// Sometimes we have to save some part of an import path, so later we can
     /// set up properties to the saved nodes.
     /// We may have several of these import paths associated to one Decl.
     using SavedImportPathsForOneDecl =
public:
     /// If not set then the original C/C++ lookup is used.
     ASTImporter(ASTContext &ToContext, FileManager &ToFileManager,
                 ASTContext &FromContext, FileManager &FromFileManager,
                 bool MinimalImport,
                 std::shared_ptr<ASTImporterSharedState> SharedState = nullptr);
 
     virtual ~ASTImporter();
 
+    /// Set a callback function for FileID import handling.
+    /// The function is invoked when a FileID is imported from the From context.
+    /// The imported FileID in the To context and the original FileID in the
+    /// From context, and the ASTImporter itself are passed to it.
+    void setFileIDImportHandler(FileIDImportHandlerType H) {
+      FileIDImportHandler = H;
+    }
+
     /// Whether the importer will perform a minimal import, creating
     /// to-be-completed forward declarations when possible.
     bool isMinimalImport() const { return Minimal; }
 
     /// \brief Import the given object, returns the result.
     ///
     /// \param To Import the object into this variable.
     /// \param From Object to import.

clang/include/clang/CrossTU/CrossTranslationUnit.h

llvm::Expected<ASTUnit *> loadExternalAST(StringRef LookupName,
                                             StringRef IndexName,
                                             bool DisplayCTUProgress = false);
 
   /// This function merges a definition from a separate AST Unit into
   ///        the current one which was created by the compiler instance that
   ///        was passed to the constructor.
   ///
   /// \return Returns the resulting definition or an error.
-  llvm::Expected<const FunctionDecl *> importDefinition(const FunctionDecl *FD);
-  llvm::Expected<const VarDecl *> importDefinition(const VarDecl *VD);
+  llvm::Expected<const FunctionDecl *> importDefinition(const FunctionDecl *FD,
+                                                        ASTUnit *Unit);
+  llvm::Expected<const VarDecl *> importDefinition(const VarDecl *VD,
+                                                   ASTUnit *Unit);
 
   /// Get a name to identify a named decl.
   static std::string getLookupName(const NamedDecl *ND);
 
   /// Emit diagnostics for the user for potential configuration errors.
   void emitCrossTUDiagnostics(const IndexError &IE);
 
+  /// Determine the original source location in the original TU for an
+  /// imported source location.
+  /// \p ToLoc Source location in the imported-to AST.
+  /// \return Source location in the imported-from AST and the corresponding
+  /// ASTUnit object (the AST was loaded from a file using an internal ASTUnit
+  /// object that is returned here).
+  /// If any error happens (ToLoc is a non-imported source location) empty is
+  /// returned.
+  llvm::Optional<std::pair<SourceLocation /*FromLoc*/, ASTUnit *>>
+  getImportedFromSourceLocation(const clang::SourceLocation &ToLoc) const;
+
 private:
+  using ImportedFileIDMap =
+      llvm::DenseMap<FileID, std::pair<FileID, ASTUnit *>>;
+
   void lazyInitImporterSharedSt(TranslationUnitDecl *ToTU);
-  ASTImporter &getOrCreateASTImporter(ASTContext &From);
+  ASTImporter &getOrCreateASTImporter(ASTUnit *Unit);
   template <typename T>
   llvm::Expected<const T *> getCrossTUDefinitionImpl(const T *D,
                                                      StringRef CrossTUDir,
                                                      StringRef IndexName,
                                                      bool DisplayCTUProgress);
   template <typename T>
   const T *findDefInDeclContext(const DeclContext *DC,
                                 StringRef LookupName);
   template <typename T>
-  llvm::Expected<const T *> importDefinitionImpl(const T *D);
+  llvm::Expected<const T *> importDefinitionImpl(const T *D, ASTUnit *Unit);
 
   llvm::StringMap<std::unique_ptr<clang::ASTUnit>> FileASTUnitMap;
   llvm::StringMap<clang::ASTUnit *> NameASTUnitMap;
   llvm::StringMap<std::string> NameFileMap;
   llvm::DenseMap<TranslationUnitDecl *, std::unique_ptr<ASTImporter>>
       ASTUnitImporterMap;
   CompilerInstance &CI;
   ASTContext &Context;
   std::shared_ptr<ASTImporterSharedState> ImporterSharedSt;
+  /// Map of imported FileID's (in "To" context) to FileID in "From" context
+  /// and the ASTUnit for the From context.
+  /// This map is used by getImportedFromSourceLocation to lookup a FileID and
+  /// its Preprocessor when knowing only the FileID in the 'To' context. The
+  /// FileID could be imported by any of multiple 'From' ASTImporter objects.
+  /// we do not want to loop over all ASTImporter's to find the one that
+  /// imported the FileID.
+  ImportedFileIDMap ImportedFileIDs;
+
   /// \p CTULoadTreshold should serve as an upper limit to the number of TUs
   /// imported in order to reduce the memory footprint of CTU analysis.
   const unsigned CTULoadThreshold;
   unsigned NumASTLoaded{0u};
 };
 
 } // namespace cross_tu
 } // namespace clang
 
 #endif // LLVM_CLANG_CROSSTU_CROSSTRANSLATIONUNIT_H

clang/lib/AST/ASTImporter.cpp

if (ToID.isInvalid() || IsBuiltin) {
       ToID = ToSM.createFileID(std::move(ToBuf),
                                FromSLoc.getFile().getFileCharacteristic());
     }
   }
 
   assert(ToID.isValid() && "Unexpected invalid fileID was created.");
 
   ImportedFileIDs[FromID] = ToID;
+
+  if (FileIDImportHandler)
+    FileIDImportHandler(ToID, FromID, *this);
+
   return ToID;
 }
 
 Expected<CXXCtorInitializer *> ASTImporter::Import(CXXCtorInitializer *From) {
   ExpectedExpr ToExprOrErr = Import(From->getInit());
   if (!ToExprOrErr)
     return ToExprOrErr.takeError();
 

clang/lib/CrossTU/CrossTranslationUnit.cpp

if (LangTo.CPlusPlus11 != LangFrom.CPlusPlus11 ||
       LangTo.CPlusPlus2a != LangFrom.CPlusPlus2a) {
     ++NumLangDialectMismatch;
     return llvm::make_error<IndexError>(
         index_error_code::lang_dialect_mismatch);
   }
 
   TranslationUnitDecl *TU = Unit->getASTContext().getTranslationUnitDecl();
   if (const T *ResultDecl = findDefInDeclContext<T>(TU, LookupName))
-    return importDefinition(ResultDecl);
+    return importDefinition(ResultDecl, Unit);
   return llvm::make_error<IndexError>(index_error_code::failed_import);
 }
 
 llvm::Expected<const FunctionDecl *>
 CrossTranslationUnitContext::getCrossTUDefinition(const FunctionDecl *FD,
                                                   StringRef CrossTUDir,
                                                   StringRef IndexName,
                                                   bool DisplayCTUProgress) {
llvm::Expected<ASTUnit *> CrossTranslationUnitContext::loadExternalAST(
   if (!Unit)
     return llvm::make_error<IndexError>(
         index_error_code::failed_to_get_external_ast);
   return Unit;
 }
 
 template <typename T>
 llvm::Expected<const T *>
-CrossTranslationUnitContext::importDefinitionImpl(const T *D) {
+CrossTranslationUnitContext::importDefinitionImpl(const T *D, ASTUnit *Unit) {
   assert(hasBodyOrInit(D) && "Decls to be imported should have body or init.");
 
-  ASTImporter &Importer = getOrCreateASTImporter(D->getASTContext());
+  assert(&D->getASTContext() == &Unit->getASTContext() &&
+         "ASTContext of Decl and the unit should match.");
+  ASTImporter &Importer = getOrCreateASTImporter(Unit);
+
   auto ToDeclOrError = Importer.Import(D);
   if (!ToDeclOrError) {
     handleAllErrors(ToDeclOrError.takeError(),
                     [&](const ImportError &IE) {
                       switch (IE.Error) {
                       case ImportError::NameConflict:
                         ++NumNameConflicts;
                          break;
CrossTranslationUnitContext::importDefinitionImpl(const T *D, ASTUnit *Unit) {
   auto *ToDecl = cast<T>(*ToDeclOrError);
   assert(hasBodyOrInit(ToDecl) && "Imported Decl should have body or init.");
   ++NumGetCTUSuccess;
 
   return ToDecl;
 }
 
 llvm::Expected<const FunctionDecl *>
-CrossTranslationUnitContext::importDefinition(const FunctionDecl *FD) {
-  return importDefinitionImpl(FD);
+CrossTranslationUnitContext::importDefinition(const FunctionDecl *FD,
+                                              ASTUnit *Unit) {
+  return importDefinitionImpl(FD, Unit);
 }
 
 llvm::Expected<const VarDecl *>
-CrossTranslationUnitContext::importDefinition(const VarDecl *VD) {
-  return importDefinitionImpl(VD);
+CrossTranslationUnitContext::importDefinition(const VarDecl *VD,
+                                              ASTUnit *Unit) {
+  return importDefinitionImpl(VD, Unit);
 }
 
 void CrossTranslationUnitContext::lazyInitImporterSharedSt(
     TranslationUnitDecl *ToTU) {
   if (!ImporterSharedSt)
     ImporterSharedSt = std::make_shared<ASTImporterSharedState>(*ToTU);
 }
 
 ASTImporter &
-CrossTranslationUnitContext::getOrCreateASTImporter(ASTContext &From) {
+CrossTranslationUnitContext::getOrCreateASTImporter(ASTUnit *Unit) {
+  ASTContext &From = Unit->getASTContext();
+
   auto I = ASTUnitImporterMap.find(From.getTranslationUnitDecl());
   if (I != ASTUnitImporterMap.end())
     return *I->second;
   lazyInitImporterSharedSt(Context.getTranslationUnitDecl());
   ASTImporter *NewImporter = new ASTImporter(
       Context, Context.getSourceManager().getFileManager(), From,
       From.getSourceManager().getFileManager(), false, ImporterSharedSt);
+  NewImporter->setFileIDImportHandler(
+      [this, Unit](FileID ToID, FileID FromID, ASTImporter &Importer) {
+        assert(ImportedFileIDs.find(ToID) == ImportedFileIDs.end() &&
+               "FileID already imported, should not happen.");
+        ImportedFileIDs[ToID] = std::make_pair(FromID, Unit);
+      });
   ASTUnitImporterMap[From.getTranslationUnitDecl()].reset(NewImporter);
   return *NewImporter;
 }
 
+llvm::Optional<std::pair<SourceLocation, ASTUnit *>>
+CrossTranslationUnitContext::getImportedFromSourceLocation(
+    const clang::SourceLocation &ToLoc) const {
+  const SourceManager &SM = Context.getSourceManager();
+  auto DecToLoc = SM.getDecomposedLoc(ToLoc);
+
+  auto I = ImportedFileIDs.find(DecToLoc.first);
+  if (I == ImportedFileIDs.end())
+    return {};
+
+  FileID FromID = I->second.first;
+  clang::ASTUnit *Unit = I->second.second;
+  SourceLocation FromLoc =
+      Unit->getSourceManager().getComposedLoc(FromID, DecToLoc.second);
+
+  return std::make_pair(FromLoc, Unit);
+}
+
 } // namespace cross_tu
 } // namespace clang

clang/unittests/CrossTU/CrossTranslationUnitTest.cpp

 namespace {
 
 class CTUASTConsumer : public clang::ASTConsumer {
 public:
   explicit CTUASTConsumer(clang::CompilerInstance &CI, bool *Success)
       : CTU(CI), Success(Success) {}
 
   void HandleTranslationUnit(ASTContext &Ctx) {
-    const TranslationUnitDecl *TU = Ctx.getTranslationUnitDecl();
+    auto FindFInTU = [](const TranslationUnitDecl *TU) {
       const FunctionDecl *FD = nullptr;
       for (const Decl *D : TU->decls()) {
         FD = dyn_cast<FunctionDecl>(D);
         if (FD && FD->getName() == "f")
           break;
       }
+      return FD;
+    };
+
+    const TranslationUnitDecl *TU = Ctx.getTranslationUnitDecl();
+    const FunctionDecl *FD = FindFInTU(TU);
     assert(FD && FD->getName() == "f");
     bool OrigFDHasBody = FD->hasBody();
 
     // Prepare the index file and the AST file.
     int ASTFD;
     llvm::SmallString<256> ASTFileName;
     ASSERT_FALSE(
         llvm::sys::fs::createTemporaryFile("f_ast", "ast", ASTFD, ASTFileName));
void HandleTranslationUnit(ASTContext &Ctx) {
     // Load the definition from the AST file.
     llvm::Expected<const FunctionDecl *> NewFDorError = handleExpected(
         CTU.getCrossTUDefinition(FD, "", IndexFileName, false),
         []() { return nullptr; }, [](IndexError &) {});
 
     if (NewFDorError) {
       const FunctionDecl *NewFD = *NewFDorError;
       *Success = NewFD && NewFD->hasBody() && !OrigFDHasBody;
+
+      if (NewFD) {
+        // Check GetImportedFromSourceLocation.
+        llvm::Optional<std::pair<SourceLocation, ASTUnit *>> SLocResult =
+            CTU.getImportedFromSourceLocation(NewFD->getLocation());
+        EXPECT_TRUE(SLocResult);
+        if (SLocResult) {
+          SourceLocation OrigSLoc = (*SLocResult).first;
+          ASTUnit *OrigUnit = (*SLocResult).second;
+          // OrigUnit is created internally by CTU (is not the
+          // ASTWithDefinition).
+          TranslationUnitDecl *OrigTU =
+              OrigUnit->getASTContext().getTranslationUnitDecl();
+          const FunctionDecl *FDWithDefinition = FindFInTU(OrigTU);
+          EXPECT_TRUE(FDWithDefinition);
+          if (FDWithDefinition) {
+            EXPECT_EQ(FDWithDefinition->getName(), "f");
+            EXPECT_TRUE(FDWithDefinition->isThisDeclarationADefinition());
+            EXPECT_EQ(OrigSLoc, FDWithDefinition->getLocation());
+          }
+        }
+      }
     }
   }
 
 private:
   CrossTranslationUnitContext CTU;
   bool *Success;
 };