MainReturns.cpp

//===---- tools/extra/MainReturns.cpp -------------------------------------===//
//
//                     The LLVM Compiler Infrastructure
//
// This file is distributed under the University of Illinois Open Source
// License. See LICENSE.TXT for details.
//
//===----------------------------------------------------------------------===//

#include "clang/ASTMatchers/ASTMatchFinder.h"
#include "clang/ASTMatchers/ASTMatchers.h"
#include "clang/Basic/SourceManager.h"
#include "clang/Frontend/FrontendActions.h"
#include "clang/Lex/Lexer.h"
#include "clang/Rewrite/Core/Rewriter.h"
#include "clang/Tooling/CommonOptionsParser.h"
#include "clang/Tooling/Execution.h"
#include "clang/Tooling/Refactoring.h"
#include "clang/Tooling/Refactoring/AtomicChange.h"
#include "clang/Tooling/Tooling.h"
#include "llvm/Support/CommandLine.h"
#include "llvm/Support/MemoryBuffer.h"
#include "llvm/Support/Signals.h"

using namespace clang;
using namespace clang::ast_matchers;
using namespace clang::tooling;
using namespace llvm;

static llvm::cl::OptionCategory MainReturnsCategory("Main should always return");
static cl::extrahelp CommonHelp(CommonOptionsParser::HelpMessage);
static cl::extrahelp MoreHelp("\nAdd return to main when it doesn't have one.\n");

auto MainWithoutReturn =
  functionDecl(allOf(
                     hasName("main"),
                     unless(hasDescendant(returnStmt()))
                     )).bind("badMain");

class MainRewriter : public MatchFinder::MatchCallback {
public :
  void run(const MatchFinder::MatchResult &Result) override {
    auto *D = Result.Nodes.getNodeAs<FunctionDecl>("badMain");
    if (D->getBeginLoc().isValid()) {
      unsigned Length = 1;
      SourceManager& source = *Result.SourceManager;
      Replacement replacement(source, D->getEndLoc(), Length, "\n  return 0;\n}");
      Rewriter rewriter(source, LangOptions());
      bool replacementSucceeded = replacement.apply(rewriter);
      if (!replacementSucceeded)
        llvm::errs() << "Failed applying replacement " << replacement.toString() << "\n";
      rewriter.overwriteChangedFiles();
    }
  }
};

int main(int argc, const char **argv) {
  llvm::sys::PrintStackTraceOnErrorSignal(argv[0]);
  auto Executor = clang::tooling::createExecutorFromCommandLineArgs(argc, argv, MainReturnsCategory);
  if (!Executor) {
    llvm::errs() << llvm::toString(Executor.takeError()) << "\n";
    return 1;
  }

  ast_matchers::MatchFinder Finder;
  MainRewriter mainRewriter;
  Finder.addMatcher(MainWithoutReturn, &mainRewriter);
  auto Err = Executor->get()->execute(newFrontendActionFactory(&Finder));
  if (Err) {
    llvm::errs() << llvm::toString(std::move(Err)) << "\n";
  }
  Executor->get()->getToolResults()->forEachResult(
      [](llvm::StringRef key, llvm::StringRef value) {
        llvm::errs() << "----" << key.str() << "\n" << value.str() << "\n";
      });
}