1 files changed, 103 insertions, 2 deletions

diff --git a/unittests/Lex/PPCallbacksTest.cpp b/unittests/Lex/PPCallbacksTest.cpp
index 36bd5f9..9405a84 100644
--- a/unittests/Lex/PPCallbacksTest.cpp
+++ b/unittests/Lex/PPCallbacksTest.cpp
@@ -18,8 +18,12 @@
 #include "clang/Lex/HeaderSearchOptions.h"
 #include "clang/Lex/ModuleLoader.h"
 #include "clang/Lex/PreprocessorOptions.h"
+#include "clang/Parse/Parser.h"
+#include "clang/Sema/Sema.h"
+#include "clang/AST/ASTContext.h"
+#include "clang/AST/ASTConsumer.h"
 #include "llvm/ADT/SmallString.h"
-#include "llvm/Support/PathV2.h"
+#include "llvm/Support/Path.h"
 #include "gtest/gtest.h"
 
 using namespace llvm;
@@ -77,6 +81,31 @@ public:
   const Module* Imported;
 };
 
+// Stub to collect data from PragmaOpenCLExtension callbacks.
+class PragmaOpenCLExtensionCallbacks : public PPCallbacks {
+public:
+  typedef struct {
+    SmallString<16> Name;
+    unsigned State;
+  } CallbackParameters;
+
+  PragmaOpenCLExtensionCallbacks() : Name("Not called."), State(99) {};
+
+  void PragmaOpenCLExtension(
+    clang::SourceLocation NameLoc, const clang::IdentifierInfo *Name,
+    clang::SourceLocation StateLoc, unsigned State) {
+      this->NameLoc = NameLoc;
+      this->Name = Name->getName();
+      this->StateLoc = StateLoc;
+      this->State = State;
+  };
+
+  SourceLocation NameLoc;
+  SmallString<16> Name;
+  SourceLocation StateLoc;
+  unsigned State;
+};
+
 // PPCallbacks test fixture.
 class PPCallbacksTest : public ::testing::Test {
 protected:
@@ -133,7 +162,8 @@ protected:
     VoidModuleLoader ModLoader;
 
     IntrusiveRefCntPtr<HeaderSearchOptions> HSOpts = new HeaderSearchOptions();
-    HeaderSearch HeaderInfo(HSOpts, FileMgr, Diags, LangOpts, Target.getPtr());
+    HeaderSearch HeaderInfo(HSOpts, SourceMgr, Diags, LangOpts,
+                            Target.getPtr());
     AddFakeHeader(HeaderInfo, HeaderPath, SystemHeader);
 
     IntrusiveRefCntPtr<PreprocessorOptions> PPOpts = new PreprocessorOptions();
@@ -159,6 +189,53 @@ protected:
     // Callbacks have been executed at this point -- return filename range.
     return Callbacks->FilenameRange;
   }
+
+  PragmaOpenCLExtensionCallbacks::CallbackParameters 
+  PragmaOpenCLExtensionCall(const char* SourceText) {
+    LangOptions OpenCLLangOpts;
+    OpenCLLangOpts.OpenCL = 1;
+
+    MemoryBuffer* sourceBuf = MemoryBuffer::getMemBuffer(SourceText, "test.cl");
+    (void)SourceMgr.createMainFileIDForMemBuffer(sourceBuf);
+
+    VoidModuleLoader ModLoader;
+    HeaderSearch HeaderInfo(new HeaderSearchOptions, SourceMgr, Diags, 
+                            OpenCLLangOpts, Target.getPtr());
+
+    Preprocessor PP(new PreprocessorOptions(), Diags, OpenCLLangOpts, 
+                    Target.getPtr(),
+                    SourceMgr, HeaderInfo, ModLoader,
+                   /*IILookup =*/ 0,
+                    /*OwnsHeaderSearch =*/false,
+                    /*DelayInitialization =*/ false);
+
+    // parser actually sets correct pragma handlers for preprocessor
+    // according to LangOptions, so we init Parser to register opencl
+    // pragma handlers
+    ASTContext Context(OpenCLLangOpts, SourceMgr, Target.getPtr(), 
+                       PP.getIdentifierTable(), PP.getSelectorTable(), 
+                       PP.getBuiltinInfo(), 0);    
+    ASTConsumer Consumer;
+    Sema S(PP, Context, Consumer);
+    Parser P(PP, S, false);
+    PragmaOpenCLExtensionCallbacks* Callbacks = new PragmaOpenCLExtensionCallbacks;
+    PP.addPPCallbacks(Callbacks); // Takes ownership.
+
+    // Lex source text.
+    PP.EnterMainSourceFile();
+    while (true) {
+      Token Tok;
+      PP.Lex(Tok);
+      if (Tok.is(tok::eof))
+        break;
+    }
+
+    PragmaOpenCLExtensionCallbacks::CallbackParameters RetVal = {
+      Callbacks->Name,
+      Callbacks->State
+    };
+    return RetVal;    
+  }
 };
 
 TEST_F(PPCallbacksTest, QuotedFilename) {
@@ -247,4 +324,28 @@ TEST_F(PPCallbacksTest, TrigraphInMacro) {
   ASSERT_EQ("\"tri\?\?-graph.h\"", GetSourceString(Range));
 }
 
+TEST_F(PPCallbacksTest, OpenCLExtensionPragmaEnabled) {
+  const char* Source =
+    "#pragma OPENCL EXTENSION cl_khr_fp64 : enable\n";
+
+  PragmaOpenCLExtensionCallbacks::CallbackParameters Parameters =
+    PragmaOpenCLExtensionCall(Source);
+
+  ASSERT_EQ("cl_khr_fp64", Parameters.Name);
+  unsigned ExpectedState = 1;
+  ASSERT_EQ(ExpectedState, Parameters.State);
+}
+
+TEST_F(PPCallbacksTest, OpenCLExtensionPragmaDisabled) {
+  const char* Source =
+    "#pragma OPENCL EXTENSION cl_khr_fp16 : disable\n";
+
+  PragmaOpenCLExtensionCallbacks::CallbackParameters Parameters =
+    PragmaOpenCLExtensionCall(Source);
+
+  ASSERT_EQ("cl_khr_fp16", Parameters.Name);
+  unsigned ExpectedState = 0;
+  ASSERT_EQ(ExpectedState, Parameters.State);
+}
+
 } // anonoymous namespace