diff --git a/common/expr.h b/common/expr.h index f6a32d4ee..18828471f 100644 --- a/common/expr.h +++ b/common/expr.h @@ -581,6 +581,27 @@ class ComprehensionExpr final { return release(iter_var_); } + ABSL_MUST_USE_RESULT const std::string& iter_var2() const + ABSL_ATTRIBUTE_LIFETIME_BOUND { + return iter_var2_; + } + + void set_iter_var2(std::string iter_var2) { + iter_var2_ = std::move(iter_var2); + } + + void set_iter_var2(absl::string_view iter_var2) { + iter_var2_.assign(iter_var2.data(), iter_var2.size()); + } + + void set_iter_var2(const char* iter_var2) { + set_iter_var2(absl::NullSafeStringView(iter_var2)); + } + + ABSL_MUST_USE_RESULT std::string release_iter_var2() { + return release(iter_var2_); + } + ABSL_MUST_USE_RESULT bool has_iter_range() const { return iter_range_ != nullptr; } @@ -685,6 +706,7 @@ class ComprehensionExpr final { friend void swap(ComprehensionExpr& lhs, ComprehensionExpr& rhs) noexcept { using std::swap; swap(lhs.iter_var_, rhs.iter_var_); + swap(lhs.iter_var2_, rhs.iter_var2_); swap(lhs.iter_range_, rhs.iter_range_); swap(lhs.accu_var_, rhs.accu_var_); swap(lhs.accu_init_, rhs.accu_init_); @@ -711,6 +733,7 @@ class ComprehensionExpr final { } std::string iter_var_; + std::string iter_var2_; std::unique_ptr iter_range_; std::string accu_var_; std::unique_ptr accu_init_; diff --git a/extensions/protobuf/internal/ast.cc b/extensions/protobuf/internal/ast.cc index e6972317c..4d2d0ed13 100644 --- a/extensions/protobuf/internal/ast.cc +++ b/extensions/protobuf/internal/ast.cc @@ -227,6 +227,7 @@ class ExprToProtoState final { auto* comprehension_proto = proto->mutable_comprehension_expr(); proto->set_id(expr.id()); comprehension_proto->set_iter_var(comprehension_expr.iter_var()); + comprehension_proto->set_iter_var2(comprehension_expr.iter_var2()); if (comprehension_expr.has_iter_range()) { Push(comprehension_expr.iter_range(), comprehension_proto->mutable_iter_range()); @@ -457,6 +458,7 @@ class ExprFromProtoState final { expr.set_id(proto.id()); auto& comprehension_expr = expr.mutable_comprehension_expr(); comprehension_expr.set_iter_var(comprehension_proto.iter_var()); + comprehension_expr.set_iter_var2(comprehension_proto.iter_var2()); comprehension_expr.set_accu_var(comprehension_proto.accu_var()); if (comprehension_proto.has_iter_range()) { Push(comprehension_proto.iter_range(), diff --git a/extensions/protobuf/internal/ast_test.cc b/extensions/protobuf/internal/ast_test.cc index 243d75920..cf8d36748 100644 --- a/extensions/protobuf/internal/ast_test.cc +++ b/extensions/protobuf/internal/ast_test.cc @@ -220,6 +220,34 @@ INSTANTIATE_TEST_SUITE_P( } } )pb"}, + {R"pb( + id: 1 + comprehension_expr { + iter_var: "foo" + iter_var2: "baz" + iter_range { + id: 2 + list_expr {} + } + accu_var: "bar" + accu_init { + id: 3 + list_expr {} + } + loop_condition { + id: 4 + const_expr { bool_value: true } + } + loop_step { + id: 4 + ident_expr { name: "bar" } + } + result { + id: 5 + ident_expr { name: "foo" } + } + } + )pb"}, })); TEST(ExprFromProto, StructFieldInMap) {