//------------------------------------------------------------------------------
// AST matching sample. Demonstrates:
//
// * How to write a simple source tool using libTooling.
// * How to use AST matchers to find interesting AST nodes.
// * How to use the Rewriter API to rewrite the source code.
//
// Eli Bendersky (eliben@gmail.com)
// This code is in the public domain
//------------------------------------------------------------------------------
#include <iostream>
#include <string>

#include <clang/AST/AST.h>
#include <clang/AST/ASTConsumer.h>
#include <clang/ASTMatchers/ASTMatchFinder.h>
#include <clang/ASTMatchers/ASTMatchers.h>
#include <clang/Frontend/CompilerInstance.h>
#include <clang/Frontend/FrontendActions.h>
#include <clang/Rewrite/Core/Rewriter.h>
#include <clang/Tooling/CommonOptionsParser.h>
#include <clang/Tooling/Tooling.h>
#include <llvm/Support/raw_ostream.h>

using namespace clang;
using namespace clang::ast_matchers;
using namespace clang::driver;
using namespace clang::tooling;

static llvm::cl::OptionCategory MatcherSampleCategory("Matcher Sample");

class MyMatchHandler : public MatchFinder::MatchCallback
{
public:
	MyMatchHandler(Rewriter& rewriter);

	virtual void run(const MatchFinder::MatchResult& result);

private:
	Rewriter& rewriter_;
};

MyMatchHandler::MyMatchHandler(Rewriter& rewriter)
	: rewriter_(rewriter)
{
}

void MyMatchHandler::run(const MatchFinder::MatchResult& result)
{
	std::cout << "MyMatchHandler::run()\n";

	const FunctionDecl* caller(result.Nodes.getNodeAs<FunctionDecl>("caller"));
	const FunctionDecl* callee(result.Nodes.getNodeAs<FunctionDecl>("callee"));

#if 0
	const CallExpr* expr(result.Nodes.getNodeAs<CallExpr>("moo"));

	std::cout << expr->getCallee() << "\n";
#endif

	const CXXMemberCallExpr* stmt(result.Nodes.getNodeAs<CXXMemberCallExpr>("moo"));

	stmt->getMethodDecl()->getParent()->dump();

	// XXX:
	// Verify that callee is a method in the base class of caller's class.

	if (caller->getNameInfo().getName() == callee->getNameInfo().getName())
	{
		rewriter_.InsertText(caller->getLocStart(),
				             "/* XXX USELESS */\n", true, true);
	}
}

// Implementation of the ASTConsumer interface for reading an AST produced
// by the Clang parser. It registers a couple of matchers and runs them on
// the AST.
class MyASTConsumer : public ASTConsumer
{
public:
	MyASTConsumer(Rewriter& r)
	    : ASTConsumer()
	    , my_handler_(r)
	{
		matcher_.addMatcher(
			functionDecl(
				hasBody(
					compoundStmt(
						statementCountIs(1)
						, hasDescendant(
							cxxMemberCallExpr(
								callee(
									functionDecl(
									).bind("callee")
								)
							).bind("moo")
						)
					)
				)
			).bind("caller"),
			&my_handler_
		);
	}

	void HandleTranslationUnit(ASTContext& context) override
	{
		// Run the matchers when we have the whole TU parsed.
		matcher_.matchAST(context);
	}

private:
	MatchFinder matcher_;

	MyMatchHandler my_handler_;
};

// For each source file provided to the tool, a new FrontendAction is created.
class MyFrontendAction : public ASTFrontendAction {
public:
	MyFrontendAction()
	{
	}

	void EndSourceFileAction() override {
		rewriter_.getEditBuffer(rewriter_.getSourceMgr().getMainFileID())
			.write(llvm::outs());
	}

	std::unique_ptr<ASTConsumer> CreateASTConsumer(CompilerInstance& CI, StringRef file) override {
		rewriter_.setSourceMgr(CI.getSourceManager(), CI.getLangOpts());

		return llvm::make_unique<MyASTConsumer>(rewriter_);
	}

private:
	Rewriter rewriter_;
};

int main(int argc, const char** argv)
{
	CommonOptionsParser op(argc, argv, MatcherSampleCategory);
	ClangTool Tool(op.getCompilations(), op.getSourcePathList());

	return Tool.run(newFrontendActionFactory<MyFrontendAction>().get());
}