parse-tree-visitor.h

//===-- include/flang/Parser/parse-tree-visitor.h ---------------*- 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
//
//===----------------------------------------------------------------------===//
 
#ifndef FORTRAN_PARSER_PARSE_TREE_VISITOR_H_
#define FORTRAN_PARSER_PARSE_TREE_VISITOR_H_
 
#include "parse-tree.h"
#include "tools.h"
#include "flang/Common/enum-set.h"
#include "flang/Common/visit.h"
#include <cstddef>
#include <optional>
#include <tuple>
#include <utility>
#include <variant>
#include <vector>

 
namespace Fortran::parser {
 
template <typename A, typename V> void Walk(const A &x, V &visitor);
template <typename A, typename M> void Walk(A &x, M &mutator);
 
namespace detail {
// A number of the Walk functions below call other Walk functions. Define
// a dummy class, and put all of them in it to ensure that name lookup for
// Walk considers all overloads (not just those defined prior to the call
// to Walk).
struct ParseTreeVisitorLookupScope {
  // Default case for visitation of non-class data members, strings, and
  // any other non-decomposable values.
  template <typename A, typename V>
  static std::enable_if_t<!std::is_class_v<A> || common::IsEnumSet<A> ||
      std::is_same_v<std::string, A> || std::is_same_v<CharBlock, A>>
  Walk(const A &x, V &visitor) {
    if (visitor.Pre(x)) {
      visitor.Post(x);
    }
  }
  template <typename A, typename M>
  static std::enable_if_t<!std::is_class_v<A> ||
      std::is_same_v<std::string, A> || std::is_same_v<CharBlock, A>>
  Walk(A &x, M &mutator) {
    if (mutator.Pre(x)) {
      mutator.Post(x);
    }
  }
 
  template <typename A, typename VM>
  static void WalkSource(A &x, VM &visitorOrMutator) {
    if constexpr (HasSource<A>::value) {
      Walk(x.source, visitorOrMutator);
    }
  }
 
  // Traversal of needed STL template classes (optional, list, tuple, variant)
  // For most lists, just traverse the elements; but when a list constitutes
  // a Block (i.e., std::list<ExecutionPartConstruct>), also invoke the
  // visitor/mutator on the list itself.
  template <typename T, typename V>
  static void Walk(const std::list<T> &x, V &visitor) {
    for (const auto &elem : x) {
      Walk(elem, visitor);
    }
  }
  template <typename T, typename M>
  static void Walk(std::list<T> &x, M &mutator) {
    for (auto &elem : x) {
      Walk(elem, mutator);
    }
  }
  template <typename V> static void Walk(const Block &x, V &visitor) {
    if (visitor.Pre(x)) {
      for (const auto &elem : x) {
        Walk(elem, visitor);
      }
      visitor.Post(x);
    }
  }
  template <typename M> static void Walk(Block &x, M &mutator) {
    if (mutator.Pre(x)) {
      for (auto &elem : x) {
        Walk(elem, mutator);
      }
      mutator.Post(x);
    }
  }
  template <typename T, typename V>
  static void Walk(const std::optional<T> &x, V &visitor) {
    if (x) {
      Walk(*x, visitor);
    }
  }
  template <typename T, typename M>
  static void Walk(std::optional<T> &x, M &mutator) {
    if (x) {
      Walk(*x, mutator);
    }
  }
  template <std::size_t I = 0, typename Func, typename T>
  static void ForEachInTuple(const T &tuple, Func func) {
    func(std::get<I>(tuple));
    if constexpr (I + 1 < std::tuple_size_v<T>) {
      ForEachInTuple<I + 1>(tuple, func);
    }
  }
  template <typename V, typename... A>
  static void Walk(const std::tuple<A...> &x, V &visitor) {
    if (sizeof...(A) > 0) {
      if (visitor.Pre(x)) {
        ForEachInTuple(x, [&](const auto &y) { Walk(y, visitor); });
        visitor.Post(x);
      }
    }
  }
  template <std::size_t I = 0, typename Func, typename T>
  static void ForEachInTuple(T &tuple, Func func) {
    func(std::get<I>(tuple));
    if constexpr (I + 1 < std::tuple_size_v<T>) {
      ForEachInTuple<I + 1>(tuple, func);
    }
  }
  template <typename M, typename... A>
  static void Walk(std::tuple<A...> &x, M &mutator) {
    if (sizeof...(A) > 0) {
      if (mutator.Pre(x)) {
        ForEachInTuple(x, [&](auto &y) { Walk(y, mutator); });
        mutator.Post(x);
      }
    }
  }
  template <typename V, typename... A>
  static void Walk(const std::variant<A...> &x, V &visitor) {
    if (visitor.Pre(x)) {
      common::visit([&](const auto &y) { Walk(y, visitor); }, x);
      visitor.Post(x);
    }
  }
  template <typename M, typename... A>
  static void Walk(std::variant<A...> &x, M &mutator) {
    if (mutator.Pre(x)) {
      common::visit([&](auto &y) { Walk(y, mutator); }, x);
      mutator.Post(x);
    }
  }
  template <typename A, typename B, typename V>
  static void Walk(const std::pair<A, B> &x, V &visitor) {
    if (visitor.Pre(x)) {
      Walk(x.first, visitor);
      Walk(x.second, visitor);
    }
  }
  template <typename A, typename B, typename M>
  static void Walk(std::pair<A, B> &x, M &mutator) {
    if (mutator.Pre(x)) {
      Walk(x.first, mutator);
      Walk(x.second, mutator);
    }
  }
 
