//===-- lib/Parser/tools.cpp ----------------------------------------------===//
//
// 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 "flang/Parser/tools.h"

namespace Fortran::parser {

const Name &GetLastName(const Name &x) { return x; }

const Name &GetLastName(const StructureComponent &x) {
  return GetLastName(x.Component());
}

const Name &GetLastName(const DataRef &x) {
  return common::visit(
      common::visitors{
          [](const Name &name) -> const Name & { return name; },
          [](const common::Indirection<StructureComponent> &sc)
              -> const Name & { return GetLastName(sc.value()); },
          [](const common::Indirection<ArrayElement> &sc) -> const Name & {
            return GetLastName(sc.value().Base());
          },
          [](const common::Indirection<CoindexedNamedObject> &ci)
              -> const Name & {
            return GetLastName(std::get<DataRef>(ci.value().t));
          },
      },
      x.u);
}

const Name &GetLastName(const Substring &x) {
  return GetLastName(std::get<DataRef>(x.t));
}

const Name &GetLastName(const Designator &x) {
  return common::visit(
      [](const auto &y) -> const Name & { return GetLastName(y); }, x.u);
}

const Name &GetLastName(const ProcComponentRef &x) {
  return GetLastName(x.v.thing);
}

const Name &GetLastName(const ProcedureDesignator &x) {
  return common::visit(
      [](const auto &y) -> const Name & { return GetLastName(y); }, x.u);
}

const Name &GetLastName(const Call &x) {
  return GetLastName(std::get<ProcedureDesignator>(x.t));
}

const Name &GetLastName(const FunctionReference &x) { return GetLastName(x.v); }

const Name &GetLastName(const Variable &x) {
  return common::visit(
      [](const auto &indirection) -> const Name & {
        return GetLastName(indirection.value());
      },
      x.u);
}

const Name &GetLastName(const AllocateObject &x) {
  return common::visit(
      [](const auto &y) -> const Name & { return GetLastName(y); }, x.u);
}

const Name &GetFirstName(const Name &x) { return x; }

const Name &GetFirstName(const StructureComponent &x) {
  return GetFirstName(x.Base());
}

const Name &GetFirstName(const DataRef &x) {
  return common::visit(
      common::visitors{
          [](const Name &name) -> const Name & { return name; },
          [](const common::Indirection<StructureComponent> &sc)
              -> const Name & { return GetFirstName(sc.value()); },
          [](const common::Indirection<ArrayElement> &sc) -> const Name & {
            return GetFirstName(sc.value().Base());
          },
          [](const common::Indirection<CoindexedNamedObject> &ci)
              -> const Name & {
            return GetFirstName(std::get<DataRef>(ci.value().t));
          },
      },
      x.u);
}

const Name &GetFirstName(const Substring &x) {
  return GetFirstName(std::get<DataRef>(x.t));
}

const Name &GetFirstName(const Designator &x) {
  return common::visit(
      [](const auto &y) -> const Name & { return GetFirstName(y); }, x.u);
}

const Name &GetFirstName(const ProcComponentRef &x) {
  return GetFirstName(x.v.thing);
}

const Name &GetFirstName(const ProcedureDesignator &x) {
  return common::visit(
      [](const auto &y) -> const Name & { return GetFirstName(y); }, x.u);
}

const Name &GetFirstName(const Call &x) {
  return GetFirstName(std::get<ProcedureDesignator>(x.t));
}

const Name &GetFirstName(const FunctionReference &x) {
  return GetFirstName(x.v);
}

const Name &GetFirstName(const Variable &x) {
  return common::visit(
      [](const auto &indirect) -> const Name & {
        return GetFirstName(indirect.value());
      },
      x.u);
}

const Name &GetFirstName(const EntityDecl &x) {
  return std::get<ObjectName>(x.t);
}

const Name &GetFirstName(const AccObject &x) {
  return common::visit(
      [](const auto &y) -> const Name & { return GetFirstName(y); }, x.u);
}

bool IsBOZLiteral(const ScalarIntConstantExpr &x) {
  const Expr &expr{UnwrapRef<Expr>(x)};
  const auto *literal{std::get_if<LiteralConstant>(&expr.u)};
  return literal && std::holds_alternative<BOZLiteralConstant>(literal->u);
}

const CoindexedNamedObject *GetCoindexedNamedObject(const DataRef &base) {
  return common::visit(
      common::visitors{
          [](const Name &) -> const CoindexedNamedObject * { return nullptr; },
          [](const common::Indirection<CoindexedNamedObject> &x)
              -> const CoindexedNamedObject * { return &x.value(); },
          [](const auto &x) -> const CoindexedNamedObject * {
            return GetCoindexedNamedObject(x.value().Base());
          },
      },
      base.u);
}
const CoindexedNamedObject *GetCoindexedNamedObject(
    const Designator &designator) {
  return common::visit(
      common::visitors{
          [](const DataRef &x) -> const CoindexedNamedObject * {
            return GetCoindexedNamedObject(x);
          },
          [](const Substring &x) -> const CoindexedNamedObject * {
            return GetCoindexedNamedObject(std::get<DataRef>(x.t));
          },
      },
      designator.u);
}
const CoindexedNamedObject *GetCoindexedNamedObject(const Variable &variable) {
  return common::visit(
      common::visitors{
          [](const common::Indirection<Designator> &designator)
              -> const CoindexedNamedObject * {
            return GetCoindexedNamedObject(designator.value());
          },
          [](const auto &) -> const CoindexedNamedObject * { return nullptr; },
      },
      variable.u);
}
const CoindexedNamedObject *GetCoindexedNamedObject(
    const AllocateObject &allocateObject) {
  return common::visit(
      common::visitors{
          [](const StructureComponent &x) -> const CoindexedNamedObject * {
            return GetCoindexedNamedObject(x.Base());
          },
          [](const auto &) -> const CoindexedNamedObject * { return nullptr; },
      },
      allocateObject.u);
}

bool CheckForSingleVariableOnRHS(const AssignmentStmt &assignmentStmt) {
  return Unwrap<Designator>(std::get<Expr>(assignmentStmt.t)) != nullptr;
}

const Name *GetDesignatorNameIfDataRef(const Designator &designator) {
  const auto *dataRef{std::get_if<DataRef>(&designator.u)};
  return dataRef ? std::get_if<Name>(&dataRef->u) : nullptr;
}

// Get the Label from a Statement<...> contained in an ExecutionPartConstruct,
// or std::nullopt, if there is no Statement<...> contained in there.
template <typename T>
static std::optional<Label> GetStatementLabelHelper(const T &stmt) {
  if constexpr (IsStatement<T>::value) {
    return stmt.label;
  } else if constexpr (WrapperTrait<T>) {
    return GetStatementLabelHelper(stmt.v);
  } else if constexpr (UnionTrait<T>) {
    return common::visit(
        [&](auto &&s) { return GetStatementLabelHelper(s); }, stmt.u);
  }
  return std::nullopt;
}

std::optional<Label> GetStatementLabel(const ExecutionPartConstruct &x) {
  return GetStatementLabelHelper(x);
}

std::optional<Label> GetFinalLabel(const Block &x) {
  if (!x.empty()) {
    const ExecutionPartConstruct &last{x.back()};
    if (auto *omp{Unwrap<OpenMPConstruct>(last)}) {
      return GetFinalLabel(*omp);
    } else if (auto *acc{Unwrap<OpenACCConstruct>(last)}) {
      return GetFinalLabel(*acc);
    } else if (auto *doLoop{Unwrap<DoConstruct>(last)}) {
      return GetFinalLabel(std::get<Block>(doLoop->t));
    } else {
      return GetStatementLabel(x.back());
    }
  } else {
    return std::nullopt;
  }
}

std::optional<Label> GetFinalLabel(const OpenMPConstruct &x) {
  return common::visit(
      [](auto &&s) -> std::optional<Label> {
        using TypeS = llvm::remove_cvref_t<decltype(s)>;
        if constexpr (std::is_same_v<TypeS, OpenMPSectionsConstruct>) {
          auto &list{std::get<std::list<OpenMPConstruct>>(s.t)};
          if (!list.empty()) {
            return GetFinalLabel(list.back());
          } else {
            return std::nullopt;
          }
        } else if constexpr ( //
            std::is_same_v<TypeS, OpenMPLoopConstruct> ||
            std::is_same_v<TypeS, OmpSectionDirective> ||
            std::is_base_of_v<OmpBlockConstruct, TypeS>) {
          return GetFinalLabel(std::get<Block>(s.t));
        } else {
          return std::nullopt;
        }
      },
      x.u);
}

// The DO loop associated with a loop or combined construct is only terminated
// by a labeled statement when the construct has no end directive.
template <typename END, typename C>
static std::optional<Label> GetFinalLabelOfAssociatedLoop(const C &x) {
  if (std::get<std::optional<END>>(x.t)) {
    return std::nullopt;
  }
  if (const auto &doLoop{std::get<std::optional<DoConstruct>>(x.t)}) {
    return GetFinalLabel(std::get<Block>(doLoop->t));
  }
  return std::nullopt;
}

std::optional<Label> GetFinalLabel(const OpenACCConstruct &x) {
  return common::visit(
      common::visitors{
          [](const OpenACCBlockConstruct &x) -> std::optional<Label> {
            return GetFinalLabel(std::get<Block>(x.t));
          },
          [](const OpenACCLoopConstruct &x) -> std::optional<Label> {
            return GetFinalLabelOfAssociatedLoop<AccEndLoop>(x);
          },
          [](const OpenACCCombinedConstruct &x) -> std::optional<Label> {
            return GetFinalLabelOfAssociatedLoop<AccEndCombinedDirective>(x);
          },
          [](const OpenACCAtomicConstruct &x) -> std::optional<Label> {
            return common::visit(
                common::visitors{
                    [](const auto &x) { // AtomicRead, AtomicWrite, AtomicUpdate
                      return std::get<Statement<AssignmentStmt>>(x.t).label;
                    },
                    [](const AccAtomicCapture &x) {
                      return std::get<AccAtomicCapture::Stmt2>(x.t).v.label;
                    },
                },
                x.u);
          },
          [](const auto &) -> std::optional<Label> { return std::nullopt; },
      },
      x.u);
}

} // namespace Fortran::parser