BoxValue.h

//===-- BoxValue.h -- internal box values -----------------------*- 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
//
//===----------------------------------------------------------------------===//
//
// Coding style: https://mlir.llvm.org/getting_started/DeveloperGuide/
//
//===----------------------------------------------------------------------===//
 
#ifndef FORTRAN_OPTIMIZER_BUILDER_BOXVALUE_H
#define FORTRAN_OPTIMIZER_BUILDER_BOXVALUE_H
 
#include "flang/Optimizer/Dialect/FIRType.h"
#include "flang/Optimizer/Support/FatalError.h"
#include "flang/Optimizer/Support/Matcher.h"
#include "mlir/IR/OperationSupport.h"
#include "mlir/IR/Value.h"
#include "llvm/ADT/SmallVector.h"
#include "llvm/Support/Compiler.h"
#include "llvm/Support/raw_ostream.h"
#include <utility>
 
namespace fir {
class FirOpBuilder;
class ArrayLoadOp;
 
class ArrayBoxValue;
class BoxValue;
class CharBoxValue;
class CharArrayBoxValue;
class MutableBoxValue;
class PolymorphicValue;
class ProcBoxValue;
 
llvm::raw_ostream &operator<<(llvm::raw_ostream &, const CharBoxValue &);
llvm::raw_ostream &operator<<(llvm::raw_ostream &, const ArrayBoxValue &);
llvm::raw_ostream &operator<<(llvm::raw_ostream &, const CharArrayBoxValue &);
llvm::raw_ostream &operator<<(llvm::raw_ostream &, const ProcBoxValue &);
llvm::raw_ostream &operator<<(llvm::raw_ostream &, const MutableBoxValue &);
llvm::raw_ostream &operator<<(llvm::raw_ostream &, const BoxValue &);
llvm::raw_ostream &operator<<(llvm::raw_ostream &, const PolymorphicValue &);
 
//===----------------------------------------------------------------------===//
//
// Boxed values
//
// Define a set of containers used internally by the lowering bridge to keep
// track of extended values associated with a Fortran subexpression. These
// associations are maintained during the construction of FIR.
//
//===----------------------------------------------------------------------===//

using UnboxedValue = mlir::Value;

class AbstractBox {
public:
  AbstractBox() = delete;
  AbstractBox(mlir::Value addr) : addr{addr} {}

  mlir::Value getAddr() const { return addr; }
 
protected:
  mlir::Value addr;
};

class CharBoxValue : public AbstractBox {
public:
  CharBoxValue(mlir::Value addr, mlir::Value len)
      : AbstractBox{addr}, len{len} {
    if (addr && mlir::isa<fir::BoxCharType>(addr.getType()))
      fir::emitFatalError(addr.getLoc(),
                          "BoxChar should not be in CharBoxValue");
  }
 
  CharBoxValue clone(mlir::Value newBase) const { return {newBase, len}; }

  mlir::Value getBuffer() const { return getAddr(); }
 
  mlir::Value getLen() const { return len; }
 
  friend llvm::raw_ostream &operator<<(llvm::raw_ostream &,
                                       const CharBoxValue &);
  LLVM_DUMP_METHOD void dump() const { llvm::errs() << *this; }
 
protected:
  mlir::Value len;
};

class PolymorphicValue : public AbstractBox {
public:
  PolymorphicValue(mlir::Value addr, mlir::Value sourceBox)
      : AbstractBox{addr}, sourceBox{sourceBox} {}
 
  PolymorphicValue clone(mlir::Value newBase) const {
    return {newBase, sourceBox};
  }
 
  mlir::Value getSourceBox() const { return sourceBox; }
 
  friend llvm::raw_ostream &operator<<(llvm::raw_ostream &,
                                       const PolymorphicValue &);
  LLVM_DUMP_METHOD void dump() const { llvm::errs() << *this; }
 
protected:
  mlir::Value sourceBox;
};

class AbstractArrayBox {
public:
  AbstractArrayBox() = default;
  AbstractArrayBox(llvm::ArrayRef<mlir::Value> extents,
                   llvm::ArrayRef<mlir::Value> lbounds)
      : extents{extents}, lbounds{lbounds} {}
 
  // Every array has extents that describe its shape.
  const llvm::SmallVectorImpl<mlir::Value> &getExtents() const {
    return extents;
  }
 
  // An array expression may have user-defined lower bound values.
  // If this vector is empty, the default in all dimensions in `1`.
  const llvm::SmallVectorImpl<mlir::Value> &getLBounds() const {
    return lbounds;
  }
 
