//===-- lib/Semantics/openmp-utils.h --------------------------------------===//
//
// 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
//
//===----------------------------------------------------------------------===//
//
// Common utilities used in OpenMP semantic checks.
//
//===----------------------------------------------------------------------===//

#ifndef FORTRAN_SEMANTICS_OPENMP_UTILS_H
#define FORTRAN_SEMANTICS_OPENMP_UTILS_H

#include "flang/Common/indirection.h"
#include "flang/Evaluate/type.h"
#include "flang/Parser/char-block.h"
#include "flang/Parser/message.h"
#include "flang/Parser/openmp-utils.h"
#include "flang/Parser/parse-tree.h"
#include "flang/Parser/tools.h"
#include "flang/Semantics/tools.h"

#include "llvm/ADT/APInt.h"
#include "llvm/ADT/ArrayRef.h"
#include "llvm/Frontend/OpenMP/OMPContext.h"

#include <memory>
#include <optional>
#include <string>
#include <tuple>
#include <type_traits>
#include <utility>
#include <vector>

namespace Fortran::semantics {
class Scope;
class SemanticsContext;
class Symbol;

// Add this namespace to avoid potential conflicts
namespace omp {
using Fortran::parser::omp::BlockRange;
using Fortran::parser::omp::ExecutionPartIterator;
using Fortran::parser::omp::is_range_v;
using Fortran::parser::omp::LoopNestIterator;
using Fortran::parser::omp::LoopRange;

template <typename T, typename U = std::remove_const_t<T>> U AsRvalue(T &t) {
  return U(t);
}

template <typename T> T &&AsRvalue(T &&t) { return std::move(t); }

const Scope &GetScopingUnit(const Scope &scope);
const Scope &GetProgramUnit(const Scope &scope);

template <typename T> struct WithSource {
  template < //
      typename U = std::remove_reference_t<T>,
      typename = std::enable_if_t<std::is_default_constructible_v<U>>>
  WithSource() : value(), source() {}
  WithSource(const WithSource<T> &) = default;
  WithSource(WithSource<T> &&) = default;
  WithSource(const T &t, parser::CharBlock s) : value(t), source(s) {}
  WithSource(T &&t, parser::CharBlock s) : value(std::move(t)), source(s) {}
  WithSource &operator=(const WithSource<T> &) = default;
  WithSource &operator=(WithSource<T> &&) = default;

  using value_type = T;
  T value;
  parser::CharBlock source;
};

// There is no consistent way to get the source of an ActionStmt, but there
// is "source" in Statement<T>. This structure keeps the ActionStmt with the
// extracted source for further use.
struct SourcedActionStmt : public WithSource<const parser::ActionStmt *> {
  using WithSource<value_type>::WithSource;
  value_type stmt() const { return value; }
  operator bool() const { return stmt() != nullptr; }
};

SourcedActionStmt GetActionStmt(const parser::ExecutionPartConstruct *x);
SourcedActionStmt GetActionStmt(const parser::Block &block);

std::string ThisVersion(unsigned version);
std::string TryVersion(unsigned version);

const Symbol *GetObjectSymbol(
    const parser::OmpObject &object, bool ultimate = false);
const Symbol *GetArgumentSymbol(
    const parser::OmpArgument &argument, bool ultimate = false);

bool IsCommonBlock(const Symbol &sym);
bool IsExtendedListItem(const Symbol &sym);
bool IsVariableListItem(const Symbol &sym);
bool IsTypeParamInquiry(const Symbol &sym);
bool IsComplexPart(const Symbol &sym);
bool IsStructureComponent(const Symbol &sym);
bool IsPrivatizable(const Symbol &sym);
bool IsVarOrFunctionRef(const MaybeExpr &expr);

bool IsWholeAssumedSizeArray(const parser::OmpObject &object);

bool IsExtendedListItem(
    const parser::OmpObject &object, SemanticsContext *semaCtx);
bool IsLocatorListItem(
    const parser::OmpObject &object, SemanticsContext *semaCtx);
bool IsVariableListItem(
    const parser::OmpObject &object, SemanticsContext *semaCtx);

bool IsSubstring(const parser::OmpObject &object, SemanticsContext *semaCtx);
bool IsArrayElement(const parser::OmpObject &object, SemanticsContext *semaCtx);

const Symbol *GetHostSymbol(const Symbol &sym);

bool IsMapEnteringType(parser::OmpMapType::Value type);
bool IsMapExitingType(parser::OmpMapType::Value type);

// Returns true if the symbol has a temporary stack-allocated descriptor.
// This includes assumed-shape and assumed-rank dummy arguments that are
// not allocatable or pointer. These descriptors are created on the caller's
// stack and become invalid after the function returns.
bool HasTemporaryStackDescriptor(const Symbol &symbol);

MaybeExpr GetEvaluateExpr(const parser::Expr &parserExpr);
template <typename T> MaybeExpr GetEvaluateExpr(const T &inp) {
  return GetEvaluateExpr(parser::UnwrapRef<parser::Expr>(inp));
}

std::optional<evaluate::DynamicType> GetDynamicType(
    const parser::Expr &parserExpr);

std::optional<bool> GetLogicalValue(const SomeExpr &expr);
std::optional<int64_t> GetIntValueFromExpr(
    const parser::Expr &parserExpr, SemanticsContext *semaCtx = nullptr);

template <typename T>
std::optional<int64_t> GetIntValueFromExpr(
    const T &wrappedExpr, SemanticsContext *semaCtx = nullptr) {
  if (auto *parserExpr{parser::Unwrap<parser::Expr>(wrappedExpr)}) {
    return GetIntValueFromExpr(*parserExpr, semaCtx);
  }
  return std::nullopt;
}

// There are several clauses that take an optional, compile-time
// constant bool argument. Those clauses are stored as std::optional, e.g.
// OmpClause::ReverseOffload -> std::optional<OmpReverseOffloadClause>.
// Retrieve the logical value if present.
template <typename ClauseTy>
std::optional<bool> GetLogicalArgument(
    const std::optional<ClauseTy> &maybeClause, SemanticsContext &semaCtx) {
  if (maybeClause) {
    // Scalar<Logical<Constant<common::Indirection<Expr>>>>
    auto &parserExpr{parser::UnwrapRef<parser::Expr>(*maybeClause)};
    evaluate::ExpressionAnalyzer ea{semaCtx};
    if (auto &&maybeExpr{ea.Analyze(parserExpr)}) {
      if (auto v{GetLogicalValue(*maybeExpr)}) {
        return *v;
      }
    }
  }
  return std::nullopt;
}

std::optional<bool> IsContiguous(
    SemanticsContext &semaCtx, const parser::OmpObject &object);

/// Non-constant user condition expression and source for runtime lowering.
struct DynamicUserCondition {
  const parser::ScalarExpr *expr;
  parser::CharBlock source;
};

/// A context-selector feature that variant matching accepts syntactically but
/// cannot yet honour during selection. Callers are expected to diagnose these
/// (a lowering \c TODO or a semantic error) before calling
/// \c MakeVariantMatchInfo, which asserts none are present.
enum class UnsupportedSelectorFeature {
  None,
  /// A `target_device={...}` selector set.
  TargetDevice,
  /// A clause property (e.g. \c simdlen(8) in \c construct={simd(simdlen(8))})
  /// or an extension property (e.g. \c foo(bar) in
  /// \c implementation={my_trait(foo(bar))}).
  ClauseOrExtensionProperty,
};

/// Scan a parsed context selector for the first feature that variant matching
/// cannot yet honour (see \c UnsupportedSelectorFeature). Pure detection: emits
/// no diagnostics and has no side effects on any match info.
UnsupportedSelectorFeature FindUnsupportedSelectorFeature(
    const parser::traits::OmpContextSelectorSpecification &ctxSel,
    SemanticsContext &semaCtx);

/// Populate \p vmi from a parsed context selector. Score modifiers are
/// honoured (including on `condition(...)` selectors). Constant user
/// conditions are folded into user_condition_true/false traits; a non-constant
/// user condition is recorded as user_condition_unknown and the first such
/// expression is returned for the caller to lower as a runtime condition.
///
/// The caller must first reject unsupported selector features (see
/// \c FindUnsupportedSelectorFeature); this function asserts none are present.
std::optional<DynamicUserCondition> MakeVariantMatchInfo(
    llvm::omp::VariantMatchInfo &vmi,
    const parser::traits::OmpContextSelectorSpecification &ctxSel,
    SemanticsContext &semaCtx);

std::vector<SomeExpr> GetTopLevelDesignators(const SomeExpr &expr);
const SomeExpr *HasStorageOverlap(
    const SomeExpr &base, llvm::ArrayRef<SomeExpr> exprs);

bool IsAssignment(const parser::ActionStmt *x);
bool IsPointerAssignment(const evaluate::Assignment &x);

MaybeExpr MakeEvaluateExpr(const parser::OmpStylizedInstance &inp);

enum struct ListItemKind : uint32_t {
  Depend,
  DirectiveName,
  DirectiveSpecification,
  Extended,
  IntegerExpression,
  Interop,
  Locator,
  Operation,
  Parameter,
  ProcedureArgument,
  Variable,
};

std::optional<ListItemKind> GetArgumentListItemKind(
    llvm::omp::Clause clause, unsigned version);

bool IsLoopTransforming(llvm::omp::Directive dir);
bool HasDataEnvironment(llvm::omp::Directive dir);

bool IsFullUnroll(const parser::OmpDirectiveSpecification &spec);

inline bool IsDoConcurrentLegal(unsigned version) {
  // DO CONCURRENT is allowed (as an alternative to a Canonical Loop Nest)
  // in OpenMP 6.0+.
  return version >= 60;
}

struct LoopControl {
  LoopControl(LoopControl &&x) = default;
  LoopControl(const LoopControl &x) = default;
  LoopControl(const parser::LoopControl::Bounds &x);
  LoopControl(const parser::ConcurrentControl &x);

