SymbolMap.h

//===-- SymbolMap.h -- lowering internal symbol map -------------*- 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_LOWER_SYMBOLMAP_H
#define FORTRAN_LOWER_SYMBOLMAP_H
 
#include "flang/Lower/Support/Utils.h"
#include "flang/Optimizer/Builder/BoxValue.h"
#include "flang/Optimizer/Dialect/FIRType.h"
#include "flang/Optimizer/Dialect/FortranVariableInterface.h"
#include "flang/Optimizer/Support/Matcher.h"
#include "flang/Semantics/symbol.h"
#include "mlir/IR/Value.h"
#include "llvm/ADT/ArrayRef.h"
#include "llvm/ADT/DenseMap.h"
#include "llvm/ADT/SmallVector.h"
#include "llvm/Support/Compiler.h"
#include <optional>
 
namespace Fortran::lower {
 
struct SymbolBox;
class SymMap;
llvm::raw_ostream &operator<<(llvm::raw_ostream &os, const SymbolBox &symMap);
llvm::raw_ostream &operator<<(llvm::raw_ostream &os, const SymMap &symMap);
 
//===----------------------------------------------------------------------===//
// Symbol information
//===----------------------------------------------------------------------===//

struct SymbolBox : public fir::details::matcher<SymbolBox> {
  // For lookups that fail, have a monostate
  using None = std::monostate;
 
  // Trivial intrinsic type
  using Intrinsic = fir::AbstractBox;
 
  // Array variable that uses bounds notation
  using FullDim = fir::ArrayBoxValue;
 
  // CHARACTER type variable with its dependent type LEN parameter
  using Char = fir::CharBoxValue;
 
  // CHARACTER array variable using bounds notation
  using CharFullDim = fir::CharArrayBoxValue;
 
  // Pointer or allocatable variable
  using PointerOrAllocatable = fir::MutableBoxValue;
 
  // Non pointer/allocatable variable that must be tracked with
  // a fir.box (either because it is not contiguous, or assumed rank, or assumed
  // type, or polymorphic, or because the fir.box is describing an optional
  // value and cannot be read into one of the other category when lowering the
  // symbol).
  using Box = fir::BoxValue;
 
  using VT =
      std::variant<Intrinsic, FullDim, Char, CharFullDim, PointerOrAllocatable,
                   Box, fir::FortranVariableOpInterface, None>;
 
  //===--------------------------------------------------------------------===//
  // Constructors
  //===--------------------------------------------------------------------===//
 
  SymbolBox() : box{None{}} {}
  template <typename A>
  SymbolBox(const A &x) : box{x} {}
 
  explicit operator bool() const { return !std::holds_alternative<None>(box); }
 
  //===--------------------------------------------------------------------===//
  // Accessors
  //===--------------------------------------------------------------------===//

  mlir::Value getAddr() const {
    return match([](const None &) { return mlir::Value{}; },
                 [](const fir::FortranVariableOpInterface &x) {
                   return fir::FortranVariableOpInterface(x).getBase();
                 },
                 [](const auto &x) { return x.getAddr(); });
  }
 
  std::optional<fir::FortranVariableOpInterface>
  getIfFortranVariableOpInterface() {
    return match(
        [](const fir::FortranVariableOpInterface &x)
            -> std::optional<fir::FortranVariableOpInterface> { return x; },
        [](const auto &x) -> std::optional<fir::FortranVariableOpInterface> {
          return std::nullopt;
        });
  }

  template <typename ON, typename RT>
  constexpr RT apply(RT (&&func)(const ON &)) const {
    if (auto *x = std::get_if<ON>(&box))
      return func(*x);
    return RT{};
  }
 
  const VT &matchee() const { return box; }
 
  friend llvm::raw_ostream &operator<<(llvm::raw_ostream &os,
                                       const SymbolBox &symBox);

  LLVM_DUMP_METHOD void dump() const;
 
private:
  VT box;
};

class ComponentMap {
public:
  void insert(const Fortran::evaluate::Component &component,
              fir::FortranVariableOpInterface definingOp) {
    auto iter = componentMap.find(&component);
    if (iter != componentMap.end()) {
      iter->second = definingOp;
      return;
    }
    componentStorage.push_back(
        std::make_unique<Fortran::evaluate::Component>(component));
    componentMap.insert({componentStorage.back().get(), definingOp});
  }
 
  std::optional<fir::FortranVariableOpInterface>
  lookup(const Fortran::evaluate::Component *component) const {
    auto iter = componentMap.find(component);
    if (iter != componentMap.end())
      return iter->second;
    return std::nullopt;
  }
 
  LLVM_DUMP_METHOD void dump() const;
 
private:
  llvm::DenseMap<const Fortran::evaluate::Component *,
                 fir::FortranVariableOpInterface>
      componentMap;
  llvm::SmallVector<std::unique_ptr<Fortran::evaluate::Component>>
      componentStorage;
};
 
//===----------------------------------------------------------------------===//
// Map of symbol information
//===----------------------------------------------------------------------===//

class SymMap {
public:
  using AcDoVar = llvm::StringRef;
  using StorageDesc = std::pair<mlir::Value, std::uint64_t>;
 