  bool lboundsAllOne() const { return lbounds.empty(); }
  std::size_t rank() const { return extents.size(); }
 
protected:
  llvm::SmallVector<mlir::Value, 4> extents;
  llvm::SmallVector<mlir::Value, 4> lbounds;
};

class ArrayBoxValue : public PolymorphicValue, public AbstractArrayBox {
public:
  ArrayBoxValue(mlir::Value addr, llvm::ArrayRef<mlir::Value> extents,
                llvm::ArrayRef<mlir::Value> lbounds = {},
                mlir::Value sourceBox = {})
      : PolymorphicValue{addr, sourceBox}, AbstractArrayBox{extents, lbounds} {}
 
  ArrayBoxValue clone(mlir::Value newBase) const {
    return {newBase, extents, lbounds};
  }
 
  friend llvm::raw_ostream &operator<<(llvm::raw_ostream &,
                                       const ArrayBoxValue &);
  LLVM_DUMP_METHOD void dump() const { llvm::errs() << *this; }
};

class CharArrayBoxValue : public CharBoxValue, public AbstractArrayBox {
public:
  CharArrayBoxValue(mlir::Value addr, mlir::Value len,
                    llvm::ArrayRef<mlir::Value> extents,
                    llvm::ArrayRef<mlir::Value> lbounds = {})
      : CharBoxValue{addr, len}, AbstractArrayBox{extents, lbounds} {}
 
  CharArrayBoxValue clone(mlir::Value newBase) const {
    return {newBase, len, extents, lbounds};
  }
 
  CharBoxValue cloneElement(mlir::Value newBase) const {
    return {newBase, len};
  }
 
  friend llvm::raw_ostream &operator<<(llvm::raw_ostream &,
                                       const CharArrayBoxValue &);
  LLVM_DUMP_METHOD void dump() const { llvm::errs() << *this; }
};

class ProcBoxValue : public AbstractBox {
public:
  ProcBoxValue(mlir::Value addr, mlir::Value context)
      : AbstractBox{addr}, hostContext{context} {}
 
  ProcBoxValue clone(mlir::Value newBase) const {
    return {newBase, hostContext};
  }
 
  mlir::Value getHostContext() const { return hostContext; }
 
  friend llvm::raw_ostream &operator<<(llvm::raw_ostream &,
                                       const ProcBoxValue &);
  LLVM_DUMP_METHOD void dump() const { llvm::errs() << *this; }
 
protected:
  mlir::Value hostContext;
};

class AbstractIrBox : public AbstractBox, public AbstractArrayBox {
public:
  AbstractIrBox(mlir::Value addr) : AbstractBox{addr} {}
  AbstractIrBox(mlir::Value addr, llvm::ArrayRef<mlir::Value> lbounds,
                llvm::ArrayRef<mlir::Value> extents)
      : AbstractBox{addr}, AbstractArrayBox(extents, lbounds) {}
  fir::BaseBoxType getBoxTy() const {
    auto type = getAddr().getType();
    if (auto pointedTy = fir::dyn_cast_ptrEleTy(type))
      type = pointedTy;
    return mlir::cast<fir::BaseBoxType>(type);
  }

  mlir::Type getBaseTy() const {
    return fir::dyn_cast_ptrOrBoxEleTy(getBoxTy());
  }

  mlir::Type getMemTy() const {
    auto ty = getBoxTy().getEleTy();
    if (fir::isa_ref_type(ty))
      return ty;
    return fir::ReferenceType::get(ty, fir::isa_volatile_type(getBoxTy()));
  }

  mlir::Type getEleTy() const {
    auto type = getBaseTy();
    if (auto seqTy = mlir::dyn_cast<fir::SequenceType>(type))
      return seqTy.getEleTy();
    return type;
  }

  bool hasRank() const { return mlir::isa<fir::SequenceType>(getBaseTy()); }
  bool hasAssumedRank() const {
    auto seqTy = mlir::dyn_cast<fir::SequenceType>(getBaseTy());
    return seqTy && seqTy.hasUnknownShape();
  }

  unsigned rank() const {
    if (auto seqTy = mlir::dyn_cast<fir::SequenceType>(getBaseTy()))
      return seqTy.getDimension();
    return 0;
  }

  bool isCharacter() const { return fir::isa_char(getEleTy()); }

  bool isDerived() const { return mlir::isa<fir::RecordType>(getEleTy()); }
 
  bool isDerivedWithLenParameters() const {
    return fir::isRecordWithTypeParameters(getEleTy());
  }

  bool isPolymorphic() const { return fir::isPolymorphicType(getBoxTy()); }

  bool isUnlimitedPolymorphic() const {
    return fir::isUnlimitedPolymorphicType(getBoxTy());
  }
 