  // Trait-determined traversal of empty, tuple, union, wrapper,
  // and constraint-checking classes.
  template <typename A, typename V>
  static std::enable_if_t<EmptyTrait<A>> Walk(const A &x, V &visitor) {
    if (visitor.Pre(x)) {
      WalkSource(x, visitor);
      visitor.Post(x);
    }
  }
  template <typename A, typename M>
  static std::enable_if_t<EmptyTrait<A>> Walk(A &x, M &mutator) {
    if (mutator.Pre(x)) {
      WalkSource(x, mutator);
      mutator.Post(x);
    }
  }
 
  template <typename A, typename V>
  static std::enable_if_t<TupleTrait<A>> Walk(const A &x, V &visitor) {
    if (visitor.Pre(x)) {
      WalkSource(x, visitor);
      Walk(x.t, visitor);
      visitor.Post(x);
    }
  }
  template <typename A, typename M>
  static std::enable_if_t<TupleTrait<A>> Walk(A &x, M &mutator) {
    if (mutator.Pre(x)) {
      WalkSource(x, mutator);
      Walk(x.t, mutator);
      mutator.Post(x);
    }
  }
 
  template <typename A, typename V>
  static std::enable_if_t<UnionTrait<A>> Walk(const A &x, V &visitor) {
    if (visitor.Pre(x)) {
      WalkSource(x, visitor);
      Walk(x.u, visitor);
      visitor.Post(x);
    }
  }
  template <typename A, typename M>
  static std::enable_if_t<UnionTrait<A>> Walk(A &x, M &mutator) {
    if (mutator.Pre(x)) {
      WalkSource(x, mutator);
      Walk(x.u, mutator);
      mutator.Post(x);
    }
  }
 
  template <typename A, typename V>
  static std::enable_if_t<WrapperTrait<A>> Walk(const A &x, V &visitor) {
    if (visitor.Pre(x)) {
      WalkSource(x, visitor);
      Walk(x.v, visitor);
      visitor.Post(x);
    }
  }
  template <typename A, typename M>
  static std::enable_if_t<WrapperTrait<A>> Walk(A &x, M &mutator) {
    if (mutator.Pre(x)) {
      WalkSource(x, mutator);
      Walk(x.v, mutator);
      mutator.Post(x);
    }
  }
 
  template <typename A, typename V>
  static std::enable_if_t<ConstraintTrait<A>> Walk(const A &x, V &visitor) {
    if (visitor.Pre(x)) {
      WalkSource(x, visitor);
      Walk(x.thing, visitor);
      visitor.Post(x);
    }
  }
  template <typename A, typename M>
  static std::enable_if_t<ConstraintTrait<A>> Walk(A &x, M &mutator) {
    if (mutator.Pre(x)) {
      WalkSource(x, mutator);
      Walk(x.thing, mutator);
      mutator.Post(x);
    }
  }
 
  template <typename T, typename V>
  static void Walk(const common::Indirection<T> &x, V &visitor) {
    Walk(x.value(), visitor);
  }
  template <typename T, typename M>
  static void Walk(common::Indirection<T> &x, M &mutator) {
    Walk(x.value(), mutator);
  }
 