  SymMap() { pushScope(); }
  SymMap(const SymMap &) = delete;
 
  void pushScope() {
    symbolMapStack.emplace_back();
    deviceSymbolMapStack.emplace_back();
    storageMapStack.emplace_back();
    componentMapStack.emplace_back();
  }
  void popScope() {
    symbolMapStack.pop_back();
    assert(symbolMapStack.size() >= 1);
    deviceSymbolMapStack.pop_back();
    assert(deviceSymbolMapStack.size() >= 1);
    storageMapStack.pop_back();
    assert(storageMapStack.size() >= 1);
    componentMapStack.pop_back();
    assert(componentMapStack.size() >= 1);
  }

  void addSymbol(semantics::SymbolRef sym, const fir::ExtendedValue &ext,
                 bool force = false);

  void addSymbol(semantics::SymbolRef sym, mlir::Value value,
                 bool force = false) {
    makeSym(sym, SymbolBox::Intrinsic(value), force);
  }

  void addCharSymbol(semantics::SymbolRef sym, mlir::Value value,
                     mlir::Value len, bool force = false) {
    makeSym(sym, SymbolBox::Char(value, len), force);
  }
  void addCharSymbol(semantics::SymbolRef sym, const SymbolBox::Char &value,
                     bool force = false) {
    makeSym(sym, value, force);
  }

  void addSymbolWithShape(semantics::SymbolRef sym, mlir::Value value,
                          llvm::ArrayRef<mlir::Value> shape,
                          bool force = false) {
    makeSym(sym, SymbolBox::FullDim(value, shape), force);
  }
  void addSymbolWithShape(semantics::SymbolRef sym,
                          const SymbolBox::FullDim &value, bool force = false) {
    makeSym(sym, value, force);
  }

  void addCharSymbolWithShape(semantics::SymbolRef sym, mlir::Value value,
                              mlir::Value len,
                              llvm::ArrayRef<mlir::Value> shape,
                              bool force = false) {
    makeSym(sym, SymbolBox::CharFullDim(value, len, shape), force);
  }
  void addCharSymbolWithShape(semantics::SymbolRef sym,
                              const SymbolBox::CharFullDim &value,
                              bool force = false) {
    makeSym(sym, value, force);
  }

  void addSymbolWithBounds(semantics::SymbolRef sym, mlir::Value value,
                           llvm::ArrayRef<mlir::Value> extents,
                           llvm::ArrayRef<mlir::Value> lbounds,
                           bool force = false) {
    makeSym(sym, SymbolBox::FullDim(value, extents, lbounds), force);
  }
  void addSymbolWithBounds(semantics::SymbolRef sym,
                           const SymbolBox::FullDim &value,
                           bool force = false) {
    makeSym(sym, value, force);
  }

  void addCharSymbolWithBounds(semantics::SymbolRef sym, mlir::Value value,
                               mlir::Value len,
                               llvm::ArrayRef<mlir::Value> extents,
                               llvm::ArrayRef<mlir::Value> lbounds,
                               bool force = false) {
    makeSym(sym, SymbolBox::CharFullDim(value, len, extents, lbounds), force);
  }
  void addCharSymbolWithBounds(semantics::SymbolRef sym,
                               const SymbolBox::CharFullDim &value,
                               bool force = false) {
    makeSym(sym, value, force);
  }
 
  void addAllocatableOrPointer(semantics::SymbolRef sym,
                               fir::MutableBoxValue box, bool force = false) {
    makeSym(sym, box, force);
  }
 
  void addBoxSymbol(semantics::SymbolRef sym, mlir::Value irBox,
                    llvm::ArrayRef<mlir::Value> lbounds,
                    llvm::ArrayRef<mlir::Value> explicitParams,
                    llvm::ArrayRef<mlir::Value> explicitExtents,
                    bool force = false) {
    makeSym(sym,
            SymbolBox::Box(irBox, lbounds, explicitParams, explicitExtents),
            force);
  }
  void addBoxSymbol(semantics::SymbolRef sym, const SymbolBox::Box &value,
                    bool force = false) {
    makeSym(sym, value, force);
  }

  SymbolBox lookupSymbol(semantics::SymbolRef sym);
  SymbolBox lookupSymbol(const semantics::Symbol *sym) {
    return lookupSymbol(*sym);
  }

  void addDeviceVariableDefinition(semantics::SymbolRef symRef,
                                   fir::FortranVariableOpInterface definingOp,
                                   bool force = false) {
    makeDeviceSym(symRef, SymbolBox(definingOp), force);
  }
  SymbolBox lookupDeviceSymbol(semantics::SymbolRef);
  bool copyDeviceBindingToCurrentScope(semantics::SymbolRef symRef) {
    if (SymbolBox box{lookupDeviceSymbol(symRef)}) {
      makeSym(symRef, box, /*force=*/true);
      return true;
    }
    return false;
  }

  const semantics::Symbol *lookupSymbolByName(llvm::StringRef symName);

  SymbolBox shallowLookupSymbol(semantics::SymbolRef sym);
  SymbolBox shallowLookupSymbol(const semantics::Symbol *sym) {
    return shallowLookupSymbol(*sym);
  }