  const parser::Name &iv;
  WithSource<MaybeExpr> lbound, ubound, step;

private:
  static WithSource<MaybeExpr> fromParserExpr(const parser::Expr &x);
};

std::vector<LoopControl> GetLoopControls(const parser::DoConstruct &x);

/// A representation of a "because" message.
struct Reason {
  Reason() = default;
  Reason(Reason &&) = default;
  Reason(const Reason &);
  Reason &operator=(Reason &&) = default;
  Reason &operator=(const Reason &);

  parser::Messages msgs;

  template <typename... Ts> Reason &Say(Ts &&...args) {
    msgs.Say(std::forward<Ts>(args)...);
    return *this;
  }
  parser::Message &AttachTo(parser::Message &msg);
  Reason &Append(const Reason &other) {
    CopyFrom(other);
    return *this;
  }
  operator bool() const { return !msgs.empty(); }

private:
  void CopyFrom(const Reason &other);
};

// A property with an explanation of its value. Both, the property and the
// reason are optional (the reason can have no messages in it).
template <typename T> struct WithReason {
  std::optional<T> value;
  Reason reason;

  WithReason() = default;
  WithReason(std::optional<T> v, const Reason &r = Reason())
      : value(v), reason(r) {}
  operator bool() const { return value.has_value(); }
};

WithReason<int64_t> GetArgumentValueWithReason(
    const parser::OmpDirectiveSpecification &spec, llvm::omp::Clause clauseId,
    unsigned version, SemanticsContext *semaCtx = nullptr);
WithReason<int64_t> GetNumArgumentsWithReason(
    const parser::OmpDirectiveSpecification &spec, llvm::omp::Clause clauseId,
    unsigned version, SemanticsContext *semaCtx = nullptr);
WithReason<int64_t> GetHeightWithReason(
    const parser::OmpDirectiveSpecification &spec, unsigned version,
    SemanticsContext *semaCtx = nullptr);

/// Return the depth of the affected nest(s):
///   {affected-depth, must-be-perfect-nest}.
std::pair<WithReason<int64_t>, bool> GetAffectedNestDepthWithReason(
    const parser::OmpDirectiveSpecification &spec, unsigned version,
    SemanticsContext *semaCtx = nullptr);
/// Return the depth of the generated nest(s):
///   {generated-depth, is-perfect-nest}
std::pair<WithReason<int64_t>, bool> GetGeneratedNestDepthWithReason(
    const parser::OmpDirectiveSpecification &spec, unsigned version,
    SemanticsContext *semaCtx = nullptr);
/// Return the range of the affected nests in the sequence:
///   {first, count}.
/// If the range is "the whole sequence", the return value will be {1, -1}.
WithReason<std::pair<int64_t, int64_t>> GetAffectedLoopRangeWithReason(
    const parser::OmpDirectiveSpecification &spec, unsigned version,
    SemanticsContext *semaCtx = nullptr);
/// Return the depth in which all loops must be rectangular.
WithReason<int64_t> GetRectangularNestDepthWithReason(
    const parser::OmpDirectiveSpecification &spec, unsigned version,
    SemanticsContext *semaCtx = nullptr);

/// Calculate the minimum length of a sequence that contains the specified
/// range. If the range's count is -1, return -1 indicating all loops in the
/// sequence.
std::optional<int64_t> GetMinimumSequenceCount(
    std::optional<int64_t> first, std::optional<int64_t> count);
std::optional<int64_t> GetMinimumSequenceCount(
    std::optional<std::pair<int64_t, int64_t>> range);

/// Collect the set of DO loops present in the source code that are directly
/// affected by the given loop construct. This does not include any DO loops
/// that are affected by any construct nested in `x`.
/// Returns std::nullopt if `x` or code nested in `x` was malformed in a
/// way that prevented the function from returning an accurate result.
std::optional<std::vector<const parser::DoConstruct *>> CollectAffectedDoLoops(
    const parser::OpenMPLoopConstruct &x, unsigned version,
    SemanticsContext *semaCtx = nullptr);

struct LoopSequence {
  LoopSequence(const parser::ExecutionPartConstruct &root, unsigned version,
      bool allowAllLoops = false, SemanticsContext *semaCtx = nullptr);