  template <typename T, typename V>
  static void Walk(const Statement<T> &x, V &visitor) {
    if (visitor.Pre(x)) {
      // N.B. The label, if any, is not visited.
      Walk(x.source, visitor);
      Walk(x.statement, visitor);
      visitor.Post(x);
    }
  }
  template <typename T, typename M>
  static void Walk(Statement<T> &x, M &mutator) {
    if (mutator.Pre(x)) {
      // N.B. The label, if any, is not visited.
      Walk(x.source, mutator);
      Walk(x.statement, mutator);
      mutator.Post(x);
    }
  }
 
  template <typename T, typename V>
  static void Walk(const UnlabeledStatement<T> &x, V &visitor) {
    if (visitor.Pre(x)) {
      Walk(x.source, visitor);
      Walk(x.statement, visitor);
      visitor.Post(x);
    }
  }
  template <typename T, typename M>
  static void Walk(UnlabeledStatement<T> &x, M &mutator) {
    if (mutator.Pre(x)) {
      Walk(x.source, mutator);
      Walk(x.statement, mutator);
      mutator.Post(x);
    }
  }
 
  template <typename V> static void Walk(const Name &x, V &visitor) {
    if (visitor.Pre(x)) {
      Walk(x.source, visitor);
      visitor.Post(x);
    }
  }
  template <typename M> static void Walk(Name &x, M &mutator) {
    if (mutator.Pre(x)) {
      Walk(x.source, mutator);
      mutator.Post(x);
    }
  }
 
  // Expr traversal uses iteration rather than recursion to avoid
  // blowing out the stack on very deep expression parse trees.
  // It replaces implementations that looked like:
  //   template <typename V> static void Walk(const Expr &x, V visitor) {
  //     if (visitor.Pre(x)) {      // Pre on the Expr
  //       Walk(x.source, visitor);
  //       // Pre on the operator, walk the operands, Post on operator
  //       Walk(x.u, visitor);
  //       visitor.Post(x);         // Post on the Expr
  //     }
  //   }
  template <typename A, typename V, typename UNARY, typename BINARY>
  static void IterativeWalk(A &start, V &visitor) {
    struct ExprWorkList {
      ExprWorkList(A &x) : expr(&x) {}
      bool doPostExpr{false}, doPostOpr{false};
      A *expr;
    };
    std::vector<ExprWorkList> stack;
    stack.emplace_back(start);
    do {
      A &expr{*stack.back().expr};
      if (stack.back().doPostOpr) {
        stack.back().doPostOpr = false;
        common::visit([&visitor](auto &y) { visitor.Post(y); }, expr.u);
      } else if (stack.back().doPostExpr) {
        visitor.Post(expr);
        stack.pop_back();
      } else if (!visitor.Pre(expr)) {
        stack.pop_back();
      } else {
        stack.back().doPostExpr = true;
        Walk(expr.source, visitor);
        UNARY *unary{nullptr};
        BINARY *binary{nullptr};
        common::visit(
            [&unary, &binary](auto &y) {
              if constexpr (std::is_convertible_v<decltype(&y), UNARY *>) {
                unary = &y;
              } else if constexpr (std::is_convertible_v<decltype(&y),
                                       BINARY *>) {
                binary = &y;
              }
            },
            expr.u);
        if (!unary && !binary) {
          Walk(expr.u, visitor);
        } else if (common::visit([&visitor](auto &y) { return visitor.Pre(y); },
                       expr.u)) {
          stack.back().doPostOpr = true;
          if (unary) {
            stack.emplace_back(unary->v.value());
          } else {
            stack.emplace_back(std::get<1>(binary->t).value());
            stack.emplace_back(std::get<0>(binary->t).value());
          }
        }
      }
    } while (!stack.empty());
  }
  template <typename V> static void Walk(const Expr &x, V &visitor) {
    IterativeWalk<const Expr, V, const Expr::IntrinsicUnary,
        const Expr::IntrinsicBinary>(x, visitor);
  }
  template <typename M> static void Walk(Expr &x, M &mutator) {
    IterativeWalk<Expr, M, Expr::IntrinsicUnary, Expr::IntrinsicBinary>(
        x, mutator);
  }
 
