//===--- USRFindingAction.cpp - Clang refactoring library -----------------===//
//
// 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
//
//===----------------------------------------------------------------------===//
///
/// \file
/// Provides an action to find USR for the symbol at <offset>, as well as
/// all additional USRs.
///
//===----------------------------------------------------------------------===//

#include "clang/Tooling/Refactoring/Rename/USRFindingAction.h"
#include "clang/AST/AST.h"
#include "clang/AST/ASTConsumer.h"
#include "clang/AST/ASTContext.h"
#include "clang/AST/Decl.h"
#include "clang/AST/RecursiveASTVisitor.h"
#include "clang/Basic/FileManager.h"
#include "clang/Frontend/CompilerInstance.h"
#include "clang/Frontend/FrontendAction.h"
#include "clang/Lex/Lexer.h"
#include "clang/Lex/Preprocessor.h"
#include "clang/Tooling/CommonOptionsParser.h"
#include "clang/Tooling/Refactoring.h"
#include "clang/Tooling/Refactoring/Rename/USRFinder.h"
#include "clang/Tooling/Tooling.h"

#include <algorithm>
#include <set>
#include <string>
#include <vector>

using namespace llvm;

namespace clang {
namespace tooling {

const NamedDecl *getCanonicalSymbolDeclaration(const NamedDecl *FoundDecl) {
  if (!FoundDecl)
    return nullptr;
  // If FoundDecl is a constructor or destructor, we want to instead take
  // the Decl of the corresponding class.
  if (const auto *CtorDecl = dyn_cast<CXXConstructorDecl>(FoundDecl))
    FoundDecl = CtorDecl->getParent();
  else if (const auto *DtorDecl = dyn_cast<CXXDestructorDecl>(FoundDecl))
    FoundDecl = DtorDecl->getParent();
  // FIXME: (Alex L): Canonicalize implicit template instantions, just like
  // the indexer does it.

  // Note: please update the declaration's doc comment every time the
  // canonicalization rules are changed.
  return FoundDecl;
}

namespace {
// NamedDeclFindingConsumer should delegate finding USRs of given Decl to
// AdditionalUSRFinder. AdditionalUSRFinder adds USRs of ctor and dtor if given
// Decl refers to class and adds USRs of all overridden methods if Decl refers
// to virtual method.
class AdditionalUSRFinder : public RecursiveASTVisitor<AdditionalUSRFinder> {
public:
  AdditionalUSRFinder(const Decl *FoundDecl, ASTContext &Context)
      : FoundDecl(FoundDecl), Context(Context) {}

  std::vector<std::string> Find() {
    // Fill OverriddenMethods and PartialSpecs storages.
    TraverseAST(Context);
    if (const auto *MethodDecl = dyn_cast<CXXMethodDecl>(FoundDecl)) {
      addUSRsOfOverridenFunctions(MethodDecl);
      for (const auto &OverriddenMethod : OverriddenMethods) {
        if (checkIfOverriddenFunctionAscends(OverriddenMethod))
          USRSet.insert(getUSRForDecl(OverriddenMethod));
      }
      addUSRsOfInstantiatedMethods(MethodDecl);
    } else if (const auto *RecordDecl = dyn_cast<CXXRecordDecl>(FoundDecl)) {
      handleCXXRecordDecl(RecordDecl);
    } else if (const auto *TemplateDecl =
                   dyn_cast<ClassTemplateDecl>(FoundDecl)) {
      handleClassTemplateDecl(TemplateDecl);
    } else if (const auto *FD = dyn_cast<FunctionDecl>(FoundDecl)) {
      USRSet.insert(getUSRForDecl(FD));
      if (const auto *FTD = FD->getPrimaryTemplate())
        handleFunctionTemplateDecl(FTD);
    } else if (const auto *FD = dyn_cast<FunctionTemplateDecl>(FoundDecl)) {
      handleFunctionTemplateDecl(FD);
    } else if (const auto *VTD = dyn_cast<VarTemplateDecl>(FoundDecl)) {
      handleVarTemplateDecl(VTD);
    } else if (const auto *VD =
                   dyn_cast<VarTemplateSpecializationDecl>(FoundDecl)) {
      // FIXME: figure out why FoundDecl can be a VarTemplateSpecializationDecl.
      handleVarTemplateDecl(VD->getSpecializedTemplate());
    } else if (const auto *VD = dyn_cast<VarDecl>(FoundDecl)) {
      USRSet.insert(getUSRForDecl(VD));
      if (const auto *VTD = VD->getDescribedVarTemplate())
        handleVarTemplateDecl(VTD);
    } else {
      USRSet.insert(getUSRForDecl(FoundDecl));
    }
    return std::vector<std::string>(USRSet.begin(), USRSet.end());
  }

