1 files changed, 47 insertions, 11 deletions

diff --git a/unittests/Lex/PPCallbacksTest.cpp b/unittests/Lex/PPCallbacksTest.cpp
index 67b56a601c71a..0fa6286cc0480 100644
--- a/unittests/Lex/PPCallbacksTest.cpp
+++ b/unittests/Lex/PPCallbacksTest.cpp
@@ -39,16 +39,18 @@ public:
                           StringRef FileName, bool IsAngled,
                           CharSourceRange FilenameRange, const FileEntry *File,
                           StringRef SearchPath, StringRef RelativePath,
-                          const Module *Imported) override {
-      this->HashLoc = HashLoc;
-      this->IncludeTok = IncludeTok;
-      this->FileName = FileName.str();
-      this->IsAngled = IsAngled;
-      this->FilenameRange = FilenameRange;
-      this->File = File;
-      this->SearchPath = SearchPath.str();
-      this->RelativePath = RelativePath.str();
-      this->Imported = Imported;
+                          const Module *Imported,
+                          SrcMgr::CharacteristicKind FileType) override {
+    this->HashLoc = HashLoc;
+    this->IncludeTok = IncludeTok;
+    this->FileName = FileName.str();
+    this->IsAngled = IsAngled;
+    this->FilenameRange = FilenameRange;
+    this->File = File;
+    this->SearchPath = SearchPath.str();
+    this->RelativePath = RelativePath.str();
+    this->Imported = Imported;
+    this->FileType = FileType;
   }
 
   SourceLocation HashLoc;
@@ -60,6 +62,7 @@ public:
   SmallString<16> SearchPath;
   SmallString<16> RelativePath;
   const Module* Imported;
+  SrcMgr::CharacteristicKind FileType;
 };
 
 // Stub to collect data from PragmaOpenCLExtension callbacks.
@@ -153,6 +156,30 @@ protected:
                     SourceMgr, PCMCache, HeaderInfo, ModLoader,
                     /*IILookup =*/nullptr,
                     /*OwnsHeaderSearch =*/false);
+    return InclusionDirectiveCallback(PP)->FilenameRange;
+  }
+
+  SrcMgr::CharacteristicKind InclusionDirectiveCharacteristicKind(
+      const char *SourceText, const char *HeaderPath, bool SystemHeader) {
+    std::unique_ptr<llvm::MemoryBuffer> Buf =
+        llvm::MemoryBuffer::getMemBuffer(SourceText);
+    SourceMgr.setMainFileID(SourceMgr.createFileID(std::move(Buf)));
+
+    TrivialModuleLoader ModLoader;
+    MemoryBufferCache PCMCache;
+
+    HeaderSearch HeaderInfo(std::make_shared<HeaderSearchOptions>(), SourceMgr,
+                            Diags, LangOpts, Target.get());
+    AddFakeHeader(HeaderInfo, HeaderPath, SystemHeader);
+
+    Preprocessor PP(std::make_shared<PreprocessorOptions>(), Diags, LangOpts,
+                    SourceMgr, PCMCache, HeaderInfo, ModLoader,
+                    /*IILookup =*/nullptr,
+                    /*OwnsHeaderSearch =*/false);
+    return InclusionDirectiveCallback(PP)->FileType;
+  }
+
+  InclusionDirectiveCallbacks *InclusionDirectiveCallback(Preprocessor &PP) {
     PP.Initialize(*Target);
     InclusionDirectiveCallbacks* Callbacks = new InclusionDirectiveCallbacks;
     PP.addPPCallbacks(std::unique_ptr<PPCallbacks>(Callbacks));
@@ -168,7 +195,7 @@ protected:
     }
 
     // Callbacks have been executed at this point -- return filename range.
-    return Callbacks->FilenameRange;
+    return Callbacks;
   }
 
   PragmaOpenCLExtensionCallbacks::CallbackParameters 
@@ -222,6 +249,15 @@ protected:
   }
 };
 
+TEST_F(PPCallbacksTest, UserFileCharacteristics) {
+  const char *Source = "#include \"quoted.h\"\n";
+
+  SrcMgr::CharacteristicKind Kind =
+      InclusionDirectiveCharacteristicKind(Source, "/quoted.h", false);
+
+  ASSERT_EQ(SrcMgr::CharacteristicKind::C_User, Kind);
+}
+
 TEST_F(PPCallbacksTest, QuotedFilename) {
   const char* Source =
     "#include \"quoted.h\"\n";