From 6e017188b7607eb0f0b248881319bfd803cb8f78 Mon Sep 17 00:00:00 2001
From: Kadir Cetinkaya <kadircet@google.com>
Date: Wed, 15 Apr 2020 22:00:19 +0200
Subject: [PATCH] [clangd] Store ppdirective in Inclusion

Summary:
This will enable PreamblePatching proposed in D77392 craft a more
informed patch.

Reviewers: sammccall

Subscribers: ilya-biryukov, MaskRay, jkorous, arphaman, usaxena95, cfe-commits

Tags: #clang

Differential Revision: https://reviews.llvm.org/D78235
---
 clang-tools-extra/clangd/Headers.cpp                |  3 ++-
 clang-tools-extra/clangd/Headers.h                  |  9 ++++++---
 clang-tools-extra/clangd/unittests/HeadersTests.cpp | 15 +++++++++++++++
 3 files changed, 23 insertions(+), 4 deletions(-)

diff --git a/clang-tools-extra/clangd/Headers.cpp b/clang-tools-extra/clangd/Headers.cpp
index 295d6e6..dd4cafd 100644
--- a/clang-tools-extra/clangd/Headers.cpp
+++ b/clang-tools-extra/clangd/Headers.cpp
@@ -28,7 +28,7 @@ public:
 
   // Record existing #includes - both written and resolved paths. Only #includes
   // in the main file are collected.
-  void InclusionDirective(SourceLocation HashLoc, const Token & /*IncludeTok*/,
+  void InclusionDirective(SourceLocation HashLoc, const Token &IncludeTok,
                           llvm::StringRef FileName, bool IsAngled,
                           CharSourceRange FilenameRange, const FileEntry *File,
                           llvm::StringRef /*SearchPath*/,
@@ -44,6 +44,7 @@ public:
       Inc.Resolved = std::string(File ? File->tryGetRealPathName() : "");
       Inc.HashOffset = SM.getFileOffset(HashLoc);
       Inc.FileKind = FileKind;
+      Inc.Directive = IncludeTok.getIdentifierInfo()->getPPKeywordID();
     }
     if (File) {
       auto *IncludingFileEntry = SM.getFileEntryForID(SM.getFileID(HashLoc));
diff --git a/clang-tools-extra/clangd/Headers.h b/clang-tools-extra/clangd/Headers.h
index c0c7ab1..1f39687 100644
--- a/clang-tools-extra/clangd/Headers.h
+++ b/clang-tools-extra/clangd/Headers.h
@@ -13,6 +13,7 @@
 #include "Protocol.h"
 #include "SourceCode.h"
 #include "index/Symbol.h"
+#include "clang/Basic/TokenKinds.h"
 #include "clang/Format/Format.h"
 #include "clang/Lex/HeaderSearch.h"
 #include "clang/Lex/PPCallbacks.h"
@@ -22,6 +23,7 @@
 #include "llvm/ADT/StringSet.h"
 #include "llvm/Support/Error.h"
 #include "llvm/Support/VirtualFileSystem.h"
+#include <string>
 
 namespace clang {
 namespace clangd {
@@ -50,9 +52,10 @@ llvm::SmallVector<llvm::StringRef, 1> getRankedIncludes(const Symbol &Sym);
 
 // An #include directive that we found in the main file.
 struct Inclusion {
-  Range R;             // Inclusion range.
-  std::string Written; // Inclusion name as written e.g. <vector>.
-  Path Resolved;       // Resolved path of included file. Empty if not resolved.
+  Range R;                      // Inclusion range.
+  tok::PPKeywordKind Directive; // Directive used for inclusion, e.g. import
+  std::string Written;          // Inclusion name as written e.g. <vector>.
+  Path Resolved; // Resolved path of included file. Empty if not resolved.
   unsigned HashOffset = 0; // Byte offset from start of file to #.
   SrcMgr::CharacteristicKind FileKind = SrcMgr::C_User;
 };
diff --git a/clang-tools-extra/clangd/unittests/HeadersTests.cpp b/clang-tools-extra/clangd/unittests/HeadersTests.cpp
index 1d7c828..cb760a3 100644
--- a/clang-tools-extra/clangd/unittests/HeadersTests.cpp
+++ b/clang-tools-extra/clangd/unittests/HeadersTests.cpp
@@ -11,6 +11,7 @@
 #include "Compiler.h"
 #include "TestFS.h"
 #include "TestTU.h"
+#include "clang/Basic/TokenKinds.h"
 #include "clang/Frontend/CompilerInstance.h"
 #include "clang/Frontend/CompilerInvocation.h"
 #include "clang/Frontend/FrontendActions.h"
@@ -127,6 +128,7 @@ protected:
 MATCHER_P(Written, Name, "") { return arg.Written == Name; }
 MATCHER_P(Resolved, Name, "") { return arg.Resolved == Name; }
 MATCHER_P(IncludeLine, N, "") { return arg.R.start.line == N; }
+MATCHER_P(Directive, D, "") { return arg.Directive == D; }
 
 MATCHER_P2(Distance, File, D, "") {
   if (arg.getKey() != File)
@@ -201,6 +203,19 @@ TEST_F(HeadersTest, UnResolvedInclusion) {
               UnorderedElementsAre(Distance(MainFile, 0u)));
 }
 
+TEST_F(HeadersTest, IncludeDirective) {
+  FS.Files[MainFile] = R"cpp(
+#include "foo.h"
+#import "foo.h"
+#include_next "foo.h"
+)cpp";
+
+  EXPECT_THAT(collectIncludes().MainFileIncludes,
+              UnorderedElementsAre(Directive(tok::pp_include),
+                                   Directive(tok::pp_import),
+                                   Directive(tok::pp_include_next)));
+}
+
 TEST_F(HeadersTest, InsertInclude) {
   std::string Path = testPath("sub/bar.h");
   FS.Files[Path] = "";
-- 
2.7.4