From 0d627160ef3e042f5c9ca259a67b9c22f3fcc7db Mon Sep 17 00:00:00 2001 From: Jonathan Tatum Date: Mon, 20 Jul 2026 13:46:44 -0700 Subject: [PATCH] Add interruptable ast traversal utility for proto ASTs (CheckedExpr). PiperOrigin-RevId: 951031828 --- common/ast_traverse.h | 2 -- eval/public/ast_traverse.cc | 42 +++++++++++++++++++++++++++++ eval/public/ast_traverse.h | 43 ++++++++++++++++++++++++++++++ eval/public/ast_traverse_test.cc | 45 ++++++++++++++++++++++++++++++++ 4 files changed, 130 insertions(+), 2 deletions(-) diff --git a/common/ast_traverse.h b/common/ast_traverse.h index 004727e49..21272a898 100644 --- a/common/ast_traverse.h +++ b/common/ast_traverse.h @@ -43,8 +43,6 @@ struct TraversalOptions { // while(!traversal.IsDone()) { // traversal.Step(visitor); // } -// -// This class is thread-hostile and should only be used in synchronous code. class AstTraversal { public: static AstTraversal Create(const cel::Expr& ast ABSL_ATTRIBUTE_LIFETIME_BOUND, diff --git a/eval/public/ast_traverse.cc b/eval/public/ast_traverse.cc index c18b806b9..a31e1d37d 100644 --- a/eval/public/ast_traverse.cc +++ b/eval/public/ast_traverse.cc @@ -14,6 +14,7 @@ #include "eval/public/ast_traverse.h" +#include #include #include "cel/expr/syntax.pb.h" @@ -344,6 +345,47 @@ void PushDependencies(const StackRecord& record, std::stack& stack, } // namespace +namespace internal { +struct AstTraversalState { + std::stack stack; +}; +} // namespace internal + +AstTraversal AstTraversal::Create(const Expr* expr, + const SourceInfo* source_info, + TraversalOptions options) { + AstTraversal instance(options); + instance.state_ = std::make_unique(); + instance.state_->stack.push(StackRecord(expr, source_info)); + return instance; +} + +AstTraversal::AstTraversal(TraversalOptions options) : options_(options) {} + +AstTraversal::~AstTraversal() = default; + +bool AstTraversal::Step(AstVisitor* visitor) { + if (IsDone()) { + return false; + } + auto& stack = state_->stack; + StackRecord& record = stack.top(); + if (!record.visited) { + PreVisit(record, visitor); + PushDependencies(record, stack, options_); + record.visited = true; + } else { + PostVisit(record, visitor); + stack.pop(); + } + + return !stack.empty(); +} + +bool AstTraversal::IsDone() { + return state_ == nullptr || state_->stack.empty(); +} + void AstTraverse(const Expr* expr, const SourceInfo* source_info, AstVisitor* visitor, TraversalOptions options) { std::stack stack; diff --git a/eval/public/ast_traverse.h b/eval/public/ast_traverse.h index f81c6f47a..a330bd28d 100644 --- a/eval/public/ast_traverse.h +++ b/eval/public/ast_traverse.h @@ -17,17 +17,60 @@ #ifndef THIRD_PARTY_CEL_CPP_EVAL_PUBLIC_AST_TRAVERSE_H_ #define THIRD_PARTY_CEL_CPP_EVAL_PUBLIC_AST_TRAVERSE_H_ +#include + #include "cel/expr/syntax.pb.h" #include "eval/public/ast_visitor.h" namespace google::api::expr::runtime { +namespace internal { +struct AstTraversalState; +} // namespace internal + struct TraversalOptions { bool use_comprehension_callbacks; TraversalOptions() : use_comprehension_callbacks(false) {} }; +// Helper class for managing the traversal of the AST. +// Allows caller to step through the traversal. +// +// Usage: +// +// AstTraversal traversal = AstTraversal::Create(expr, source_info); +// +// MyVisitor visitor(); +// while (!traversal.IsDone()) { +// traversal.Step(&visitor); +// } +class AstTraversal { + public: + static AstTraversal Create(const cel::expr::Expr* expr, + const cel::expr::SourceInfo* source_info, + TraversalOptions options = TraversalOptions()); + + ~AstTraversal(); + + AstTraversal(const AstTraversal&) = delete; + AstTraversal& operator=(const AstTraversal&) = delete; + AstTraversal(AstTraversal&&) = default; + AstTraversal& operator=(AstTraversal&&) = default; + + // Advances the traversal. Returns true if there is more work to do. This is a + // no-op if the traversal is done and IsDone() is true. + bool Step(AstVisitor* visitor); + + // Returns true if there is no work left to do. + bool IsDone(); + + private: + explicit AstTraversal(TraversalOptions options); + TraversalOptions options_; + std::unique_ptr state_; +}; + // Traverses the AST representation in an expr proto. // // expr: root node of the tree. diff --git a/eval/public/ast_traverse_test.cc b/eval/public/ast_traverse_test.cc index ca6d81b72..dfc4a7700 100644 --- a/eval/public/ast_traverse_test.cc +++ b/eval/public/ast_traverse_test.cc @@ -465,6 +465,51 @@ TEST(AstCrawlerTest, CheckExprHandlers) { AstTraverse(&expr, &source_info, &handler); } +TEST(AstTraversal, Interrupt) { + SourceInfo source_info; + MockAstVisitor handler; + + Expr expr; + auto* select_expr = expr.mutable_select_expr(); + auto* operand = select_expr->mutable_operand(); + auto* ident_expr = operand->mutable_ident_expr(); + + testing::InSequence seq; + + auto traversal = AstTraversal::Create(&expr, &source_info); + + EXPECT_CALL(handler, PreVisitExpr(_, _)).Times(2); + + EXPECT_CALL(handler, PostVisitIdent(ident_expr, operand, _)).Times(1); + EXPECT_CALL(handler, PostVisitSelect(select_expr, &expr, _)).Times(0); + + EXPECT_TRUE(traversal.Step(&handler)); + EXPECT_TRUE(traversal.Step(&handler)); + EXPECT_TRUE(traversal.Step(&handler)); + + EXPECT_FALSE(traversal.IsDone()); +} + +TEST(AstTraversal, NoInterrupt) { + SourceInfo source_info; + MockAstVisitor handler; + + Expr expr; + auto* select_expr = expr.mutable_select_expr(); + auto* operand = select_expr->mutable_operand(); + auto* ident_expr = operand->mutable_ident_expr(); + + testing::InSequence seq; + + auto traversal = AstTraversal::Create(&expr, &source_info); + + EXPECT_CALL(handler, PostVisitIdent(ident_expr, operand, _)).Times(1); + EXPECT_CALL(handler, PostVisitSelect(select_expr, &expr, _)).Times(1); + + while (traversal.Step(&handler)) continue; + EXPECT_TRUE(traversal.IsDone()); +} + } // namespace } // namespace google::api::expr::runtime