#include "clang/Frontend/FrontendPluginRegistry.h"
#include "clang/AST/ASTConsumer.h"
#include "clang/AST/AST.h"
#include "clang/Frontend/CompilerInstance.h"
#include "llvm/Support/raw_ostream.h"
using namespace clang;

namespace {

class OverrideConsumer : public ASTConsumer {
public:
  virtual void HandleTagDeclDefinition(TagDecl* tag) {
    if (CXXRecordDecl* record = dyn_cast<CXXRecordDecl>(tag)) {
      for (CXXRecordDecl::method_iterator method = record->method_begin();
           method != record->method_end(); ++method) {
        if (method->isVirtual()) {
          bool has_override_attr = method->getAttr<OverrideAttr>();
          bool is_override = method->begin_overridden_methods() != method->end_overridden_methods();

          if (is_override && !has_override_attr) {
            Diagnostic& D = method->getASTContext().getDiagnostics();
            unsigned id = D.getCustomDiagID(Diagnostic::Warning,
                "method overrides parent class; consider adding OVERRIDE");
            D.Report(method->getASTContext().getFullLoc(method->getLocation()), id);
          } else if (has_override_attr && !is_override) {
            Diagnostic& D = method->getASTContext().getDiagnostics();
            unsigned id = D.getCustomDiagID(Diagnostic::Error,
                "method claims OVERRIDE but fails to override any superclass method");
            D.Report(method->getASTContext().getFullLoc(method->getLocation()), id);
          }
        }
      }
    }
  }
};

class CheckOverrideAction : public PluginASTAction {
protected:
  ASTConsumer *CreateASTConsumer(CompilerInstance &CI, llvm::StringRef) {
    return new OverrideConsumer();
  }

  bool ParseArgs(const CompilerInstance &CI,
                 const std::vector<std::string>& args) {
    return true;
  }
};

}

static FrontendPluginRegistry::Add<CheckOverrideAction>
X("check-overrides", "check usage of __attribute__((override))");