Add non-recursive FlattenedErase for Expr.

PiperOrigin-RevId: 850184035

Author: jnthntatum

common/BUILD

@@ -54,8 +54,8 @@ cc_library(
         "@com_google_absl//absl/base:core_headers",
         "@com_google_absl//absl/base:no_destructor",
         "@com_google_absl//absl/functional:overload",
+        "@com_google_absl//absl/log:absl_check",
         "@com_google_absl//absl/strings:string_view",
-        "@com_google_absl//absl/types:optional",
         "@com_google_absl//absl/types:span",
         "@com_google_absl//absl/types:variant",
     ],

common/expr.cc

@@ -19,6 +19,7 @@
 
 #include "absl/base/no_destructor.h"
 #include "absl/functional/overload.h"
+#include "absl/log/absl_check.h"
 #include "absl/types/variant.h"
 #include "common/constant.h"
 
@@ -201,4 +202,119 @@ Expr& Expr::operator=(const Expr& other) {
 
 Expr::Expr(const Expr& other) { CloneImpl(other, *this); }
 
+namespace common_internal {
+struct ExprEraseTag {};
+}  // namespace common_internal
+
+namespace {
+void Expand(Expr** tail, Expr* cur) {
+  static common_internal::ExprEraseTag tag;
+  switch (cur->kind_case()) {
+    case ExprKindCase::kSelectExpr: {
+      SelectExpr& select = cur->mutable_select_expr();
+      if (select.has_operand()) {
+        select.mutable_operand().SetNext(tag, *tail);
+        *tail = &select.mutable_operand();
+      }
+      break;
+    }
+    case ExprKindCase::kCallExpr: {
+      CallExpr& call = cur->mutable_call_expr();
+      if (call.has_target()) {
+        call.mutable_target().SetNext(tag, *tail);
+        *tail = &call.mutable_target();
+      }
+      for (auto& arg : call.mutable_args()) {
+        arg.SetNext(tag, *tail);
+        *tail = &arg;
+      }
+      break;
+    }
+    case ExprKindCase::kListExpr: {
+      for (auto& arg : cur->mutable_list_expr().mutable_elements()) {
+        arg.mutable_expr().SetNext(tag, *tail);
+        *tail = &arg.mutable_expr();
+      }
+      break;
+    }
+    case ExprKindCase::kStructExpr: {
+      for (auto& field : cur->mutable_struct_expr().mutable_fields()) {
+        field.mutable_value().SetNext(tag, *tail);
+        *tail = &field.mutable_value();
+      }
+      break;
+    }
+    case ExprKindCase::kMapExpr: {
+      for (auto& entry : cur->mutable_map_expr().mutable_entries()) {
+        entry.mutable_key().SetNext(tag, *tail);
+        *tail = &entry.mutable_key();
+        entry.mutable_value().SetNext(tag, *tail);
+        *tail = &entry.mutable_value();
+      }
+      break;
+    }
+    case ExprKindCase::kComprehensionExpr: {
+      if (cur->comprehension_expr().has_accu_init()) {
+        cur->mutable_comprehension_expr().mutable_accu_init().SetNext(tag,
+                                                                      *tail);
+        *tail = &cur->mutable_comprehension_expr().mutable_accu_init();
+      }
+      if (cur->comprehension_expr().has_iter_range()) {
+        cur->mutable_comprehension_expr().mutable_iter_range().SetNext(tag,
+                                                                       *tail);
+        *tail = &cur->mutable_comprehension_expr().mutable_iter_range();
+      }
+      if (cur->comprehension_expr().has_loop_condition()) {
+        cur->mutable_comprehension_expr().mutable_loop_condition().SetNext(
+            tag, *tail);
+        *tail = &cur->mutable_comprehension_expr().mutable_loop_condition();
+      }
+      if (cur->comprehension_expr().has_loop_step()) {
+        cur->mutable_comprehension_expr().mutable_loop_step().SetNext(tag,
+                                                                      *tail);
+        *tail = &cur->mutable_comprehension_expr().mutable_loop_step();
+      }
+      if (cur->comprehension_expr().has_result()) {
+        cur->mutable_comprehension_expr().mutable_result().SetNext(tag, *tail);
+        *tail = &cur->mutable_comprehension_expr().mutable_result();
+      }
+      break;
+    }
+    default:
+      // Leaf node, nothing to expand.
+      // Also a fallback in case we add a new node type.
+      // Note: already in the deleter list so we can't delete now, will be
+      // deleted after ordering the AST.
+      break;
+  }
+}
+}  // namespace
+
+void Expr::FlattenedErase() {
+  // High level idea is to build a topological ordering of the AST, then erase
+  // leaf to root.
+  this->u_.next = nullptr;
+  Expr* prev_tail = nullptr;
+  Expr* tail = this;
+
+  while (tail != prev_tail) {
+    Expr* next_prev_tail = tail;
+    Expr* expand_ptr = tail;
+    while (expand_ptr != prev_tail) {
+      ABSL_DCHECK(expand_ptr != nullptr);  // Linked list is broken or changed.
+      Expr* next_expand_ptr = expand_ptr->u_.next;
+      Expand(&tail, expand_ptr);
+      expand_ptr = next_expand_ptr;
+    }
+    prev_tail = next_prev_tail;
+  }
+
+  Expr* node = tail;
+  while (node != nullptr) {
+    Expr* next = node->u_.next;
+    node->Clear();
+    node = next;
+  }
+}
+
 }  // namespace cel