  bool shouldVisitTemplateInstantiations() const { return true; }

  bool VisitCXXMethodDecl(const CXXMethodDecl *MethodDecl) {
    if (MethodDecl->isVirtual())
      OverriddenMethods.push_back(MethodDecl);
    if (MethodDecl->getInstantiatedFromMemberFunction())
      InstantiatedMethods.push_back(MethodDecl);
    return true;
  }

private:
  void handleCXXRecordDecl(const CXXRecordDecl *RecordDecl) {
    if (!RecordDecl->getDefinition()) {
      USRSet.insert(getUSRForDecl(RecordDecl));
      return;
    }
    RecordDecl = RecordDecl->getDefinition();
    if (const auto *ClassTemplateSpecDecl =
            dyn_cast<ClassTemplateSpecializationDecl>(RecordDecl))
      handleClassTemplateDecl(ClassTemplateSpecDecl->getSpecializedTemplate());
    addUSRsOfCtorDtors(RecordDecl);
  }

  void handleClassTemplateDecl(const ClassTemplateDecl *TemplateDecl) {
    for (const auto *Specialization : TemplateDecl->specializations())
      addUSRsOfCtorDtors(Specialization);
    SmallVector<ClassTemplatePartialSpecializationDecl *, 4> PartialSpecs;
    TemplateDecl->getPartialSpecializations(PartialSpecs);
    for (const auto *Spec : PartialSpecs)
      addUSRsOfCtorDtors(Spec);
    addUSRsOfCtorDtors(TemplateDecl->getTemplatedDecl());
  }

  void handleFunctionTemplateDecl(const FunctionTemplateDecl *FTD) {
    USRSet.insert(getUSRForDecl(FTD));
    USRSet.insert(getUSRForDecl(FTD->getTemplatedDecl()));
    for (const auto *S : FTD->specializations())
      USRSet.insert(getUSRForDecl(S));
  }

  void handleVarTemplateDecl(const VarTemplateDecl *VTD) {
    USRSet.insert(getUSRForDecl(VTD));
    USRSet.insert(getUSRForDecl(VTD->getTemplatedDecl()));
    llvm::for_each(VTD->specializations(), [&](const auto *Spec) {
      USRSet.insert(getUSRForDecl(Spec));
    });
    SmallVector<VarTemplatePartialSpecializationDecl *, 4> PartialSpecs;
    VTD->getPartialSpecializations(PartialSpecs);
    llvm::for_each(PartialSpecs, [&](const auto *Spec) {
      USRSet.insert(getUSRForDecl(Spec));
    });
  }

  void addUSRsOfCtorDtors(const CXXRecordDecl *RD) {
    const auto* RecordDecl = RD->getDefinition();

    // Skip if the CXXRecordDecl doesn't have definition.
    if (!RecordDecl) {
      USRSet.insert(getUSRForDecl(RD));
      return;
    }

    for (const auto *CtorDecl : RecordDecl->ctors())
      USRSet.insert(getUSRForDecl(CtorDecl));
    // Add template constructor decls, they are not in ctors() unfortunately.
    if (RecordDecl->hasUserDeclaredConstructor())
      for (const auto *D : RecordDecl->decls())
        if (const auto *FTD = dyn_cast<FunctionTemplateDecl>(D))
          if (const auto *Ctor =
                  dyn_cast<CXXConstructorDecl>(FTD->getTemplatedDecl()))
            USRSet.insert(getUSRForDecl(Ctor));

    USRSet.insert(getUSRForDecl(RecordDecl->getDestructor()));
    USRSet.insert(getUSRForDecl(RecordDecl));
  }

  void addUSRsOfOverridenFunctions(const CXXMethodDecl *MethodDecl) {
    USRSet.insert(getUSRForDecl(MethodDecl));
    // Recursively visit each OverridenMethod.
    for (const auto &OverriddenMethod : MethodDecl->overridden_methods())
      addUSRsOfOverridenFunctions(OverriddenMethod);
  }

  void addUSRsOfInstantiatedMethods(const CXXMethodDecl *MethodDecl) {
    // For renaming a class template method, all references of the instantiated
    // member methods should be renamed too, so add USRs of the instantiated
    // methods to the USR set.
    USRSet.insert(getUSRForDecl(MethodDecl));
    if (const auto *FT = MethodDecl->getInstantiatedFromMemberFunction())
      USRSet.insert(getUSRForDecl(FT));
    for (const auto *Method : InstantiatedMethods) {
      if (USRSet.find(getUSRForDecl(
              Method->getInstantiatedFromMemberFunction())) != USRSet.end())
        USRSet.insert(getUSRForDecl(Method));
    }
  }

