Update AST conversion to support iter_var2

PiperOrigin-RevId: 693014456

Author: jcking

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<Expr> iter_range_;
   std::string accu_var_;
   std::unique_ptr<Expr> accu_init_;

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(),

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) {