//===-- VectorSubscripts.h -- vector subscripts tools -----------*- C++ -*-===//
//
// 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
/// \brief Defines a compiler internal representation for lowered designators
///  containing vector subscripts. This representation allows working on such
///  designators in custom ways while ensuring the designator subscripts are
///  only evaluated once. It is mainly intended for cases that do not fit in
///  the array expression lowering framework like input IO in presence of
///  vector subscripts.
///
//===----------------------------------------------------------------------===//

#ifndef FORTRAN_LOWER_VECTORSUBSCRIPTS_H
#define FORTRAN_LOWER_VECTORSUBSCRIPTS_H

#include "flang/Optimizer/Builder/BoxValue.h"

namespace fir {
class FirOpBuilder;
}

namespace Fortran {

namespace evaluate {
template <typename>
class Expr;
struct SomeType;
} // namespace evaluate

namespace lower {

class AbstractConverter;
class StatementContext;

/// VectorSubscriptBox is a lowered representation for any Designator<T> that
/// contain at least one vector subscript.
///
/// A designator `x%a(i,j)%b(1:foo():1, vector, k)%c%d(m)%e1
/// Is lowered into:
///   - an ExtendedValue for ranked base (x%a(i,j)%b)
///   - mlir:Values and ExtendedValues for the triplet, vector subscript and
///     scalar subscripts of the ranked array reference (1:foo():1, vector, k)
///   - a list of fir.field_index and scalar integers mlir::Value for the
///   component
///     path at the right of the ranked array ref (%c%d(m)%e).
///
/// This representation allows later creating loops over the designator elements
/// and fir.array_coor to get the element addresses without re-evaluating any
/// sub-expressions.
class VectorSubscriptBox {
public:
  /// Type of the callbacks that can be passed to work with the element
  /// addresses.
  using ElementalGenerator = std::function<void(const fir::ExtendedValue &)>;
  using ElementalGeneratorWithBoolReturn =
      std::function<mlir::Value(const fir::ExtendedValue &)>;
  struct LoweredVectorSubscript {
    LoweredVectorSubscript(fir::ExtendedValue &&vector, mlir::Value size)
        : vector{std::move(vector)}, size{size} {}
    fir::ExtendedValue vector;
    // Vector size, guaranteed to be of indexType.
    mlir::Value size;
  };
  struct LoweredTriplet {
    // Triplets value, guaranteed to be of indexType.
    mlir::Value lb;
    mlir::Value ub;
    mlir::Value stride;
  };
  using LoweredSubscript =
      std::variant<mlir::Value, LoweredTriplet, LoweredVectorSubscript>;
  using MaybeSubstring = llvm::SmallVector<mlir::Value, 2>;
  VectorSubscriptBox(
      fir::ExtendedValue &&loweredBase,
      llvm::SmallVector<LoweredSubscript, 16> &&loweredSubscripts,
      llvm::SmallVector<mlir::Value> &&componentPath,
      MaybeSubstring substringBounds, mlir::Type elementType)
      : loweredBase{std::move(loweredBase)}, loweredSubscripts{std::move(
                                                 loweredSubscripts)},
        componentPath{std::move(componentPath)},
        substringBounds{substringBounds}, elementType{elementType} {};

  /// Loop over the elements described by the VectorSubscriptBox, and call
  /// \p elementalGenerator inside the loops with the element addresses.
  void loopOverElements(fir::FirOpBuilder &builder, mlir::Location loc,
                        const ElementalGenerator &elementalGenerator);

  /// Loop over the elements described by the VectorSubscriptBox while a
  /// condition is true, and call \p elementalGenerator inside the loops with
  /// the element addresses. The initial condition value is \p initialCondition,
  /// and then it is the result of \p elementalGenerator. The value of the
  /// condition after the loops is returned.
  mlir::Value loopOverElementsWhile(
      fir::FirOpBuilder &builder, mlir::Location loc,
      const ElementalGeneratorWithBoolReturn &elementalGenerator,
      mlir::Value initialCondition);

  /// Return the type of the elements of the array section.
  mlir::Type getElementType() { return elementType; }

private:
  /// Common implementation for DoLoop and IterWhile loop creations.
  template <typename LoopType, typename Generator>
  mlir::Value loopOverElementsBase(fir::FirOpBuilder &builder,
                                   mlir::Location loc,
                                   const Generator &elementalGenerator,
                                   mlir::Value initialCondition);
  /// Create sliceOp for the designator.
  mlir::Value createSlice(fir::FirOpBuilder &builder, mlir::Location loc);

  /// Create ExtendedValue the element inside the loop.
  fir::ExtendedValue getElementAt(fir::FirOpBuilder &builder,
                                  mlir::Location loc, mlir::Value shape,
                                  mlir::Value slice,
                                  mlir::ValueRange inductionVariables);

  /// Generate the [lb, ub, step] to loop over the section (in loop order, not
  /// Fortran dimension order).
  llvm::SmallVector<std::tuple<mlir::Value, mlir::Value, mlir::Value>>
  genLoopBounds(fir::FirOpBuilder &builder, mlir::Location loc);

  /// Lowered base of the ranked array ref.
  fir::ExtendedValue loweredBase;
  /// Subscripts values of the rank arrayRef part.
  llvm::SmallVector<LoweredSubscript, 16> loweredSubscripts;
  /// Scalar subscripts and components at the right of the ranked
  /// array ref part of any.
  llvm::SmallVector<mlir::Value> componentPath;
  /// List of substring bounds if this is a substring (only the lower bound if
  /// the upper is implicit).
  MaybeSubstring substringBounds;
  /// Type of the elements described by this array section.
  mlir::Type elementType;
};

/// Lower \p expr, that must be an designator containing vector subscripts, to a
/// VectorSubscriptBox representation. This causes evaluation of all the
/// subscripts. Any required clean-ups from subscript expression are added to \p
/// stmtCtx.
VectorSubscriptBox genVectorSubscriptBox(
    mlir::Location loc, Fortran::lower::AbstractConverter &converter,
    Fortran::lower::StatementContext &stmtCtx,
    const Fortran::evaluate::Expr<Fortran::evaluate::SomeType> &expr);

} // namespace lower
} // namespace Fortran

#endif // FORTRAN_LOWER_VECTORSUBSCRIPTS_H