common/expr.h

@@ -884,6 +884,10 @@ enum class ExprKindCase {
   kComprehensionExpr,
 };
 
+namespace common_internal {
+struct ExprEraseTag;
+}  // namespace common_internal
+
 // `Expr` is a node in the Common Expression Language's abstract syntax tree. It
 // is composed of a numeric ID and a kind variant.
 class Expr final {
@@ -897,9 +901,9 @@ class Expr final {
 
   void Clear();
 
-  ABSL_MUST_USE_RESULT ExprId id() const { return id_; }
+  ABSL_MUST_USE_RESULT ExprId id() const { return u_.id; }
 
-  void set_id(ExprId id) { id_ = id; }
+  void set_id(ExprId id) { u_.id = id; }
 
   ABSL_MUST_USE_RESULT const ExprKind& kind() const
       ABSL_ATTRIBUTE_LIFETIME_BOUND {
@@ -1071,6 +1075,11 @@ class Expr final {
 
   friend void swap(Expr& lhs, Expr& rhs) noexcept;
 
+  // Erases the expr in place without recursion.
+  void FlattenedErase();
+
+  inline void SetNext(common_internal::ExprEraseTag&, Expr* next);
+
  private:
   friend class IdentExpr;
   friend class SelectExpr;
@@ -1105,7 +1114,12 @@ class Expr final {
   template <typename T>
   ABSL_MUST_USE_RESULT T release_kind();
 
-  ExprId id_ = 0;
+  union {
+    ExprId id = 0;
+    // Intrusive pointer to the next element in the destructor chain.
+    // Only set from FlattenedErase.
+    Expr* next;
+  } u_;
   ExprKind kind_;
 };
 
@@ -1354,7 +1368,7 @@ inline std::vector<MapExprEntry> MapExpr::release_entries() {
 }
 
 inline void Expr::Clear() {
-  id_ = 0;
+  u_.id = 0;
   mutable_kind().emplace<UnspecifiedExpr>();
 }
 
@@ -1416,7 +1430,7 @@ inline ExprKindCase Expr::kind_case() const {
 
 inline void swap(Expr& lhs, Expr& rhs) noexcept {
   using std::swap;
-  swap(lhs.id_, rhs.id_);
+  swap(lhs.u_, rhs.u_);
   swap(lhs.kind_, rhs.kind_);
 }
 
@@ -1695,6 +1709,10 @@ inline Expr MapExprEntry::release(std::unique_ptr<Expr>& property) {
   return Expr{};
 }
 
+inline void Expr::SetNext(common_internal::ExprEraseTag&, Expr* next) {
+  u_.next = next;
+}
+
 }  // namespace cel
 
 #endif  // THIRD_PARTY_CEL_CPP_COMMON_EXPR_H_

common/expr_test.cc

@@ -580,6 +580,89 @@ TEST(Expr, CopyAssignChildReference) {
   EXPECT_EQ(expr.call_expr().args()[1].ident_expr().name(), "qux");
 }
 
+TEST(Expr, FlattenedErase) {
+  Expr expr;
+  auto& list_expr = expr.mutable_list_expr();
+  list_expr.mutable_elements()
+      .emplace_back()
+      .mutable_expr()
+      .mutable_ident_expr()
+      .set_name("foo");
+
+  list_expr.mutable_elements()
+      .emplace_back()
+      .mutable_expr()
+      .mutable_select_expr()
+      .mutable_operand()
+      .mutable_ident_expr()
+      .set_name("foo");
+
+  auto& call_expr = list_expr.mutable_elements()
+                        .emplace_back()
+                        .mutable_expr()
+                        .mutable_call_expr();
+  call_expr.set_function("foo");
+  call_expr.mutable_target().mutable_ident_expr().set_name("bar");
+  call_expr.mutable_args().emplace_back().mutable_ident_expr().set_name("baz");
+
+  auto& struct_expr = list_expr.mutable_elements()
+                          .emplace_back()
+                          .mutable_expr()
+                          .mutable_struct_expr();
+  struct_expr.set_name("foo");
+  auto& field = struct_expr.mutable_fields().emplace_back();
+  field.set_name("bar");
+  field.mutable_value().mutable_ident_expr().set_name("baz");
+
+  auto& map_expr = list_expr.mutable_elements()
+                       .emplace_back()
+                       .mutable_expr()
+                       .mutable_map_expr();
+  auto& map_entry = map_expr.mutable_entries().emplace_back();
+  map_entry.mutable_key().mutable_const_expr().set_string_value("foo");
+  map_entry.mutable_value().mutable_ident_expr().set_name("bar");
+
+  auto& comprehension_expr = list_expr.mutable_elements()
+                                 .emplace_back()
+                                 .mutable_expr()
+                                 .mutable_comprehension_expr();
+  comprehension_expr.set_iter_var("foo");
+  comprehension_expr.set_accu_var("bar");
+  comprehension_expr.set_iter_range(Expr{});
+  comprehension_expr.set_accu_init(Expr{});
+  comprehension_expr.set_loop_condition(Expr{});
+  comprehension_expr.set_loop_step(Expr{});
+  comprehension_expr.set_result(Expr{});
+
+  expr.FlattenedErase();
+  EXPECT_EQ(expr.kind_case(), ExprKindCase::kUnspecifiedExpr);
+}
+
+Expr MakeNestedList(int size) {
+  Expr e;
+  Expr* node = &e;
+  e.set_id(1);
+  for (int i = 0; i < size; ++i) {
+    node = &node->mutable_list_expr()
+                .mutable_elements()
+                .emplace_back()
+                .mutable_expr();
+    node->set_id(i + 2);
+  }
+  return e;
+}
+
+TEST(Expr, FlattenedErase256k) {
+  // Large expr to ensure we're not recursing. Would likely hit stack limits
+  // with default destructor.
+  constexpr int size = 256 * 1024;
+
+  Expr expr = MakeNestedList(size);
+
+  expr.FlattenedErase();
+  EXPECT_EQ(expr.kind_case(), ExprKindCase::kUnspecifiedExpr);
+}
+
 TEST(Expr, Id) {
   Expr expr;
   EXPECT_THAT(expr.id(), Eq(0));