  unsigned corank() const { return fir::getBoxCorank(getBoxTy()); }
};

class BoxValue : public AbstractIrBox {
public:
  BoxValue(mlir::Value addr) : AbstractIrBox{addr} { assert(verify()); }
  BoxValue(mlir::Value addr, llvm::ArrayRef<mlir::Value> lbounds,
           llvm::ArrayRef<mlir::Value> explicitParams,
           llvm::ArrayRef<mlir::Value> explicitExtents = {})
      : AbstractIrBox{addr, lbounds, explicitExtents},
        explicitParams{explicitParams} {
    assert(verify());
  }
  // TODO: check contiguous attribute of addr
  bool isContiguous() const { return false; }
 
  // Replace the fir.box, keeping any non-deferred parameters.
  BoxValue clone(mlir::Value newBox) const {
    return {newBox, lbounds, explicitParams, extents};
  }
 
  friend llvm::raw_ostream &operator<<(llvm::raw_ostream &, const BoxValue &);
  LLVM_DUMP_METHOD void dump() const { llvm::errs() << *this; }
 
  llvm::ArrayRef<mlir::Value> getLBounds() const { return lbounds; }
 
  // The extents member is not guaranteed to be field for arrays. It is only
  // guaranteed to be field for explicit shape arrays. In general,
  // explicit-shape will not come as descriptors, so this field will be empty in
  // most cases. The exception are derived types with LEN parameters and
  // polymorphic dummy argument arrays. It may be possible for the explicit
  // extents to conflict with the shape information that is in the box according
  // to 15.5.2.11 sequence association rules.
  llvm::ArrayRef<mlir::Value> getExplicitExtents() const { return extents; }
 
  llvm::ArrayRef<mlir::Value> getExplicitParameters() const {
    return explicitParams;
  }
 
protected:
  // Verify constructor invariants.
  bool verify() const;
 
  // Only field when the BoxValue has explicit LEN parameters.
  // Otherwise, the LEN parameters are in the fir.box.
  llvm::SmallVector<mlir::Value, 2> explicitParams;
};

class MutableProperties {
public:
  bool isEmpty() const { return !addr; }
  mlir::Value addr;
  llvm::SmallVector<mlir::Value, 2> extents;
  llvm::SmallVector<mlir::Value, 2> lbounds;
  llvm::SmallVector<mlir::Value, 2> deferredParams;
};

class MutableBoxValue : public AbstractIrBox {
public:
  MutableBoxValue(mlir::Value addr, mlir::ValueRange lenParameters,
                  MutableProperties mutableProperties)
      : AbstractIrBox(addr), lenParams{lenParameters.begin(),
                                       lenParameters.end()},
        mutableProperties{mutableProperties} {
    // Currently only accepts fir.(ref/ptr/heap)<fir.box<type>> mlir::Value for
    // the address. This may change if we accept
    // fir.(ref/ptr/heap)<fir.heap<type>> for scalar without LEN parameters.
    assert(verify() &&
           "MutableBoxValue requires mem ref to fir.box<fir.[heap|ptr]<type>>");
  }

  bool isPointer() const {
    return mlir::isa<fir::PointerType>(getBoxTy().getEleTy());
  }

  bool isAllocatable() const {
    return mlir::isa<fir::HeapType>(getBoxTy().getEleTy());
  }
  // Replace the fir.ref<fir.box>, keeping any non-deferred parameters.
  MutableBoxValue clone(mlir::Value newBox) const {
    return {newBox, lenParams, mutableProperties};
  }
  bool hasNonDeferredLenParams() const { return !lenParams.empty(); }
  llvm::ArrayRef<mlir::Value> nonDeferredLenParams() const { return lenParams; }
  friend llvm::raw_ostream &operator<<(llvm::raw_ostream &,
                                       const MutableBoxValue &);
  LLVM_DUMP_METHOD void dump() const { llvm::errs() << *this; }

  bool isDescribedByVariables() const { return !mutableProperties.isEmpty(); }
 
