Refactor codes and tests to make it work in PCH builds

[psionic12]

Summary

The former codes won't work for cases which PCH involves 
in a build workflow, this patch fixes this. Besides, add
some tips on how to deal with PCH when traversing nodes.

clang/examples/CallSuperAttribute/CallSuperAttrInfo.cpp

 //===- AnnotateFunctions.cpp ----------------------------------------------===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 //
 // Attribute plugin to mark a virtual method as ``call_super``, subclasses must
 // call it in the overridden method.
 //
 // This example shows that attribute plugins combined with ``PluginASTAction``
 // in Clang can do some of the same things which Java Annotations do.
 //
-// Unlike the other attribute plugin examples, this one does not attach an
-// attribute AST node to the declaration AST node. Instead, it keeps a separate
-// list of attributed declarations, which may be faster than using
-// ``Decl::getAttr<T>()`` in some cases. The disadvantage of this approach is
-// that the attribute is not part of the AST, which means that dumping the AST
-// will lose the attribute information, pretty printing the AST won't write the
-// attribute back out to source, and AST matchers will not be able to match
-// against the attribute on the declaration.
-//
+// Notice the part of traversing decls, it only focuses on nodes which belong to
+// the current parsing file, all other nodes are not loaded, this is especially
+// useful to improve the performance when using PCH, or the upcoming module
+// feature in C++20.
 //===----------------------------------------------------------------------===//
 
 #include "clang/AST/ASTContext.h"
 #include "clang/AST/Attr.h"
 #include "clang/AST/RecursiveASTVisitor.h"
 #include "clang/Frontend/FrontendPluginRegistry.h"
 #include "clang/Sema/ParsedAttr.h"
 #include "clang/Sema/Sema.h"
 #include "clang/Sema/SemaDiagnostic.h"
 #include "llvm/ADT/SmallPtrSet.h"
 using namespace clang;
 
 namespace {
-// Cached methods which are marked as 'call_super'.
-llvm::SmallPtrSet<const CXXMethodDecl *, 16> MarkedMethods;
+constexpr char AnnotateStr[] = "call_super";
 bool isMarkedAsCallSuper(const CXXMethodDecl *D) {
-  // Uses this way to avoid add an annotation attr to the AST.
-  return MarkedMethods.contains(D);
+  for (auto It = D->specific_attr_begin<AnnotateAttr>();
+       It != D->specific_attr_end<AnnotateAttr>(); It++) {
+    if (AnnotateStr == It->getAnnotation())
+      return true;
+  }
+  return false;
 }
 
 class MethodUsageVisitor : public RecursiveASTVisitor<MethodUsageVisitor> {
 public:
   bool IsOverriddenUsed = false;
   explicit MethodUsageVisitor(
       llvm::SmallPtrSet<const CXXMethodDecl *, 16> &MustCalledMethods)
       : MustCalledMethods(MustCalledMethods) {}
CallSuperVisitor(DiagnosticsEngine &Diags) : Diags(Diags) {
     WarningSuperNotCalled = Diags.getCustomDiagID(
         DiagnosticsEngine::Warning,
         "virtual function %q0 is marked as 'call_super' but this overriding "
         "method does not call the base version");
     NotePreviousCallSuperDeclaration = Diags.getCustomDiagID(
         DiagnosticsEngine::Note, "function marked 'call_super' here");
   }
   bool VisitCXXMethodDecl(CXXMethodDecl *MethodDecl) {
+    if (isMarkedAsCallSuper(MethodDecl) &&
+        MethodDecl->isFirstDecl() /* avoid warning multiple times */
+    ) {
+      lateDiagAppertainsToDecl(MethodDecl);
+    }
+
     if (MethodDecl->isThisDeclarationADefinition() && MethodDecl->hasBody()) {
       // First find out which overridden methods are marked as 'call_super'
       llvm::SmallPtrSet<const CXXMethodDecl *, 16> OverriddenMarkedMethods;
       for (const auto *Overridden : MethodDecl->overridden_methods()) {
         if (isMarkedAsCallSuper(Overridden)) {
           OverriddenMarkedMethods.insert(Overridden);
         }
       }
bool VisitCXXMethodDecl(CXXMethodDecl *MethodDecl) {
     }
     return true;
   }
 
 private:
   DiagnosticsEngine &Diags;
   unsigned WarningSuperNotCalled;
   unsigned NotePreviousCallSuperDeclaration;
-};
-
-class CallSuperConsumer : public ASTConsumer {
-public:
-  void HandleTranslationUnit(ASTContext &Context) override {
-    auto &Diags = Context.getDiagnostics();
-    for (const auto *Method : MarkedMethods) {
-      lateDiagAppertainsToDecl(Diags, Method);
-    }
-
-    CallSuperVisitor Visitor(Context.getDiagnostics());
-    Visitor.TraverseDecl(Context.getTranslationUnitDecl());
-  }
-
-private:
   // This function does checks which cannot be done in `diagAppertainsToDecl()`,
   // typical example is checking Attributes (such as `FinalAttr`), on the time
   // when `diagAppertainsToDecl()` is called, `FinalAttr` is not added into
   // the AST yet.
-  void lateDiagAppertainsToDecl(DiagnosticsEngine &Diags,
-                                const CXXMethodDecl *MethodDecl) {
+  void lateDiagAppertainsToDecl(const CXXMethodDecl *MethodDecl) {
     if (MethodDecl->hasAttr<FinalAttr>()) {
       unsigned ID = Diags.getCustomDiagID(
           DiagnosticsEngine::Warning,
           "'call_super' attribute marked on a final method");
       Diags.Report(MethodDecl->getLocation(), ID);
     }
   }
 };
 
+class CallSuperConsumer : public ASTConsumer {
+public:
+  void HandleTranslationUnit(ASTContext &Context) override {
+    CallSuperVisitor Visitor(Context.getDiagnostics());
+    // Do not load external sources (usually a PCH), to speed up the traverse.
+    auto *TU = Context.getTranslationUnitDecl();
+    std::vector<Decl *> V;
+    for (auto *D : TU->noload_decls()) {
+      V.push_back(D);
+    }
+    Context.setTraversalScope(V);
+    Visitor.TraverseAST(Context);
+  }
+};
+
 class CallSuperAction : public PluginASTAction {
 public:
   std::unique_ptr<ASTConsumer> CreateASTConsumer(CompilerInstance &CI,
                                                  llvm::StringRef) override {
     return std::make_unique<CallSuperConsumer>();
   }
 
   bool ParseArgs(const CompilerInstance &CI,
struct CallSuperAttrInfo : public ParsedAttrInfo {
   bool diagAppertainsToDecl(Sema &S, const ParsedAttr &Attr,
                             const Decl *D) const override {
     const auto *TheMethod = dyn_cast_or_null<CXXMethodDecl>(D);
     if (!TheMethod || !TheMethod->isVirtual()) {
       S.Diag(Attr.getLoc(), diag::warn_attribute_wrong_decl_type_str)
           << Attr << "virtual functions";
       return false;
     }
-    MarkedMethods.insert(TheMethod);
     return true;
   }
   AttrHandling handleDeclAttribute(Sema &S, Decl *D,
                                    const ParsedAttr &Attr) const override {
-    // No need to add an attr object (usually an `AnnotateAttr` is added).
-    // Save the address of the Decl in a set, it maybe faster than compare to
-    // strings.
+    D->addAttr(AnnotateAttr::Create(S.Context, AnnotateStr, nullptr, 0,
+                                    Attr.getRange()));
     return AttributeNotApplied;
   }
 };
 
 } // namespace
 static FrontendPluginRegistry::Add<CallSuperAction>
     X("call_super_plugin", "clang plugin, checks every overridden virtual "
                            "function whether called this function or not.");
 static ParsedAttrInfoRegistry::Add<CallSuperAttrInfo>
     Y("call_super_attr", "Attr plugin to define 'call_super' attribute");

clang/test/Frontend/plugin-call-super-pch.cpp

+// RUN: split-file %s %t
+// RUN: %clang -cc1 -verify=verify-pch -x c++-header %t/call_super_test.h -emit-pch -o %t/call_super_test.h.pch -load %llvmshlibdir/CallSuperAttr%pluginext
+// RUN: %clang -cc1 -verify=verify-source -fsyntax-only -include-pch %t/call_super_test.h.pch %t/call_super_test.cpp  -load %llvmshlibdir/CallSuperAttr%pluginext
+//--- call_super_test.h
+struct Foo {
+  [[clang::call_super]]
+  virtual void test() {}
+  [[clang::call_super]]
+  virtual void test2() final {} // verify-pch-warning {{'call_super' attribute marked on a final method}}
+};
+
+struct Bar : public Foo {
+  void test() override final;
+};
+
+//--- call_super_test.cpp
+void Bar::test() { // verify-source-warning {{virtual function 'Foo::test' is marked as 'call_super' but this overriding method does not call the base version}}
+// verify-source-note@call_super_test.h:7 {{function marked 'call_super' here}}
+}
 No newline at end of file