  template <typename R, typename = std::enable_if_t<is_range_v<R>>>
  LoopSequence(const R &range, unsigned version, bool allowAllLoops = false,
      SemanticsContext *semaCtx = nullptr)
      : version_(version), allowAllLoops_(allowAllLoops), semaCtx_(semaCtx) {
    entry_ = std::make_unique<Construct>(range, nullptr);
    createChildrenFromRange(entry_->location);
    precalculate();
  }

  struct Depth {
    // If this sequence is a nest, the depth of the Canonical Loop Nest rooted
    // at this sequence. Otherwise unspecified.
    WithReason<int64_t> semantic;
    // If this sequence is a nest, the depth of the perfect Canonical Loop Nest
    // rooted at this sequence. Otherwise unspecified.
    WithReason<int64_t> perfect;
  };

  bool isNest() const { return length_.value == 1; }
  const WithReason<int64_t> &length() const { return length_; }
  const WithReason<int64_t> &height() const { return height_; }
  const Depth &depth() const { return depth_; }
  const std::vector<LoopSequence> &children() const { return children_; }
  const parser::ExecutionPartConstruct *owner() const { return entry_->owner; }

  WithReason<bool> isWellFormedSequence() const;
  WithReason<bool> isWellFormedNest() const;

  /// Return the first DO CONCURRENT loop contained in this sequence.
  /// If there are no such loops, return nullptr.
  const LoopSequence *getNestedDoConcurrent() const;

