[llvm-branch-commits] [flang] [flang][Evaluate] Implement rewriting framework for evaluate::Expr (PR #153037)

[Krzysztof Parzyszek via llvm-branch-commits]

https://github.com/kparzysz created https://github.com/llvm/llvm-project/pull/153037

The structure of evaluate::Expr is highly customized for the specific operation or entity that it represents. The different cases are expressed with different types, which makes the traversal and modifications somewhat complicated. There exists a framework for read-only traversal (traverse.h), but there is nothing that helps with modifying evaluate::Expr.

It's rare that evaluate::Expr needs to be modified, but for the cases where it needs to be, this code will make it easier.

>From 2120125ad74caa41402f782a93a5a153520fdc1c Mon Sep 17 00:00:00 2001
From: Krzysztof Parzyszek <Krzysztof.Parzyszek at amd.com>
Date: Tue, 29 Jul 2025 16:17:29 -0500
Subject: [PATCH] [flang][Evaluate] Implement rewriting framework for
 evaluate::Expr

The structure of evaluate::Expr is highly customized for the specific
operation or entity that it represents. The different cases are expressed
with different types, which makes the traversal and modifications somewhat
complicated. There exists a framework for read-only traversal (traverse.h),
but there is nothing that helps with modifying evaluate::Expr.

It's rare that evaluate::Expr needs to be modified, but for the cases
where it needs to be, this code will make it easier.
---
 flang/include/flang/Evaluate/rewrite.h | 160 +++++++++++++++++++++++++
 1 file changed, 160 insertions(+)
 create mode 100644 flang/include/flang/Evaluate/rewrite.h

diff --git a/flang/include/flang/Evaluate/rewrite.h b/flang/include/flang/Evaluate/rewrite.h
new file mode 100644
index 0000000000000..034b1efa21977
--- /dev/null
+++ b/flang/include/flang/Evaluate/rewrite.h
@@ -0,0 +1,160 @@
+//===-- include/flang/Evaluate/rewrite.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_EVALUATE_REWRITE_H_
+#define FORTRAN_EVALUATE_REWRITE_H_
+
+#include "flang/Common/visit.h"
+#include "flang/Evaluate/expression.h"
+#include "flang/Support/Fortran.h"
+#include "llvm/ADT/STLExtras.h"
+
+#include <tuple>
+#include <type_traits>
+#include <utility>
+#include <variant>
+
+namespace Fortran::evaluate {
+namespace rewrite {
+namespace detail {
+template <typename, typename = void> //
+struct IsOperation {
+  static constexpr bool value{false};
+};
+
+template <typename T>
+struct IsOperation<T, std::void_t<decltype(T::operands)>> {
+  static constexpr bool value{true};
+};
+} // namespace detail
+
+template <typename T>
+constexpr bool is_operation_v{detail::IsOperation<T>::value};
+
+/// Individual Expr<T> rewriter that simply constructs an expression that is
+/// identical to the input. This is a suitable base class for all user-defined
+/// rewriters.
+struct Identity {
+  template <typename T, typename U>
+  Expr<T> operator()(Expr<T> &&x, const U &op) {
+    return std::move(x);
+  }
+};
+
+/// Bottom-up Expr<T> rewriter.
+///
+/// The Mutator traverses and reconstructs given Expr<T>. Going bottom-up,
+/// whenever the traversal visits a sub-node of type Expr<U> (for some U),
+/// it will invoke the user-provided rewriter via the () operator.
+///
+/// If x is of type Expr<U>, it will call (in pseudo-code):
+///   rewriter_(x, active_member_of(x.u))
+/// The second parameter is there to make it easier to overload the () operator
+/// for specific operations in Expr<...>.
+///
+/// The user rewriter is only invoked for Expr<U>, not for Operation, nor any
+/// other subobject.
+template <typename Rewriter> struct Mutator {
+  Mutator(Rewriter &rewriter) : rewriter_(rewriter) {}
+
+  template <typename T, typename U = llvm::remove_cvref_t<T>>
+  U operator()(T &&x) {
+    if constexpr (std::is_lvalue_reference_v<T>) {
+      return Mutate(U(x));
+    } else {
+      return Mutate(std::move(x));
+    }
+  }
+
+private:
+  template <typename T> struct LambdaWithRvalueCapture {
+    LambdaWithRvalueCapture(Rewriter &r, Expr<T> &&c)
+        : rewriter_(r), capture_(std::move(c)) {}
+    template <typename S> Expr<T> operator()(const S &s) {
+      return rewriter_(std::move(capture_), s);
+    }
+
+  private:
+    Rewriter &rewriter_;
+    Expr<T> &&capture_;
+  };
+
+  template <typename T, typename = std::enable_if_t<!is_operation_v<T>>>
+  T Mutate(T &&x) const {
+    return std::move(x);
+  }
+
+  template <typename D, typename = std::enable_if_t<is_operation_v<D>>>
+  D Mutate(D &&op, std::make_index_sequence<D::operands> t = {}) const {
+    return MutateOp(std::move(op), t);
+  }
+
+  template <typename T> //
+  Expr<T> Mutate(Expr<T> &&x) const {
+    // First construct the new expression with the rewritten op.
+    Expr<T> n{common::visit(
+        [&](auto &&s) { //
+          return Expr<T>(Mutate(std::move(s)));
+        },
+        std::move(x.u))};
+    // Return the rewritten expression. The second visit it to make sure
+    // that the second argument in the call to the rewriter is a part of
+    // the Expr<T> passed to it.
+    return common::visit(
+        LambdaWithRvalueCapture<T>(rewriter_, std::move(n)), std::move(n.u));
+  }
+
+  template <typename... Ts>
+  std::variant<Ts...> Mutate(std::variant<Ts...> &&u) const {
+    return common::visit(
+        [this](auto &&s) { return Mutate(std::move(s)); }, std::move(u));
+  }
+
+  template <typename... Ts>
+  std::tuple<Ts...> Mutate(std::tuple<Ts...> &&t) const {
+    return MutateTuple(std::move(t), std::index_sequence_for<Ts...>{});
+  }
+
+  template <typename... Ts, size_t... Is>
+  std::tuple<Ts...> MutateTuple(
+      std::tuple<Ts...> &&t, std::index_sequence<Is...>) const {
+    return std::make_tuple(Mutate(std::move(std::get<Is>(t))...));
+  }
+
+  template <typename D, size_t... Is>
+  D MutateOp(D &&op, std::index_sequence<Is...>) const {
+    return D(Mutate(std::move(op.template operand<Is>()))...);
+  }
+
+  template <typename T, size_t... Is>
+  Extremum<T> MutateOp(Extremum<T> &&op, std::index_sequence<Is...>) const {
+    return Extremum<T>(
+        op.ordering, Mutate(std::move(op.template operand<Is>()))...);
+  }
+
+  template <int K, size_t... Is>
+  ComplexComponent<K> MutateOp(
+      ComplexComponent<K> &&op, std::index_sequence<Is...>) const {
+    return ComplexComponent<K>(
+        op.isImaginaryPart, Mutate(std::move(op.template operand<Is>()))...);
+  }
+
+  template <int K, size_t... Is>
+  LogicalOperation<K> MutateOp(
+      LogicalOperation<K> &&op, std::index_sequence<Is...>) const {
+    return LogicalOperation<K>(
+        op.logicalOperator, Mutate(std::move(op.template operand<Is>()))...);
+  }
+
+  Rewriter &rewriter_;
+};
+
+template <typename Rewriter> Mutator(Rewriter &) -> Mutator<Rewriter>;
+} // namespace rewrite
+} // namespace Fortran::evaluate
+
+#endif // FORTRAN_EVALUATE_REWRITE_H_