  bool checkIfOverriddenFunctionAscends(const CXXMethodDecl *MethodDecl) {
    for (const auto &OverriddenMethod : MethodDecl->overridden_methods()) {
      if (USRSet.find(getUSRForDecl(OverriddenMethod)) != USRSet.end())
        return true;
      return checkIfOverriddenFunctionAscends(OverriddenMethod);
    }
    return false;
  }

  const Decl *FoundDecl;
  ASTContext &Context;
  std::set<std::string> USRSet;
  std::vector<const CXXMethodDecl *> OverriddenMethods;
  std::vector<const CXXMethodDecl *> InstantiatedMethods;
};
} // namespace

std::vector<std::string> getUSRsForDeclaration(const NamedDecl *ND,
                                               ASTContext &Context) {
  AdditionalUSRFinder Finder(ND, Context);
  return Finder.Find();
}

class NamedDeclFindingConsumer : public ASTConsumer {
public:
  NamedDeclFindingConsumer(ArrayRef<unsigned> SymbolOffsets,
                           ArrayRef<std::string> QualifiedNames,
                           std::vector<std::string> &SpellingNames,
                           std::vector<std::vector<std::string>> &USRList,
                           bool Force, bool &ErrorOccurred)
      : SymbolOffsets(SymbolOffsets), QualifiedNames(QualifiedNames),
        SpellingNames(SpellingNames), USRList(USRList), Force(Force),
        ErrorOccurred(ErrorOccurred) {}

private:
  bool FindSymbol(ASTContext &Context, const SourceManager &SourceMgr,
                  unsigned SymbolOffset, const std::string &QualifiedName) {
    DiagnosticsEngine &Engine = Context.getDiagnostics();
    const FileID MainFileID = SourceMgr.getMainFileID();

    if (SymbolOffset >= SourceMgr.getFileIDSize(MainFileID)) {
      ErrorOccurred = true;
      unsigned InvalidOffset = Engine.getCustomDiagID(
          DiagnosticsEngine::Error,
          "SourceLocation in file %0 at offset %1 is invalid");
      Engine.Report(SourceLocation(), InvalidOffset)
          << SourceMgr.getFileEntryForID(MainFileID)->getName() << SymbolOffset;
      return false;
    }

    const SourceLocation Point = SourceMgr.getLocForStartOfFile(MainFileID)
                                     .getLocWithOffset(SymbolOffset);
    const NamedDecl *FoundDecl = QualifiedName.empty()
                                     ? getNamedDeclAt(Context, Point)
                                     : getNamedDeclFor(Context, QualifiedName);

    if (FoundDecl == nullptr) {
      if (QualifiedName.empty()) {
        FullSourceLoc FullLoc(Point, SourceMgr);
        unsigned CouldNotFindSymbolAt = Engine.getCustomDiagID(
            DiagnosticsEngine::Error,
            "clang-rename could not find symbol (offset %0)");
        Engine.Report(Point, CouldNotFindSymbolAt) << SymbolOffset;
        ErrorOccurred = true;
        return false;
      }

      if (Force) {
        SpellingNames.push_back(std::string());
        USRList.push_back(std::vector<std::string>());
        return true;
      }

      unsigned CouldNotFindSymbolNamed = Engine.getCustomDiagID(
          DiagnosticsEngine::Error, "clang-rename could not find symbol %0");
      Engine.Report(CouldNotFindSymbolNamed) << QualifiedName;
      ErrorOccurred = true;
      return false;
    }

    FoundDecl = getCanonicalSymbolDeclaration(FoundDecl);
    SpellingNames.push_back(FoundDecl->getNameAsString());
    AdditionalUSRFinder Finder(FoundDecl, Context);
    USRList.push_back(Finder.Find());
    return true;
  }

  void HandleTranslationUnit(ASTContext &Context) override {
    const SourceManager &SourceMgr = Context.getSourceManager();
    for (unsigned Offset : SymbolOffsets) {
      if (!FindSymbol(Context, SourceMgr, Offset, ""))
        return;
    }
    for (const std::string &QualifiedName : QualifiedNames) {
      if (!FindSymbol(Context, SourceMgr, 0, QualifiedName))
        return;
    }
  }

  ArrayRef<unsigned> SymbolOffsets;
  ArrayRef<std::string> QualifiedNames;
  std::vector<std::string> &SpellingNames;
  std::vector<std::vector<std::string>> &USRList;
  bool Force;
  bool &ErrorOccurred;
};

std::unique_ptr<ASTConsumer> USRFindingAction::newASTConsumer() {
  return std::make_unique<NamedDeclFindingConsumer>(
      SymbolOffsets, QualifiedNames, SpellingNames, USRList, Force,
      ErrorOccurred);
}

} // end namespace tooling
} // end namespace clang