1 files changed, 97 insertions, 22 deletions

diff --git a/unittests/Lex/PPCallbacksTest.cpp b/unittests/Lex/PPCallbacksTest.cpp
index 0fa6286cc0480..838e033e3df49 100644
--- a/unittests/Lex/PPCallbacksTest.cpp
+++ b/unittests/Lex/PPCallbacksTest.cpp
@@ -65,6 +65,29 @@ public:
   SrcMgr::CharacteristicKind FileType;
 };
 
+class CondDirectiveCallbacks : public PPCallbacks {
+public:
+  struct Result {
+    SourceRange ConditionRange;
+    ConditionValueKind ConditionValue;
+
+    Result(SourceRange R, ConditionValueKind K)
+        : ConditionRange(R), ConditionValue(K) {}
+  };
+
+  std::vector<Result> Results;
+
+  void If(SourceLocation Loc, SourceRange ConditionRange,
+          ConditionValueKind ConditionValue) override {
+    Results.emplace_back(ConditionRange, ConditionValue);
+  }
+
+  void Elif(SourceLocation Loc, SourceRange ConditionRange,
+            ConditionValueKind ConditionValue, SourceLocation IfLoc) override {
+    Results.emplace_back(ConditionRange, ConditionValue);
+  }
+};
+
 // Stub to collect data from PragmaOpenCLExtension callbacks.
 class PragmaOpenCLExtensionCallbacks : public PPCallbacks {
 public:
@@ -95,7 +118,7 @@ public:
 class PPCallbacksTest : public ::testing::Test {
 protected:
   PPCallbacksTest()
-      : InMemoryFileSystem(new vfs::InMemoryFileSystem),
+      : InMemoryFileSystem(new llvm::vfs::InMemoryFileSystem),
         FileMgr(FileSystemOptions(), InMemoryFileSystem),
         DiagID(new DiagnosticIDs()), DiagOpts(new DiagnosticOptions()),
         Diags(DiagID, DiagOpts.get(), new IgnoringDiagConsumer()),
@@ -104,7 +127,7 @@ protected:
     Target = TargetInfo::CreateTargetInfo(Diags, TargetOpts);
   }
 
-  IntrusiveRefCntPtr<vfs::InMemoryFileSystem> InMemoryFileSystem;
+  IntrusiveRefCntPtr<llvm::vfs::InMemoryFileSystem> InMemoryFileSystem;
   FileManager FileMgr;
   IntrusiveRefCntPtr<DiagnosticIDs> DiagID;
   IntrusiveRefCntPtr<DiagnosticOptions> DiagOpts;
@@ -116,17 +139,17 @@ protected:
 
   // Register a header path as a known file and add its location
   // to search path.
-  void AddFakeHeader(HeaderSearch& HeaderInfo, const char* HeaderPath, 
-    bool IsSystemHeader) {
-      // Tell FileMgr about header.
-      InMemoryFileSystem->addFile(HeaderPath, 0,
-                                  llvm::MemoryBuffer::getMemBuffer("\n"));
-
-      // Add header's parent path to search path.
-      StringRef SearchPath = llvm::sys::path::parent_path(HeaderPath);
-      const DirectoryEntry *DE = FileMgr.getDirectory(SearchPath);
-      DirectoryLookup DL(DE, SrcMgr::C_User, false);
-      HeaderInfo.AddSearchPath(DL, IsSystemHeader);
+  void AddFakeHeader(HeaderSearch &HeaderInfo, const char *HeaderPath,
+                     bool IsSystemHeader) {
+    // Tell FileMgr about header.
+    InMemoryFileSystem->addFile(HeaderPath, 0,
+                                llvm::MemoryBuffer::getMemBuffer("\n"));
+
+    // Add header's parent path to search path.
+    StringRef SearchPath = llvm::sys::path::parent_path(HeaderPath);
+    const DirectoryEntry *DE = FileMgr.getDirectory(SearchPath);
+    DirectoryLookup DL(DE, SrcMgr::C_User, false);
+    HeaderInfo.AddSearchPath(DL, IsSystemHeader);
   }
 
   // Get the raw source string of the range.
@@ -137,10 +160,20 @@ protected:
     return StringRef(B, E - B);
   }
 
+  StringRef GetSourceStringToEnd(CharSourceRange Range) {
+    const char *B = SourceMgr.getCharacterData(Range.getBegin());
+    const char *E = SourceMgr.getCharacterData(Range.getEnd());
+
+    return StringRef(
+        B,
+        E - B + Lexer::MeasureTokenLength(Range.getEnd(), SourceMgr, LangOpts));
+  }
+
   // Run lexer over SourceText and collect FilenameRange from
   // the InclusionDirective callback.
-  CharSourceRange InclusionDirectiveFilenameRange(const char* SourceText, 
-      const char* HeaderPath, bool SystemHeader) {
+  CharSourceRange InclusionDirectiveFilenameRange(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)));
@@ -198,8 +231,38 @@ protected:
     return Callbacks;
   }
 
