#include <iostream>

#include "clang/AST/ASTConsumer.h"
#include "clang/AST/RecursiveASTVisitor.h"
#include "clang/Frontend/CompilerInstance.h"
#include "clang/Frontend/FrontendAction.h"
#include "clang/Tooling/Tooling.h"

using namespace clang;

class FindNamedClassVisitor : public RecursiveASTVisitor<FindNamedClassVisitor>
{
public:
    explicit FindNamedClassVisitor(ASTContext *) {}

    bool VisitCXXRecordDecl(CXXRecordDecl *Decl)
    {
        if (Decl->getQualifiedNameAsString() != "A")
            return true;

        // For each method of A
        for (const auto &M : Decl->methods())
        {
            assert(M->getAccess() == AS_public);

            // Check the visibility of the return type
            const auto *C = M->getReturnType()->getAsCXXRecordDecl();
            assert(C != nullptr);

            std::cout << C->getNameAsString()
                      << ": Acc=" << AccToString(C->getAccess())
                      << "; Visible=" << IsVisible(C) << std::endl;

            // just to be sure, check also the canonical decl
            C = C->getCanonicalDecl();
            std::cout << C->getNameAsString()
                      << ": Acc=" << AccToString(C->getAccess())
                      << "; Visible=" << IsVisible(C) << std::endl;
        }
        return true;
    }

private:
    bool IsVisible(const CXXRecordDecl *R) const
    {
        assert(R != nullptr);
        // R is not public? return false
        if (R->getAccess() != AS_public && R->getAccess() != AS_none)
            return false;

        // if R is public (or AS_none), check that also the parent is public
        const auto *Parent = R->getParent();
        if (llvm::isa<CXXRecordDecl>(Parent))
            return IsVisible(llvm::cast<CXXRecordDecl>(Parent));
        return true;
    }

    // return a printable version of an access specifier
    std::string AccToString(AccessSpecifier A) const
    {
        switch (A)
        {
            case AccessSpecifier::AS_none:
                return "AS_none";
            case AccessSpecifier::AS_public:
                return "AS_public";
            case AccessSpecifier::AS_private:
                return "AS_private";
            case AccessSpecifier::AS_protected:
                return "AS_protected";
        }
    }
};

class FindNamedClassConsumer : public clang::ASTConsumer
{
public:
    explicit FindNamedClassConsumer(ASTContext *Context) : Visitor(Context) {}

    virtual void HandleTranslationUnit(clang::ASTContext &Context)
    {
        Visitor.TraverseDecl(Context.getTranslationUnitDecl());
    }

private:
    FindNamedClassVisitor Visitor;
};

class FindNamedClassAction : public clang::ASTFrontendAction
{
public:
    virtual std::unique_ptr<clang::ASTConsumer> CreateASTConsumer(
        clang::CompilerInstance &Compiler, llvm::StringRef InFile)
    {
        return std::unique_ptr<clang::ASTConsumer>(
            new FindNamedClassConsumer(&Compiler.getASTContext()));
    }
};

int main(int argc, char **argv)
{
    // Demo code
    const char code[] = R"__(
class A {
private:
  class B { public: class B2 {}; };
  template<typename T> class C {};

public:
  class D {};

  // ---
  B getB();
  B::B2 getB2();
  C<int> getC();
  D getD();
};

)__";
    clang::tooling::runToolOnCode(new FindNamedClassAction(), code);
}