  template <typename V> static void Walk(const ReadStmt &x, V &visitor) {
    if (visitor.Pre(x)) {
      Walk(x.iounit, visitor);
      Walk(x.format, visitor);
      Walk(x.controls, visitor);
      Walk(x.items, visitor);
      visitor.Post(x);
    }
  }
  template <typename M> static void Walk(ReadStmt &x, M &mutator) {
    if (mutator.Pre(x)) {
      Walk(x.iounit, mutator);
      Walk(x.format, mutator);
      Walk(x.controls, mutator);
      Walk(x.items, mutator);
      mutator.Post(x);
    }
  }
  template <typename V> static void Walk(const UseStmt &x, V &visitor) {
    if (visitor.Pre(x)) {
      Walk(x.nature, visitor);
      Walk(x.moduleName, visitor);
      Walk(x.u, visitor);
      visitor.Post(x);
    }
  }
  template <typename M> static void Walk(UseStmt &x, M &mutator) {
    if (mutator.Pre(x)) {
      Walk(x.nature, mutator);
      Walk(x.moduleName, mutator);
      Walk(x.u, mutator);
      mutator.Post(x);
    }
  }
  template <typename V> static void Walk(const WriteStmt &x, V &visitor) {
    if (visitor.Pre(x)) {
      Walk(x.iounit, visitor);
      Walk(x.format, visitor);
      Walk(x.controls, visitor);
      Walk(x.items, visitor);
      visitor.Post(x);
    }
  }
  template <typename M> static void Walk(WriteStmt &x, M &mutator) {
    if (mutator.Pre(x)) {
      Walk(x.iounit, mutator);
      Walk(x.format, mutator);
      Walk(x.controls, mutator);
      Walk(x.items, mutator);
      mutator.Post(x);
    }
  }
  template <typename V>
  static void Walk(const format::ControlEditDesc &x, V &visitor) {
    if (visitor.Pre(x)) {
      Walk(x.kind, visitor);
      visitor.Post(x);
    }
  }
  template <typename M>
  static void Walk(format::ControlEditDesc &x, M &mutator) {
    if (mutator.Pre(x)) {
      Walk(x.kind, mutator);
      mutator.Post(x);
    }
  }
  template <typename V>
  static void Walk(const format::DerivedTypeDataEditDesc &x, V &visitor) {
    if (visitor.Pre(x)) {
      Walk(x.type, visitor);
      Walk(x.parameters, visitor);
      visitor.Post(x);
    }
  }
  template <typename M>
  static void Walk(format::DerivedTypeDataEditDesc &x, M &mutator) {
    if (mutator.Pre(x)) {
      Walk(x.type, mutator);
      Walk(x.parameters, mutator);
      mutator.Post(x);
    }
  }
  template <typename V>
  static void Walk(const format::FormatItem &x, V &visitor) {
    if (visitor.Pre(x)) {
      Walk(x.repeatCount, visitor);
      Walk(x.u, visitor);
      visitor.Post(x);
    }
  }
  template <typename M> static void Walk(format::FormatItem &x, M &mutator) {
    if (mutator.Pre(x)) {
      Walk(x.repeatCount, mutator);
      Walk(x.u, mutator);
      mutator.Post(x);
    }
  }
  template <typename V>
  static void Walk(const format::FormatSpecification &x, V &visitor) {
    if (visitor.Pre(x)) {
      Walk(x.items, visitor);
      Walk(x.unlimitedItems, visitor);
      visitor.Post(x);
    }
  }
  template <typename M>
  static void Walk(format::FormatSpecification &x, M &mutator) {
    if (mutator.Pre(x)) {
      Walk(x.items, mutator);
      Walk(x.unlimitedItems, mutator);
      mutator.Post(x);
    }
  }
  template <typename V>
  static void Walk(const format::IntrinsicTypeDataEditDesc &x, V &visitor) {
    if (visitor.Pre(x)) {
      Walk(x.kind, visitor);
      Walk(x.width, visitor);
      Walk(x.digits, visitor);
      Walk(x.exponentWidth, visitor);
      visitor.Post(x);
    }
  }
  template <typename M>
  static void Walk(format::IntrinsicTypeDataEditDesc &x, M &mutator) {
    if (mutator.Pre(x)) {
      Walk(x.kind, mutator);
      Walk(x.width, mutator);
      Walk(x.digits, mutator);
      Walk(x.exponentWidth, mutator);
      mutator.Post(x);
    }
  }
};
} // namespace detail
 
template <typename A, typename V> void Walk(const A &x, V &visitor) {
  detail::ParseTreeVisitorLookupScope::Walk(x, visitor);
}
 
template <typename A, typename M> void Walk(A &x, M &mutator) {
  detail::ParseTreeVisitorLookupScope::Walk(x, mutator);
}
 
} // namespace Fortran::parser
#endif // FORTRAN_PARSER_PARSE_TREE_VISITOR_H_