//===----------------------------------------------------------------------===//
//
// 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
//
//===----------------------------------------------------------------------===//

#include "ElseAfterReturnCheck.h"
#include "clang/AST/ASTContext.h"
#include "clang/ASTMatchers/ASTMatchFinder.h"
#include "clang/Lex/Lexer.h"
#include "clang/Lex/Preprocessor.h"
#include "clang/Tooling/FixIt.h"
#include "llvm/ADT/SmallVector.h"

using namespace clang::ast_matchers;

namespace clang::tidy::readability {

namespace {

class PPConditionalCollector : public PPCallbacks {
public:
  PPConditionalCollector(
      ElseAfterReturnCheck::ConditionalBranchMap &Collections,
      const SourceManager &SM)
      : Collections(Collections), SM(SM) {}
  void Endif(SourceLocation Loc, SourceLocation IfLoc) override {
    if (!SM.isWrittenInSameFile(Loc, IfLoc))
      return;
    SmallVectorImpl<SourceRange> &Collection = Collections[SM.getFileID(Loc)];
    assert(Collection.empty() || Collection.back().getEnd() < Loc);
    Collection.emplace_back(IfLoc, Loc);
  }

private:
  ElseAfterReturnCheck::ConditionalBranchMap &Collections;
  const SourceManager &SM;
};

AST_MATCHER_P(Stmt, stripLabelLikeStatements,
              ast_matchers::internal::Matcher<Stmt>, InnerMatcher) {
  const Stmt *S = Node.stripLabelLikeStatements();
  return InnerMatcher.matches(*S, Finder, Builder);
}

AST_MATCHER_P(Stmt, hasFinalStmt, ast_matchers::internal::Matcher<Stmt>,
              InnerMatcher) {
  for (const Stmt *S = &Node;;) {
    S = S->stripLabelLikeStatements();
    if (const auto *Compound = dyn_cast<CompoundStmt>(S)) {
      if (Compound->body_empty())
        return false;
      S = Compound->body_back();
    } else {
      return InnerMatcher.matches(*S, Finder, Builder);
    }
  }
}

} // namespace

static constexpr char InterruptingStr[] = "interrupting";
static constexpr char WarningMessage[] = "do not use 'else' after %0";
static constexpr char WarnOnUnfixableStr[] = "WarnOnUnfixable";
static constexpr char WarnOnConditionVariablesStr[] =
    "WarnOnConditionVariables";

static const DeclRefExpr *findUsage(const Stmt *Node, const Decl *D) {
  if (!Node)
    return nullptr;
  if (const auto *DeclRef = dyn_cast<DeclRefExpr>(Node)) {
    if (DeclRef->getDecl() == D)
      return DeclRef;
  } else {
    for (const Stmt *ChildNode : Node->children())
      if (const DeclRefExpr *Result = findUsage(ChildNode, D))
        return Result;
  }
  return nullptr;
}

static const DeclRefExpr *findUsageRange(const Stmt *Node,
                                         DeclStmt::decl_const_range Decls) {
  if (!Node)
    return nullptr;
  if (const auto *DeclRef = dyn_cast<DeclRefExpr>(Node)) {
    if (llvm::is_contained(Decls, DeclRef->getDecl()))
      return DeclRef;
  } else {
    for (const Stmt *ChildNode : Node->children())
      if (const DeclRefExpr *Result = findUsageRange(ChildNode, Decls))
        return Result;
  }
  return nullptr;
}

static const DeclRefExpr *checkInitDeclUsageInElse(const IfStmt *If) {
  const auto *InitDeclStmt = dyn_cast_or_null<DeclStmt>(If->getInit());
  if (!InitDeclStmt)
    return nullptr;
  if (InitDeclStmt->isSingleDecl()) {
    const Decl *InitDecl = InitDeclStmt->getSingleDecl();
    assert(isa<VarDecl>(InitDecl) && "SingleDecl must be a VarDecl");
    return findUsage(If->getElse(), InitDecl);
  }
  return findUsageRange(If->getElse(), InitDeclStmt->decls());
}

static const DeclRefExpr *checkConditionVarUsageInElse(const IfStmt *If) {
  if (const VarDecl *CondVar = If->getConditionVariable())
    return findUsage(If->getElse(), CondVar);
  return nullptr;
}

static bool containsDeclInScope(const Stmt *Node) {
  if (isa<DeclStmt>(Node))
    return true;
  if (const auto *Compound = dyn_cast<CompoundStmt>(Node))
    return llvm::any_of(Compound->body(), [](const Stmt *SubNode) {
      return isa<DeclStmt>(SubNode);
    });
  return false;
}

static void removeElseAndBrackets(DiagnosticBuilder &Diag, ASTContext &Context,
                                  const Stmt *Else, SourceLocation ElseLoc) {
  auto Remap = [&](SourceLocation Loc) {
    return Context.getSourceManager().getExpansionLoc(Loc);
  };

  if (const auto *CS = dyn_cast<CompoundStmt>(Else)) {
    Diag << tooling::fixit::createRemoval(ElseLoc)
         << tooling::fixit::createRemoval(Remap(CS->getLBracLoc()))
         << tooling::fixit::createRemoval(Remap(CS->getRBracLoc()));
  } else {
    Diag << tooling::fixit::createRemoval(Remap(ElseLoc));
  }
}

ElseAfterReturnCheck::ElseAfterReturnCheck(StringRef Name,
                                           ClangTidyContext *Context)
    : ClangTidyCheck(Name, Context),
      WarnOnUnfixable(Options.get(WarnOnUnfixableStr, true)),
      WarnOnConditionVariables(Options.get(WarnOnConditionVariablesStr, true)) {
}

void ElseAfterReturnCheck::storeOptions(ClangTidyOptions::OptionMap &Opts) {
  Options.store(Opts, WarnOnUnfixableStr, WarnOnUnfixable);
  Options.store(Opts, WarnOnConditionVariablesStr, WarnOnConditionVariables);
}

void ElseAfterReturnCheck::registerPPCallbacks(const SourceManager &SM,
                                               Preprocessor *PP,
                                               Preprocessor *ModuleExpanderPP) {
  PP->addPPCallbacks(
      std::make_unique<PPConditionalCollector>(this->PPConditionals, SM));
}

void ElseAfterReturnCheck::registerMatchers(MatchFinder *Finder) {
  const auto InterruptsControlFlow =
      stmt(anyOf(returnStmt(), continueStmt(), breakStmt(), cxxThrowExpr(),
                 callExpr(callee(functionDecl(isNoReturn())))));

  const auto IfWithInterruptingThenElse =
      ifStmt(unless(isConstexpr()), unless(isConsteval()),
             hasThen(hasFinalStmt(InterruptsControlFlow.bind(InterruptingStr))),
             hasElse(stmt().bind("else")))
          .bind("if");

  Finder->addMatcher(compoundStmt(forEach(stripLabelLikeStatements(
                                      IfWithInterruptingThenElse)))
                         .bind("cs"),
                     this);
}

static bool hasPreprocessorBranchEndBetweenLocations(
    const ElseAfterReturnCheck::ConditionalBranchMap &ConditionalBranchMap,
    const SourceManager &SM, SourceLocation StartLoc, SourceLocation EndLoc) {
  const SourceLocation ExpandedStartLoc = SM.getExpansionLoc(StartLoc);
  const SourceLocation ExpandedEndLoc = SM.getExpansionLoc(EndLoc);
  if (!SM.isWrittenInSameFile(ExpandedStartLoc, ExpandedEndLoc))
    return false;

  // StartLoc and EndLoc expand to the same macro.
  if (ExpandedStartLoc == ExpandedEndLoc)
    return false;

  assert(ExpandedStartLoc < ExpandedEndLoc);

  auto Iter = ConditionalBranchMap.find(SM.getFileID(ExpandedEndLoc));

  if (Iter == ConditionalBranchMap.end() || Iter->getSecond().empty())
    return false;

  const SmallVectorImpl<SourceRange> &ConditionalBranches = Iter->getSecond();

  assert(llvm::is_sorted(ConditionalBranches,
                         [](const SourceRange &LHS, const SourceRange &RHS) {
                           return LHS.getEnd() < RHS.getEnd();
                         }));

  // First conditional block that ends after ExpandedStartLoc.
  const auto *Begin =
      llvm::lower_bound(ConditionalBranches, ExpandedStartLoc,
                        [](const SourceRange &LHS, const SourceLocation &RHS) {
                          return LHS.getEnd() < RHS;
                        });
  const auto *End = ConditionalBranches.end();
  for (; Begin != End && Begin->getEnd() < ExpandedEndLoc; ++Begin)
    if (Begin->getBegin() < ExpandedStartLoc)
      return true;
  return false;
}

static StringRef getControlFlowString(const Stmt &Stmt) {
  if (isa<ReturnStmt>(Stmt))
    return "'return'";
  if (isa<ContinueStmt>(Stmt))
    return "'continue'";
  if (isa<BreakStmt>(Stmt))
    return "'break'";
  if (isa<CXXThrowExpr>(Stmt))
    return "'throw'";
  if (isa<CallExpr>(Stmt))
    return "calling a function that doesn't return";
  llvm_unreachable("Unknown control flow interrupter");
}

void ElseAfterReturnCheck::check(const MatchFinder::MatchResult &Result) {
  const auto *If = Result.Nodes.getNodeAs<IfStmt>("if");
  const auto *Else = Result.Nodes.getNodeAs<Stmt>("else");
  const auto *OuterScope = Result.Nodes.getNodeAs<CompoundStmt>("cs");
  const auto *Interrupt = Result.Nodes.getNodeAs<Stmt>(InterruptingStr);
  const SourceLocation ElseLoc = If->getElseLoc();

  if (hasPreprocessorBranchEndBetweenLocations(
          PPConditionals, *Result.SourceManager, Interrupt->getBeginLoc(),
          ElseLoc))
    return;

  const bool IsLastInScope = OuterScope->body_back() == If;
  const StringRef ControlFlowInterrupter = getControlFlowString(*Interrupt);

  if (!IsLastInScope && containsDeclInScope(Else)) {
    if (WarnOnUnfixable) {
      // Warn, but don't attempt an autofix.
      diag(ElseLoc, WarningMessage) << ControlFlowInterrupter;
    }
    return;
  }

  if (checkConditionVarUsageInElse(If) != nullptr) {
    if (!WarnOnConditionVariables)
      return;
    if (IsLastInScope) {
      // If the if statement is the last statement of its enclosing statements
      // scope, we can pull the decl out of the if statement.
      DiagnosticBuilder Diag = diag(ElseLoc, WarningMessage)
                               << ControlFlowInterrupter
                               << SourceRange(ElseLoc);
      if (checkInitDeclUsageInElse(If) != nullptr) {
        Diag << tooling::fixit::createReplacement(
                    SourceRange(If->getIfLoc()),
                    (tooling::fixit::getText(*If->getInit(), *Result.Context) +
                     StringRef("\n"))
                        .str())
             << tooling::fixit::createRemoval(If->getInit()->getSourceRange());
      }
      const DeclStmt *VDeclStmt = If->getConditionVariableDeclStmt();
      const VarDecl *VDecl = If->getConditionVariable();
      const std::string Repl =
          (tooling::fixit::getText(*VDeclStmt, *Result.Context) +
           StringRef(";\n") +
           tooling::fixit::getText(If->getIfLoc(), *Result.Context))
              .str();
      Diag << tooling::fixit::createReplacement(SourceRange(If->getIfLoc()),
                                                Repl)
           << tooling::fixit::createReplacement(VDeclStmt->getSourceRange(),
                                                VDecl->getName());
      removeElseAndBrackets(Diag, *Result.Context, Else, ElseLoc);
    } else if (WarnOnUnfixable) {
      // Warn, but don't attempt an autofix.
      diag(ElseLoc, WarningMessage) << ControlFlowInterrupter;
    }
    return;
  }

  if (checkInitDeclUsageInElse(If) != nullptr) {
    if (!WarnOnConditionVariables)
      return;
    if (IsLastInScope) {
      // If the if statement is the last statement of its enclosing statements
      // scope, we can pull the decl out of the if statement.
      DiagnosticBuilder Diag = diag(ElseLoc, WarningMessage)
                               << ControlFlowInterrupter
                               << SourceRange(ElseLoc);
      Diag << tooling::fixit::createReplacement(
                  SourceRange(If->getIfLoc()),
                  (tooling::fixit::getText(*If->getInit(), *Result.Context) +
                   "\n" +
                   tooling::fixit::getText(If->getIfLoc(), *Result.Context))
                      .str())
           << tooling::fixit::createRemoval(If->getInit()->getSourceRange());
      removeElseAndBrackets(Diag, *Result.Context, Else, ElseLoc);
    } else if (WarnOnUnfixable) {
      // Warn, but don't attempt an autofix.
      diag(ElseLoc, WarningMessage) << ControlFlowInterrupter;
    }
    return;
  }

  DiagnosticBuilder Diag = diag(ElseLoc, WarningMessage)
                           << ControlFlowInterrupter << SourceRange(ElseLoc);
  removeElseAndBrackets(Diag, *Result.Context, Else, ElseLoc);
}

} // namespace clang::tidy::readability