  const MutableProperties &getMutableProperties() const {
    return mutableProperties;
  }
 
protected:
  bool verify() const;
  llvm::SmallVector<mlir::Value, 2> lenParams;
  MutableProperties mutableProperties;
};
 
class ExtendedValue;

mlir::Value getBase(const ExtendedValue &exv);

mlir::Value getLen(const ExtendedValue &exv);

llvm::raw_ostream &operator<<(llvm::raw_ostream &, const ExtendedValue &);

ExtendedValue substBase(const ExtendedValue &exv, mlir::Value base);

bool isArray(const ExtendedValue &exv);

llvm::SmallVector<mlir::Value> getTypeParams(const ExtendedValue &exv);
 
//===----------------------------------------------------------------------===//
// Functions that may generate IR to recover properties from extended values.
//===----------------------------------------------------------------------===//
namespace factory {

llvm::SmallVector<mlir::Value> getTypeParams(mlir::Location loc,
                                             FirOpBuilder &builder,
                                             const ExtendedValue &exv);

llvm::SmallVector<mlir::Value>
getTypeParams(mlir::Location loc, FirOpBuilder &builder, ArrayLoadOp load);
 
// The generalized function to get a vector of extents is
llvm::SmallVector<mlir::Value>
getExtents(mlir::Location loc, FirOpBuilder &builder, const ExtendedValue &box);

mlir::Value getExtentAtDimension(mlir::Location loc, FirOpBuilder &builder,
                                 const ExtendedValue &exv, unsigned dim);
 
} // namespace factory

class ExtendedValue : public details::matcher<ExtendedValue> {
public:
  using VT =
      std::variant<UnboxedValue, CharBoxValue, ArrayBoxValue, CharArrayBoxValue,
                   ProcBoxValue, BoxValue, MutableBoxValue, PolymorphicValue>;
 
  ExtendedValue() : box{UnboxedValue{}} {}
  template <typename A, typename = std::enable_if_t<
                            !std::is_same_v<std::decay_t<A>, ExtendedValue>>>
  constexpr ExtendedValue(A &&a) : box{std::forward<A>(a)} {
    if (const auto *b = getUnboxed()) {
      if (*b) {
        auto type = b->getType();
        if (mlir::isa<fir::BoxCharType>(type))
          fir::emitFatalError(b->getLoc(), "BoxChar should be unboxed");
        type = fir::unwrapSequenceType(fir::unwrapRefType(type));
        if (fir::isa_char(type))
          fir::emitFatalError(b->getLoc(),
                              "character buffer should be in CharBoxValue");
      }
    }
  }
 
  template <typename A>
  constexpr const A *getBoxOf() const {
    return std::get_if<A>(&box);
  }
 
  constexpr const CharBoxValue *getCharBox() const {
    return getBoxOf<CharBoxValue>();
  }
 
  constexpr const UnboxedValue *getUnboxed() const {
    return getBoxOf<UnboxedValue>();
  }
 
  unsigned rank() const {
    return match([](const fir::UnboxedValue &box) -> unsigned { return 0; },
                 [](const fir::CharBoxValue &box) -> unsigned { return 0; },
                 [](const fir::ProcBoxValue &box) -> unsigned { return 0; },
                 [](const fir::PolymorphicValue &box) -> unsigned { return 0; },
                 [](const auto &box) -> unsigned { return box.rank(); });
  }
 
  bool isPolymorphic() const {
    return match([](const fir::PolymorphicValue &box) -> bool { return true; },
                 [](const fir::ArrayBoxValue &box) -> bool {
                   return box.getSourceBox() ? true : false;
                 },
                 [](const auto &box) -> bool { return false; });
  }
 
  bool hasAssumedRank() const {
    return match(
        [](const fir::BoxValue &box) -> bool { return box.hasAssumedRank(); },
        [](const fir::MutableBoxValue &box) -> bool {
          return box.hasAssumedRank();
        },
        [](const auto &box) -> bool { return false; });
  }

  LLVM_DUMP_METHOD void dump() const { llvm::errs() << *this << '\n'; }
 
  friend llvm::raw_ostream &operator<<(llvm::raw_ostream &,
                                       const ExtendedValue &);
 
  const VT &matchee() const { return box; }
 
private:
  VT box;
};

inline bool isUnboxedValue(const ExtendedValue &exv) {
  return exv.match(
      [](const fir::UnboxedValue &box) { return box ? true : false; },
      [](const auto &) { return false; });
}

inline mlir::Type getBaseTypeOf(const ExtendedValue &exv) {
  return exv.match(
      [](const fir::MutableBoxValue &box) { return box.getBaseTy(); },
      [](const fir::BoxValue &box) { return box.getBaseTy(); },
      [&](const auto &) {
        return fir::unwrapRefType(fir::getBase(exv).getType());
      });
}

inline mlir::Type getElementTypeOf(const ExtendedValue &exv) {
  return fir::unwrapSequenceType(getBaseTypeOf(exv));
}

inline bool isDerivedWithLenParameters(const ExtendedValue &exv) {
  return fir::isRecordWithTypeParameters(getElementTypeOf(exv));
}
 
} // namespace fir
 
#endif // FORTRAN_OPTIMIZER_BUILDER_BOXVALUE_H