  SymbolBox lookupOneLevelUpSymbol(semantics::SymbolRef sym);
  SymbolBox lookupOneLevelUpSymbol(const semantics::Symbol *sym) {
    return lookupOneLevelUpSymbol(*sym);
  }

  void pushImpliedDoBinding(AcDoVar var, mlir::Value value) {
    impliedDoStack.emplace_back(var, value);
  }

  void popImpliedDoBinding() {
    assert(!impliedDoStack.empty());
    impliedDoStack.pop_back();
  }

  mlir::Value lookupImpliedDo(AcDoVar var);

  void clear() {
    symbolMapStack.clear();
    symbolMapStack.emplace_back();
    deviceSymbolMapStack.clear();
    deviceSymbolMapStack.emplace_back();
    assert(symbolMapStack.size() == 1);
    impliedDoStack.clear();
    storageMapStack.clear();
    storageMapStack.emplace_back();
    componentMapStack.clear();
    componentMapStack.emplace_back();
  }
 
  friend llvm::raw_ostream &operator<<(llvm::raw_ostream &os,
                                       const SymMap &symMap);

  LLVM_DUMP_METHOD void dump() const;
 
  void addVariableDefinition(semantics::SymbolRef symRef,
                             fir::FortranVariableOpInterface definingOp,
                             bool force = false) {
    makeSym(symRef, SymbolBox(definingOp), force);
  }
 
  void copySymbolBinding(semantics::SymbolRef src,
                         semantics::SymbolRef target) {
    auto symBox = lookupSymbol(src);
    assert(symBox && "source binding does not exists");
    makeSym(target, symBox, /*force=*/false);
  }
 
  std::optional<fir::FortranVariableOpInterface>
  lookupVariableDefinition(semantics::SymbolRef sym) {
    if (auto symBox = lookupSymbol(sym))
      return symBox.getIfFortranVariableOpInterface();
    return std::nullopt;
  }

  void addComponentOverride(const Fortran::evaluate::Component &component,
                            fir::FortranVariableOpInterface definingOp) {
    assert(!componentMapStack.empty() && "component map stack is empty");
    if (!componentMapStack.back())
      componentMapStack.back() = std::make_unique<ComponentMap>();
    componentMapStack.back().value()->insert(component, definingOp);
  }

  std::optional<fir::FortranVariableOpInterface>
  lookupComponentOverride(const Fortran::evaluate::Component &component) const {
    for (auto jmap = componentMapStack.rbegin(),
              jend = componentMapStack.rend();
         jmap != jend; ++jmap) {
      if (*jmap) {
        auto iter = (**jmap)->lookup(&component);
        if (iter != std::nullopt)
          return iter;
      }
    }
    return std::nullopt;
  }

  void registerStorage(semantics::SymbolRef sym, StorageDesc storage);
  StorageDesc lookupStorage(semantics::SymbolRef sym);
  StorageDesc lookupStorage(const semantics::Symbol *sym) {
    return lookupStorage(*sym);
  }
 
private:
  void makeSym(semantics::SymbolRef symRef, const SymbolBox &box,
               bool force = false) {
    auto *sym = symRef->HasLocalLocality() ? &*symRef : &symRef->GetUltimate();
    if (force)
      symbolMapStack.back().erase(sym);
    assert(box && "cannot add an undefined symbol box");
    symbolMapStack.back().try_emplace(sym, box);
  }
 
  void makeDeviceSym(semantics::SymbolRef symRef, const SymbolBox &box,
                     bool force = false) {
    auto *sym = symRef->HasLocalLocality() ? &*symRef : &symRef->GetUltimate();
    if (force)
      deviceSymbolMapStack.back().erase(sym);
    assert(box && "cannot add an undefined device symbol box");
    deviceSymbolMapStack.back().try_emplace(sym, box);
  }
 
  llvm::SmallVector<llvm::DenseMap<const semantics::Symbol *, SymbolBox>>
      symbolMapStack;
  llvm::SmallVector<llvm::DenseMap<const semantics::Symbol *, SymbolBox>>
      deviceSymbolMapStack;
 
  // Implied DO induction variables are not represented as Se::Symbol in
  // Ev::Expr. Keep the variable markers in their own stack.
  llvm::SmallVector<std::pair<AcDoVar, mlir::Value>> impliedDoStack;
 
  // A stack of maps between the symbols and their storage descriptors.
  llvm::SmallVector<llvm::DenseMap<const semantics::Symbol *, StorageDesc>>
      storageMapStack;
 
  // A stack of maps from front-end component references to the FIR variables
  // that should be used to implement them. This allows overriding component
  // references in specific lowering contexts.
  llvm::SmallVector<std::optional<std::unique_ptr<ComponentMap>>>
      componentMapStack;
};

class SymMapScope {
public:
  explicit SymMapScope(SymMap &map) : map(map) { map.pushScope(); }
  ~SymMapScope() { map.popScope(); }
 
private:
  SymMap &map;
};
 
} // namespace Fortran::lower
 
#endif // FORTRAN_LOWER_SYMBOLMAP_H