-  PragmaOpenCLExtensionCallbacks::CallbackParameters 
-  PragmaOpenCLExtensionCall(const char* SourceText) {
+  std::vector<CondDirectiveCallbacks::Result>
+  DirectiveExprRange(StringRef SourceText) {
+    TrivialModuleLoader ModLoader;
+    MemoryBufferCache PCMCache;
+    std::unique_ptr<llvm::MemoryBuffer> Buf =
+        llvm::MemoryBuffer::getMemBuffer(SourceText);
+    SourceMgr.setMainFileID(SourceMgr.createFileID(std::move(Buf)));
+    HeaderSearch HeaderInfo(std::make_shared<HeaderSearchOptions>(), SourceMgr,
+                            Diags, LangOpts, Target.get());
+    Preprocessor PP(std::make_shared<PreprocessorOptions>(), Diags, LangOpts,
+                    SourceMgr, PCMCache, HeaderInfo, ModLoader,
+                    /*IILookup =*/nullptr,
+                    /*OwnsHeaderSearch =*/false);
+    PP.Initialize(*Target);
+    auto *Callbacks = new CondDirectiveCallbacks;
+    PP.addPPCallbacks(std::unique_ptr<PPCallbacks>(Callbacks));
+
+    // Lex source text.
+    PP.EnterMainSourceFile();
+
+    while (true) {
+      Token Tok;
+      PP.Lex(Tok);
+      if (Tok.is(tok::eof))
+        break;
+    }
+
+    return Callbacks->Results;
+  }
+
+  PragmaOpenCLExtensionCallbacks::CallbackParameters
+  PragmaOpenCLExtensionCall(const char *SourceText) {
     LangOptions OpenCLLangOpts;
     OpenCLLangOpts.OpenCL = 1;
 
@@ -221,9 +284,8 @@ protected:
     // parser actually sets correct pragma handlers for preprocessor
     // according to LangOptions, so we init Parser to register opencl
     // pragma handlers
-    ASTContext Context(OpenCLLangOpts, SourceMgr,
-                       PP.getIdentifierTable(), PP.getSelectorTable(), 
-                       PP.getBuiltinInfo());
+    ASTContext Context(OpenCLLangOpts, SourceMgr, PP.getIdentifierTable(),
+                       PP.getSelectorTable(), PP.getBuiltinInfo());
     Context.InitBuiltinTypes(*Target);
 
     ASTConsumer Consumer;
@@ -245,7 +307,7 @@ protected:
       Callbacks->Name,
       Callbacks->State
     };
-    return RetVal;    
+    return RetVal;
   }
 };
 
@@ -368,4 +430,17 @@ TEST_F(PPCallbacksTest, OpenCLExtensionPragmaDisabled) {
   ASSERT_EQ(ExpectedState, Parameters.State);
 }
 
-} // anonoymous namespace
+TEST_F(PPCallbacksTest, DirectiveExprRanges) {
+  const auto &Results8 =
+      DirectiveExprRange("#define FLOOFY 0\n#if __FILE__ > FLOOFY\n#endif\n");
+  EXPECT_EQ(Results8.size(), 1U);
+  EXPECT_EQ(
+      GetSourceStringToEnd(CharSourceRange(Results8[0].ConditionRange, false)),
+      " __FILE__ > FLOOFY\n#");
+  EXPECT_EQ(
+      Lexer::getSourceText(CharSourceRange(Results8[0].ConditionRange, false),
+                           SourceMgr, LangOpts),
+      " __FILE__ > FLOOFY\n");
+}
+
+} // namespace