Add testing for CommentHandler, and fix a bug where trailing comments in #else
and #endif in non-skipped blocks were not passed to the CommentHandler. Patch
by Andy Gibbs!


git-svn-id: https://llvm.org/svn/llvm-project/cfe/trunk@159119 91177308-0d34-0410-b5e6-96231b3b80d8
diff --git a/lib/Lex/PPDirectives.cpp b/lib/Lex/PPDirectives.cpp
index b12eb81..815c474 100644
--- a/lib/Lex/PPDirectives.cpp
+++ b/lib/Lex/PPDirectives.cpp
@@ -329,7 +329,11 @@
 
         // If we popped the outermost skipping block, we're done skipping!
         if (!CondInfo.WasSkipping) {
+          // Restore the value of LexingRawMode so that trailing comments
+          // are handled correctly, if we've reached the outermost block.
+          CurPPLexer->LexingRawMode = false;
           CheckEndOfDirective("endif");
+          CurPPLexer->LexingRawMode = true;
           if (Callbacks)
             Callbacks->Endif(Tok.getLocation(), CondInfo.IfLoc);
           break;
@@ -352,7 +356,11 @@
         // entered, enter the #else block now.
         if (!CondInfo.WasSkipping && !CondInfo.FoundNonSkip) {
           CondInfo.FoundNonSkip = true;
+          // Restore the value of LexingRawMode so that trailing comments
+          // are handled correctly.
+          CurPPLexer->LexingRawMode = false;
           CheckEndOfDirective("else");
+          CurPPLexer->LexingRawMode = true;
           if (Callbacks)
             Callbacks->Else(Tok.getLocation(), CondInfo.IfLoc);
           break;
diff --git a/unittests/Tooling/CMakeLists.txt b/unittests/Tooling/CMakeLists.txt
index 876a9c1..5905049 100644
--- a/unittests/Tooling/CMakeLists.txt
+++ b/unittests/Tooling/CMakeLists.txt
@@ -1,4 +1,5 @@
 add_clang_unittest(ToolingTests
+  CommentHandlerTest.cpp
   CompilationDatabaseTest.cpp
   ToolingTest.cpp
   RecursiveASTVisitorTest.cpp
diff --git a/unittests/Tooling/CommentHandlerTest.cpp b/unittests/Tooling/CommentHandlerTest.cpp
new file mode 100644
index 0000000..f0f7797
--- /dev/null
+++ b/unittests/Tooling/CommentHandlerTest.cpp
@@ -0,0 +1,221 @@
+//===- unittest/Tooling/CommentHandlerTest.cpp -----------------------===//
+//
+//                     The LLVM Compiler Infrastructure
+//
+// This file is distributed under the University of Illinois Open Source
+// License. See LICENSE.TXT for details.
+//
+//===----------------------------------------------------------------------===//
+
+#include "TestVisitor.h"
+#include "clang/Lex/Preprocessor.h"
+
+namespace clang {
+
+struct Comment {
+  Comment(const std::string &Message, unsigned Line, unsigned Col)
+    : Message(Message), Line(Line), Col(Col) { }
+
+  std::string Message;
+  unsigned Line, Col;
+};
+
+class CommentVerifier;
+typedef std::vector<Comment> CommentList;
+
+class CommentHandlerVisitor : public TestVisitor<CommentHandlerVisitor>,
+                              public CommentHandler {
+  typedef TestVisitor<CommentHandlerVisitor> base;
+
+public:
+  CommentHandlerVisitor() : base(), PP(0), Verified(false) { }
+
+  ~CommentHandlerVisitor() {
+    EXPECT_TRUE(Verified) << "CommentVerifier not accessed";
+  }
+
+  virtual bool HandleComment(Preprocessor &PP, SourceRange Loc) {
+    assert(&PP == this->PP && "Preprocessor changed!");
+
+    SourceLocation Start = Loc.getBegin();
+    SourceManager &SM = PP.getSourceManager();
+    std::string C(SM.getCharacterData(Start),
+                  SM.getCharacterData(Loc.getEnd()));
+
+    bool Invalid;
+    unsigned CLine = SM.getSpellingLineNumber(Start, &Invalid);
+    EXPECT_TRUE(!Invalid) << "Invalid line number on comment " << C;
+
+    unsigned CCol = SM.getSpellingColumnNumber(Start, &Invalid);
+    EXPECT_TRUE(!Invalid) << "Invalid column number on comment " << C;
+
+    Comments.push_back(Comment(C, CLine, CCol));
+    return false;
+  }
+
+  CommentVerifier GetVerifier();
+
+protected:
+  virtual ASTFrontendAction* CreateTestAction() {
+    return new CommentHandlerAction(this);
+  }
+
+private:
+  Preprocessor *PP;
+  CommentList Comments;
+  bool Verified;
+
+  class CommentHandlerAction : public base::TestAction {
+  public:
+    CommentHandlerAction(CommentHandlerVisitor *Visitor)
+        : TestAction(Visitor) { }
+
+    virtual bool BeginSourceFileAction(CompilerInstance &CI,
+                                       StringRef FileName) {
+      CommentHandlerVisitor *V =
+          static_cast<CommentHandlerVisitor*>(this->Visitor);
+      V->PP = &CI.getPreprocessor();
+      V->PP->addCommentHandler(V);
+      return true;
+    }
+
+    virtual void EndSourceFileAction() {
+      CommentHandlerVisitor *V =
+          static_cast<CommentHandlerVisitor*>(this->Visitor);
+      V->PP->removeCommentHandler(V);
+    }
+  };
+};
+
+class CommentVerifier {
+  CommentList::const_iterator Current;
+  CommentList::const_iterator End;
+  Preprocessor *PP;
+
+public:
+  CommentVerifier(const CommentList &Comments, Preprocessor *PP)
+      : Current(Comments.begin()), End(Comments.end()), PP(PP)
+    { }
+
+  ~CommentVerifier() {
+    if (Current != End) {
+      EXPECT_TRUE(Current == End) << "Unexpected comment \""
+        << Current->Message << "\" at line " << Current->Line << ", column "
+        << Current->Col;
+    }
+  }
+
+  void Match(const char *Message, unsigned Line, unsigned Col) {
+    EXPECT_TRUE(Current != End) << "Comment " << Message << " not found";
+    if (Current == End) return;
+
+    const Comment &C = *Current;
+    EXPECT_TRUE(C.Message == Message && C.Line == Line && C.Col == Col)
+      <<   "Expected comment \"" << Message
+      << "\" at line " << Line   << ", column " << Col
+      << "\nActual comment   \"" << C.Message
+      << "\" at line " << C.Line << ", column " << C.Col;
+
+    ++Current;
+  }
+};
+
+CommentVerifier CommentHandlerVisitor::GetVerifier() {
+  Verified = true;
+  return CommentVerifier(Comments, PP);
+}
+
+
+TEST(CommentHandlerTest, BasicTest1) {
+  CommentHandlerVisitor Visitor;
+  EXPECT_TRUE(Visitor.runOver("class X {}; int main() { return 0; }"));
+  CommentVerifier Verifier = Visitor.GetVerifier();
+}
+
+TEST(CommentHandlerTest, BasicTest2) {
+  CommentHandlerVisitor Visitor;
+  EXPECT_TRUE(Visitor.runOver(
+        "class X {}; int main() { /* comment */ return 0; }"));
+  CommentVerifier Verifier = Visitor.GetVerifier();
+  Verifier.Match("/* comment */", 1, 26);
+}
+
+TEST(CommentHandlerTest, BasicTest3) {
+  CommentHandlerVisitor Visitor;
+  EXPECT_TRUE(Visitor.runOver(
+        "class X {}; // comment 1\n"
+        "int main() {\n"
+        "  // comment 2\n"
+        "  return 0;\n"
+        "}"));
+  CommentVerifier Verifier = Visitor.GetVerifier();
+  Verifier.Match("// comment 1", 1, 13);
+  Verifier.Match("// comment 2", 3, 3);
+}
+
+TEST(CommentHandlerTest, IfBlock1) {
+  CommentHandlerVisitor Visitor;
+  EXPECT_TRUE(Visitor.runOver(
+        "#if 0\n"
+        "// ignored comment\n"
+        "#endif\n"
+        "// visible comment\n"));
+  CommentVerifier Verifier = Visitor.GetVerifier();
+  Verifier.Match("// visible comment", 4, 1);
+}
+
+TEST(CommentHandlerTest, IfBlock2) {
+  CommentHandlerVisitor Visitor;
+  EXPECT_TRUE(Visitor.runOver(
+        "#define TEST        // visible_1\n"
+        "#ifndef TEST        // visible_2\n"
+        "                    // ignored_3\n"
+        "# ifdef UNDEFINED   // ignored_4\n"
+        "# endif             // ignored_5\n"
+        "#elif defined(TEST) // visible_6\n"
+        "# if 1              // visible_7\n"
+        "                    // visible_8\n"
+        "# else              // visible_9\n"
+        "                    // ignored_10\n"
+        "#  ifndef TEST      // ignored_11\n"
+        "#  endif            // ignored_12\n"
+        "# endif             // visible_13\n"
+        "#endif              // visible_14\n"));
+
+  CommentVerifier Verifier = Visitor.GetVerifier();
+  Verifier.Match("// visible_1", 1, 21);
+  Verifier.Match("// visible_2", 2, 21);
+  Verifier.Match("// visible_6", 6, 21);
+  Verifier.Match("// visible_7", 7, 21);
+  Verifier.Match("// visible_8", 8, 21);
+  Verifier.Match("// visible_9", 9, 21);
+  Verifier.Match("// visible_13", 13, 21);
+  Verifier.Match("// visible_14", 14, 21);
+}
+
+TEST(CommentHandlerTest, IfBlock3) {
+  const char *Source =
+        "/* commented out ...\n"
+        "#if 0\n"
+        "// enclosed\n"
+        "#endif */";
+
+  CommentHandlerVisitor Visitor;
+  EXPECT_TRUE(Visitor.runOver(Source));
+  CommentVerifier Verifier = Visitor.GetVerifier();
+  Verifier.Match(Source, 1, 1);
+}
+
+TEST(CommentHandlerTest, PPDirectives) {
+  CommentHandlerVisitor Visitor;
+  EXPECT_TRUE(Visitor.runOver(
+        "#warning Y   // ignored_1\n" // #warning takes whole line as message
+        "#undef MACRO // visible_2\n"
+        "#line 1      // visible_3\n"));
+
+  CommentVerifier Verifier = Visitor.GetVerifier();
+  Verifier.Match("// visible_2", 2, 14);
+  Verifier.Match("// visible_3", 3, 14);
+}
+
+} // end namespace clang
diff --git a/unittests/Tooling/RecursiveASTVisitorTest.cpp b/unittests/Tooling/RecursiveASTVisitorTest.cpp
index 9ef86b5..f3ba646 100644
--- a/unittests/Tooling/RecursiveASTVisitorTest.cpp
+++ b/unittests/Tooling/RecursiveASTVisitorTest.cpp
@@ -7,125 +7,10 @@
 //
 //===----------------------------------------------------------------------===//
 
-#include "clang/AST/ASTConsumer.h"
-#include "clang/AST/RecursiveASTVisitor.h"
-#include "clang/Frontend/FrontendAction.h"
-#include "clang/Frontend/CompilerInstance.h"
-#include "clang/Tooling/Tooling.h"
-#include "gtest/gtest.h"
+#include "TestVisitor.h"
 
 namespace clang {
 
-/// \brief Base class for sipmle RecursiveASTVisitor based tests.
-///
-/// This is a drop-in replacement for RecursiveASTVisitor itself, with the
-/// additional capability of running it over a snippet of code.
-///
-/// Visits template instantiations by default.
-///
-/// FIXME: Put into a common location.
-template <typename T>
-class TestVisitor : public clang::RecursiveASTVisitor<T> {
-public:
-  /// \brief Runs the current AST visitor over the given code.
-  bool runOver(StringRef Code) {
-    return tooling::runToolOnCode(new TestAction(this), Code);
-  }
-
-  bool shouldVisitTemplateInstantiations() const {
-    return true;
-  }
-
-protected:
-  clang::ASTContext *Context;
-
-private:
-  class FindConsumer : public clang::ASTConsumer {
-  public:
-    FindConsumer(TestVisitor *Visitor) : Visitor(Visitor) {}
-
-    virtual void HandleTranslationUnit(clang::ASTContext &Context) {
-      Visitor->TraverseDecl(Context.getTranslationUnitDecl());
-    }
-
-  private:
-    TestVisitor *Visitor;
-  };
-
-  class TestAction : public clang::ASTFrontendAction {
-  public:
-    TestAction(TestVisitor *Visitor) : Visitor(Visitor) {}
-
-    virtual clang::ASTConsumer* CreateASTConsumer(
-        clang::CompilerInstance& compiler, llvm::StringRef dummy) {
-      Visitor->Context = &compiler.getASTContext();
-      /// TestConsumer will be deleted by the framework calling us.
-      return new FindConsumer(Visitor);
-    }
-
-  private:
-    TestVisitor *Visitor;
-  };
-};
-
-/// \brief A RecursiveASTVisitor for testing the RecursiveASTVisitor itself.
-///
-/// Allows simple creation of test visitors running matches on only a small
-/// subset of the Visit* methods.
-template <typename T>
-class ExpectedLocationVisitor : public TestVisitor<T> {
-public:
-  ExpectedLocationVisitor()
-    : ExpectedLine(0), ExpectedColumn(0), Found(false) {}
-
-  ~ExpectedLocationVisitor() {
-    EXPECT_TRUE(Found)
-      << "Expected \"" << ExpectedMatch << "\" at " << ExpectedLine
-      << ":" << ExpectedColumn << PartialMatches;
-  }
-
-  /// \brief Expect 'Match' to occur at the given 'Line' and 'Column'.
-  void ExpectMatch(Twine Match, unsigned Line, unsigned Column) {
-    ExpectedMatch = Match.str();
-    ExpectedLine = Line;
-    ExpectedColumn = Column;
-  }
-
-protected:
-  /// \brief Convenience method to simplify writing test visitors.
-  ///
-  /// Sets 'Found' to true if 'Name' and 'Location' match the expected
-  /// values. If only a partial match is found, record the information
-  /// to produce nice error output when a test fails.
-  ///
-  /// Implementations are required to call this with appropriate values
-  /// for 'Name' during visitation.
-  void Match(StringRef Name, SourceLocation Location) {
-    FullSourceLoc FullLocation = this->Context->getFullLoc(Location);
-    if (Name == ExpectedMatch &&
-        FullLocation.isValid() &&
-        FullLocation.getSpellingLineNumber() == ExpectedLine &&
-        FullLocation.getSpellingColumnNumber() == ExpectedColumn) {
-      EXPECT_TRUE(!Found);
-      Found = true;
-    } else if (Name == ExpectedMatch ||
-               (FullLocation.isValid() &&
-                FullLocation.getSpellingLineNumber() == ExpectedLine &&
-                FullLocation.getSpellingColumnNumber() == ExpectedColumn)) {
-      // If we did not match, record information about partial matches.
-      llvm::raw_string_ostream Stream(PartialMatches);
-      Stream << ", partial match: \"" << Name << "\" at ";
-      Location.print(Stream, this->Context->getSourceManager());
-    }
-  }
-
-  std::string ExpectedMatch;
-  unsigned ExpectedLine;
-  unsigned ExpectedColumn;
-  std::string PartialMatches;
-  bool Found;
-};
-
 class TypeLocVisitor : public ExpectedLocationVisitor<TypeLocVisitor> {
 public:
   bool VisitTypeLoc(TypeLoc TypeLocation) {
diff --git a/unittests/Tooling/TestVisitor.h b/unittests/Tooling/TestVisitor.h
new file mode 100644
index 0000000..35098d9
--- /dev/null
+++ b/unittests/Tooling/TestVisitor.h
@@ -0,0 +1,141 @@
+//===--- TestVisitor.h ------------------------------------------*- C++ -*-===//
+//
+//                     The LLVM Compiler Infrastructure
+//
+// This file is distributed under the University of Illinois Open Source
+// License. See LICENSE.TXT for details.
+//
+//===----------------------------------------------------------------------===//
+//
+//  This file defines a utility class for RecursiveASTVisitor related tests.
+//
+//===----------------------------------------------------------------------===//
+
+#ifndef LLVM_CLANG_TEST_VISITOR_H
+#define LLVM_CLANG_TEST_VISITOR_H
+
+#include "clang/AST/ASTConsumer.h"
+#include "clang/AST/RecursiveASTVisitor.h"
+#include "clang/Frontend/FrontendAction.h"
+#include "clang/Frontend/CompilerInstance.h"
+#include "clang/Tooling/Tooling.h"
+#include "gtest/gtest.h"
+
+namespace clang {
+
+/// \brief Base class for sipmle RecursiveASTVisitor based tests.
+///
+/// This is a drop-in replacement for RecursiveASTVisitor itself, with the
+/// additional capability of running it over a snippet of code.
+///
+/// Visits template instantiations by default.
+template <typename T>
+class TestVisitor : public RecursiveASTVisitor<T> {
+public:
+  TestVisitor() { }
+
+  /// \brief Runs the current AST visitor over the given code.
+  bool runOver(StringRef Code) {
+    return tooling::runToolOnCode(CreateTestAction(), Code);
+  }
+
+  bool shouldVisitTemplateInstantiations() const {
+    return true;
+  }
+
+protected:
+  virtual ASTFrontendAction* CreateTestAction() {
+    return new TestAction(this);
+  }
+
+  class FindConsumer : public ASTConsumer {
+  public:
+    FindConsumer(TestVisitor *Visitor) : Visitor(Visitor) {}
+
+    virtual void HandleTranslationUnit(clang::ASTContext &Context) {
+      Visitor->TraverseDecl(Context.getTranslationUnitDecl());
+    }
+
+  private:
+    TestVisitor *Visitor;
+  };
+
+  class TestAction : public ASTFrontendAction {
+  public:
+    TestAction(TestVisitor *Visitor) : Visitor(Visitor) {}
+
+    virtual clang::ASTConsumer* CreateASTConsumer(
+        CompilerInstance& compiler, llvm::StringRef dummy) {
+      Visitor->Context = &compiler.getASTContext();
+      /// TestConsumer will be deleted by the framework calling us.
+      return new FindConsumer(Visitor);
+    }
+
+  protected:
+    TestVisitor *Visitor;
+  };
+
+  ASTContext *Context;
+};
+
+
+/// \brief A RecursiveASTVisitor for testing the RecursiveASTVisitor itself.
+///
+/// Allows simple creation of test visitors running matches on only a small
+/// subset of the Visit* methods.
+template <typename T, template <typename> class Visitor = TestVisitor>
+class ExpectedLocationVisitor : public Visitor<T> {
+public:
+  ExpectedLocationVisitor()
+    : ExpectedLine(0), ExpectedColumn(0), Found(false) {}
+
+  ~ExpectedLocationVisitor() {
+    EXPECT_TRUE(Found)
+      << "Expected \"" << ExpectedMatch << "\" at " << ExpectedLine
+      << ":" << ExpectedColumn << PartialMatches;
+  }
+
+  /// \brief Expect 'Match' to occur at the given 'Line' and 'Column'.
+  void ExpectMatch(Twine Match, unsigned Line, unsigned Column) {
+    ExpectedMatch = Match.str();
+    ExpectedLine = Line;
+    ExpectedColumn = Column;
+  }
+
+protected:
+  /// \brief Convenience method to simplify writing test visitors.
+  ///
+  /// Sets 'Found' to true if 'Name' and 'Location' match the expected
+  /// values. If only a partial match is found, record the information
+  /// to produce nice error output when a test fails.
+  ///
+  /// Implementations are required to call this with appropriate values
+  /// for 'Name' during visitation.
+  void Match(StringRef Name, SourceLocation Location) {
+    FullSourceLoc FullLocation = this->Context->getFullLoc(Location);
+    if (Name == ExpectedMatch &&
+        FullLocation.isValid() &&
+        FullLocation.getSpellingLineNumber() == ExpectedLine &&
+        FullLocation.getSpellingColumnNumber() == ExpectedColumn) {
+      EXPECT_TRUE(!Found);
+      Found = true;
+    } else if (Name == ExpectedMatch ||
+               (FullLocation.isValid() &&
+                FullLocation.getSpellingLineNumber() == ExpectedLine &&
+                FullLocation.getSpellingColumnNumber() == ExpectedColumn)) {
+      // If we did not match, record information about partial matches.
+      llvm::raw_string_ostream Stream(PartialMatches);
+      Stream << ", partial match: \"" << Name << "\" at ";
+      Location.print(Stream, this->Context->getSourceManager());
+    }
+  }
+
+  std::string ExpectedMatch;
+  unsigned ExpectedLine;
+  unsigned ExpectedColumn;
+  std::string PartialMatches;
+  bool Found;
+};
+}
+
+#endif /* LLVM_CLANG_TEST_VISITOR_H */