  std::vector<LoopControl> getLoopControls() const;
  // Check if this loop's bounds are invariant in each of the `outer`
  // constructs.
  WithReason<bool> isRectangular(
      const std::vector<const LoopSequence *> &outer) const;

private:
  using Construct = ExecutionPartIterator::Construct;

  LoopSequence(std::unique_ptr<Construct> entry, unsigned version,
      bool allowAllLoops, SemanticsContext *semaCtx = nullptr);

  template <typename R, typename = std::enable_if_t<is_range_v<R>>>
  void createChildrenFromRange(const R &range) {
    createChildrenFromRange(range.begin(), range.end());
  }

  std::unique_ptr<Construct> createConstructEntry(
      const parser::ExecutionPartConstruct &code);

  void createChildrenFromRange( //
      ExecutionPartIterator::IteratorType begin,
      ExecutionPartIterator::IteratorType end);

  /// Precalculate length and depth.
  void precalculate();

  WithReason<int64_t> calculateLength() const;
  WithReason<int64_t> getNestedLength() const;
  Depth calculateDepths() const;
  Depth getNestedDepths() const;
  WithReason<int64_t> calculateHeight() const;

  /// The construct that is not a loop or a loop-transforming construct,
  /// that is also not a valid intervening code. Unset if no such code is
  /// present.
  const parser::ExecutionPartConstruct *invalidIC_{nullptr};
  /// The construct that is not a loop or a loop-transforming construct,
  /// whose presence would prevent perfect nesting of loops (i.e. code that
  /// is not "transparent" to a perfect nest).
  const parser::ExecutionPartConstruct *opaqueIC_{nullptr};

