1 files changed, 38 insertions, 18 deletions

diff --git a/clang-tools-extra/clangd/Headers.cpp b/clang-tools-extra/clangd/Headers.cpp
index 3cdacffc452..3d82fe5ee6b 100644
--- a/clang-tools-extra/clangd/Headers.cpp
+++ b/clang-tools-extra/clangd/Headers.cpp
@@ -24,36 +24,50 @@ namespace {
 
 class RecordHeaders : public PPCallbacks {
 public:
-  RecordHeaders(std::set<std::string> &Headers) : Headers(Headers) {}
+  RecordHeaders(llvm::StringSet<> &WrittenHeaders,
+                llvm::StringSet<> &ResolvedHeaders)
+      : WrittenHeaders(WrittenHeaders), ResolvedHeaders(ResolvedHeaders) {}
 
   void InclusionDirective(SourceLocation /*HashLoc*/,
                           const Token & /*IncludeTok*/,
-                          llvm::StringRef /*FileName*/, bool /*IsAngled*/,
+                          llvm::StringRef FileName, bool IsAngled,
                           CharSourceRange /*FilenameRange*/,
                           const FileEntry *File, llvm::StringRef /*SearchPath*/,
                           llvm::StringRef /*RelativePath*/,
                           const Module * /*Imported*/) override {
+    WrittenHeaders.insert(
+        (IsAngled ? "<" + FileName + ">" : "\"" + FileName + "\"").str());
     if (File != nullptr && !File->tryGetRealPathName().empty())
-      Headers.insert(File->tryGetRealPathName());
+      ResolvedHeaders.insert(File->tryGetRealPathName());
   }
 
 private:
-  std::set<std::string> &Headers;
+  llvm::StringSet<> &WrittenHeaders;
+  llvm::StringSet<> &ResolvedHeaders;
 };
 
 } // namespace
 
+bool isLiteralInclude(llvm::StringRef Include) {
+  return Include.startswith("<") || Include.startswith("\"");
+}
+
+bool HeaderFile::valid() const {
+  return (Verbatim && isLiteralInclude(File)) ||
+         (!Verbatim && llvm::sys::path::is_absolute(File));
+}
+
 /// FIXME(ioeric): we might not want to insert an absolute include path if the
 /// path is not shortened.
 llvm::Expected<std::string>
 calculateIncludePath(llvm::StringRef File, llvm::StringRef Code,
-                     llvm::StringRef Header,
+                     const HeaderFile &DeclaringHeader,
+                     const HeaderFile &InsertedHeader,
                      const tooling::CompileCommand &CompileCommand,
                      IntrusiveRefCntPtr<vfs::FileSystem> FS) {
-  assert(llvm::sys::path::is_absolute(File) &&
-         llvm::sys::path::is_absolute(Header));
-
-  if (File == Header)
+  assert(llvm::sys::path::is_absolute(File));
+  assert(DeclaringHeader.valid() && InsertedHeader.valid());
+  if (File == DeclaringHeader.File || File == InsertedHeader.File)
     return "";
   FS->setCurrentWorkingDirectory(CompileCommand.Directory);
 
@@ -95,29 +109,35 @@ calculateIncludePath(llvm::StringRef File, llvm::StringRef Code,
     return llvm::make_error<llvm::StringError>(
         "Failed to begin preprocessor only action for file " + File,
         llvm::inconvertibleErrorCode());
-  std::set<std::string> ExistingHeaders;
+  llvm::StringSet<> WrittenHeaders;
+  llvm::StringSet<> ResolvedHeaders;
   Clang->getPreprocessor().addPPCallbacks(
-      llvm::make_unique<RecordHeaders>(ExistingHeaders));
+      llvm::make_unique<RecordHeaders>(WrittenHeaders, ResolvedHeaders));
   if (!Action.Execute())
     return llvm::make_error<llvm::StringError>(
         "Failed to execute preprocessor only action for file " + File,
         llvm::inconvertibleErrorCode());
-  if (ExistingHeaders.find(Header) != ExistingHeaders.end()) {
-    return llvm::make_error<llvm::StringError>(
-        Header + " is already included in " + File,
-        llvm::inconvertibleErrorCode());
-  }
+  auto Included = [&](llvm::StringRef Header) {
+    return WrittenHeaders.find(Header) != WrittenHeaders.end() ||
+           ResolvedHeaders.find(Header) != ResolvedHeaders.end();
+  };
+  if (Included(DeclaringHeader.File) || Included(InsertedHeader.File))
+    return "";
 
   auto &HeaderSearchInfo = Clang->getPreprocessor().getHeaderSearchInfo();
   bool IsSystem = false;
+
+  if (InsertedHeader.Verbatim)
+    return InsertedHeader.File;
+
   std::string Suggested = HeaderSearchInfo.suggestPathToFileForDiagnostics(
-      Header, CompileCommand.Directory, &IsSystem);
+      InsertedHeader.File, CompileCommand.Directory, &IsSystem);
   if (IsSystem)
     Suggested = "<" + Suggested + ">";
   else
     Suggested = "\"" + Suggested + "\"";
 
-  log("Suggested #include for " + Header + " is: " + Suggested);
+  log("Suggested #include for " + InsertedHeader.File + " is: " + Suggested);
   return Suggested;
 }