Add a data structure for a tree of arithmetic operations

The basic structure just represents a (leaf-type-agnostic) tree of operations
and provides a mechanism to create a visitor using std::visit. This is needed
for lower_eval to separate the construction of the lowered arithmetic tree
from the materialization of the IR. However, I think it will (with adaptations)
be useful for other situations as well:

- Symbolic noise analysis in #1817
- To simplify the core routine in operation-balancer (https://github.com/google/heir/blob/b0cf72da113e6c7282733f8ba6bfcb7754a7495c/lib/Transforms/OperationBalancer/OperationBalancer.cpp#L74)

PiperOrigin-RevId: 770267746

Author: j2kun

lib/Utils/ArithmeticTree.h

@@ -0,0 +1,131 @@
+#ifndef LIB_UTILS_ARITHMETICTREE_H_
+#define LIB_UTILS_ARITHMETICTREE_H_
+
+#include <cstddef>
+#include <memory>
+#include <variant>
+
+namespace mlir {
+namespace heir {
+
+// This file contains a generic tree structure that can be used for representing
+// arithmetic trees with leaf nodes of various types.
+template <typename T>
+struct ArithmeticTreeNode;
+
+// A leaf node for the tree
+template <typename T>
+struct LeafNode {
+  T value;
+};
+
+using ConstantNode = LeafNode<double>;
+
+template <typename T>
+struct AddNode {
+  std::unique_ptr<ArithmeticTreeNode<T>> left;
+  std::unique_ptr<ArithmeticTreeNode<T>> right;
+};
+
+template <typename T>
+struct MultiplyNode {
+  std::unique_ptr<ArithmeticTreeNode<T>> left;
+  std::unique_ptr<ArithmeticTreeNode<T>> right;
+};
+
+template <typename T>
+struct PowerNode {
+  std::unique_ptr<ArithmeticTreeNode<T>> base;
+  size_t exponent;
+};
+
+template <typename T>
+struct ArithmeticTreeNode {
+ public:
+  std::variant<ConstantNode, LeafNode<T>, AddNode<T>, MultiplyNode<T>,
+               PowerNode<T>>
+      node_variant;
+
+  explicit ArithmeticTreeNode(double constant)
+      : node_variant(ConstantNode{constant}) {}
+  explicit ArithmeticTreeNode(const T& value)
+      : node_variant(LeafNode<T>{value}) {}
+  explicit ArithmeticTreeNode(T&& value)
+      : node_variant(LeafNode<T>{std::move(value)}) {}
+
+ private:
+  ArithmeticTreeNode() = default;
+
+ public:
+  // Static factory methods
+  static std::unique_ptr<ArithmeticTreeNode<T>> constant(double constant) {
+    return std::unique_ptr<ArithmeticTreeNode<T>>(
+        new ArithmeticTreeNode<T>(constant));
+  }
+
+  static std::unique_ptr<ArithmeticTreeNode<T>> leaf(const T& value) {
+    return std::unique_ptr<ArithmeticTreeNode<T>>(
+        new ArithmeticTreeNode<T>(value));
+  }
+
+  static std::unique_ptr<ArithmeticTreeNode<T>> add(
+      std::unique_ptr<ArithmeticTreeNode<T>> lhs,
+      std::unique_ptr<ArithmeticTreeNode<T>> rhs) {
+    assert(lhs && rhs && "invalid add");
+    auto node =
+        std::unique_ptr<ArithmeticTreeNode<T>>(new ArithmeticTreeNode<T>());
+    // Note, to satisfy variant we need to use aggregate initialization inside
+    // emplace
+    node->node_variant.template emplace<AddNode<T>>(
+        AddNode<T>{std::move(lhs), std::move(rhs)});
+    return node;
+  }
+
+  static std::unique_ptr<ArithmeticTreeNode<T>> mul(
+      std::unique_ptr<ArithmeticTreeNode<T>> lhs,
+      std::unique_ptr<ArithmeticTreeNode<T>> rhs) {
+    assert(lhs && rhs && "invalid mul");
+    auto node =
+        std::unique_ptr<ArithmeticTreeNode<T>>(new ArithmeticTreeNode<T>());
+    node->node_variant.template emplace<MultiplyNode<T>>(
+        MultiplyNode<T>{std::move(lhs), std::move(rhs)});
+    return node;
+  }
+
+  static std::unique_ptr<ArithmeticTreeNode<T>> power(
+      std::unique_ptr<ArithmeticTreeNode<T>> base, size_t exponent) {
+    assert(base && "invalid base for power");
+    auto node =
+        std::unique_ptr<ArithmeticTreeNode<T>>(new ArithmeticTreeNode<T>());
+    node->node_variant.template emplace<PowerNode<T>>(
+        PowerNode<T>{std::move(base), exponent});
+    return node;
+  }
+  // The presence of std::unique_ptr in AddNode, MultiplyNode, and PowerNode
+  // makes these types non-copyable. Consequently, ArithmeticTreeNode
+  // itself is move-only by default.
+  ArithmeticTreeNode(const ArithmeticTreeNode&) = delete;  // No copying
+  ArithmeticTreeNode& operator=(const ArithmeticTreeNode&) =
+      delete;  // No copy assignment
+
+  ArithmeticTreeNode(ArithmeticTreeNode&& other) noexcept =
+      default;  // Allow move construction
+  ArithmeticTreeNode& operator=(ArithmeticTreeNode&& other) noexcept =
+      default;  // Allow move assignment
+
+  // Visitor pattern
+  template <typename VisitorFunc>
+  decltype(auto) visit(VisitorFunc&& visitor) {
+    return std::visit(std::forward<VisitorFunc>(visitor), node_variant);
+  }
+
+  template <typename VisitorFunc>
+  decltype(auto) visit(VisitorFunc&& visitor) const {
+    return std::visit(std::forward<VisitorFunc>(visitor), node_variant);
+  }
+};
+
+}  // namespace heir
+}  // namespace mlir
+
+#endif  // LIB_UTILS_ARITHMETICTREE_H_

lib/Utils/ArithmeticTreeTest.cpp

@@ -0,0 +1,59 @@
+#include <iomanip>
+#include <ios>
+#include <sstream>
+#include <string>
+
+#include "gtest/gtest.h"  // from @googletest
+#include "lib/Utils/ArithmeticTree.h"
+
+namespace mlir {
+namespace heir {
+namespace {
+
+using StringLeavedTree = ArithmeticTreeNode<std::string>;
+
+struct FlattenedStringVisitor {
+  std::string operator()(const ConstantNode& node) const {
+    std::stringstream ss;
+    ss << std::fixed << std::setprecision(2) << node.value;
+    return ss.str();
+  }
+
+  std::string operator()(const LeafNode<std::string>& node) const {
+    return node.value;
+  }
+
+  std::string operator()(const AddNode<std::string>& node) const {
+    std::stringstream ss;
+    ss << "(" << node.left->visit(*this) << " + " << node.right->visit(*this)
+       << ")";
+    return ss.str();
+  }
+
+  std::string operator()(const MultiplyNode<std::string>& node) const {
+    std::stringstream ss;
+    ss << node.left->visit(*this) << " * " << node.right->visit(*this);
+    return ss.str();
+  }
+
+  std::string operator()(const PowerNode<std::string>& node) const {
+    std::stringstream ss;
+    ss << "(" << node.base->visit(*this) << " ^ " << node.exponent << ")";
+    return ss.str();
+  }
+};
+
+TEST(ArithmeticTreeTest, TestPrint) {
+  auto root = StringLeavedTree::mul(
+      StringLeavedTree::add(StringLeavedTree::leaf("x"),
+                            StringLeavedTree::constant(3.0)),
+      StringLeavedTree::power(StringLeavedTree::leaf("y"), 2));
+
+  FlattenedStringVisitor visitor;
+  std::string result = root->visit(visitor);
+  EXPECT_EQ(result, "(x + 3.00) * (y ^ 2)");
+}
+
+}  // namespace
+}  // namespace heir
+}  // namespace mlir

lib/Utils/BUILD

@@ -153,3 +153,18 @@ cc_library(
         "@llvm-project//mlir:Support",
     ],
 )
+
+cc_library(
+    name = "ArithmeticTree",
+    srcs = ["ArithmeticTree.h"],
+    hdrs = ["ArithmeticTree.h"],
+)
+
+cc_test(
+    name = "ArithmeticTreeTest",
+    srcs = ["ArithmeticTreeTest.cpp"],
+    deps = [
+        ":ArithmeticTree",
+        "@googletest//:gtest_main",
+    ],
+)