Add interruptable ast traversal utility for proto ASTs (CheckedExpr).

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,

eval/public/ast_traverse.cc

@@ -14,6 +14,7 @@
 
 #include "eval/public/ast_traverse.h"
 
+#include <memory>
 #include <stack>
 
 #include "cel/expr/syntax.pb.h"
@@ -344,6 +345,47 @@ void PushDependencies(const StackRecord& record, std::stack<StackRecord>& stack,
 
 }  // namespace
 
+namespace internal {
+struct AstTraversalState {
+  std::stack<StackRecord> stack;
+};
+}  // namespace internal
+
+AstTraversal AstTraversal::Create(const Expr* expr,
+                                  const SourceInfo* source_info,
+                                  TraversalOptions options) {
+  AstTraversal instance(options);
+  instance.state_ = std::make_unique<internal::AstTraversalState>();
+  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<StackRecord> stack;

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 <memory>
+
 #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<internal::AstTraversalState> state_;
+};
+
 // Traverses the AST representation in an expr proto.
 //
 // expr: root node of the tree.

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