  /// Precalculated length of the sequence. Note that this is different from
  /// the number of children because a child may result in a sequence, for
  /// example a fuse with a reduced loop range. The length of that sequence
  /// adds to the length of the owning LoopSequence.
  WithReason<int64_t> length_;
  /// Precalculated depths. Only meaningful if the sequence is a nest.
  Depth depth_;
  /// Precalculated height of the sequence. The height is the difference
  /// in the nesting level between "this" and any of the children (should
  /// be the same for each child). Intuitively it is the number of nested
  /// loops that are added by this construct. If this->depth_ included
  /// child->depth_ for some child, then
  ///   height_ = this->depth_ - child->depth_
  WithReason<int64_t> height_;

  // The core structure of the class:
  unsigned version_; // Needed for GetXyzWithReason
  bool allowAllLoops_;
  std::unique_ptr<Construct> entry_;
  std::vector<LoopSequence> children_;
  SemanticsContext *semaCtx_{nullptr};
};

// ---------------------------------------------------------------------------
// Trait-matching helpers shared between metadirective lowering and
// declare-variant semantic recording.
// ---------------------------------------------------------------------------

/// Map a parsed trait-set name to the corresponding LLVM OMP TraitSet enum.
llvm::omp::TraitSet MapTraitSet(parser::OmpTraitSetSelectorName::Value name);

/// Map a parsed trait-selector name (plus its containing set) to the
/// corresponding LLVM OMP TraitSelector enum.
llvm::omp::TraitSelector MapTraitSelector(
    const parser::OmpTraitSelectorName &name, llvm::omp::TraitSet set);

/// Try to constant-fold a user condition expression to a boolean.
std::optional<bool> EvaluateUserCondition(
    SemanticsContext &semaCtx, const parser::ScalarExpr &scalarExpr);

/// Extract the optional score value from trait properties.
llvm::APInt *GetTraitScore(
    const std::optional<parser::OmpTraitSelector::Properties> &props,
    SemanticsContext &semaCtx, std::optional<llvm::APInt> &scoreStorage);

/// Collect trait property names (vendor, kind, arch, isa, etc.) into a VMI.
/// Non-name properties (clause, extension) are silently skipped; the caller is
/// responsible for diagnosing them before invoking this function.
void ProcessTraitProperties(llvm::omp::VariantMatchInfo &vmi,
    llvm::omp::TraitSet set, llvm::omp::TraitSelector selector,
    const std::optional<parser::OmpTraitSelector::Properties> &props,
    llvm::APInt *scorePtr);

} // namespace omp
} // namespace Fortran::semantics

#endif // FORTRAN_SEMANTICS_OPENMP_UTILS_H
