From 5ac7ae287de30f545b682169307b4c7f5d928c73 Mon Sep 17 00:00:00 2001
From: Antoine Pietri <seirl@google.com>
Date: Fri, 3 Oct 2025 02:50:42 -0700
Subject: [PATCH] Fix source range offsets for lists, maps and messages.

Update the parser to derive the source range for list, map, and struct/message creation expressions from the full parser rule context rather than just the opening token. This ensures the recorded offsets in `EnrichedSourceInfo` span the entire composite expression.

PiperOrigin-RevId: 814612738
---
 parser/parser.cc      |  6 +++---
 parser/parser_test.cc | 23 +++++++++++++++++++++++
 2 files changed, 26 insertions(+), 3 deletions(-)

diff --git a/parser/parser.cc b/parser/parser.cc
index d430e3169..22efedda0 100644
--- a/parser/parser.cc
+++ b/parser/parser.cc
@@ -1089,7 +1089,7 @@ std::any ParserVisitor::visitCreateMessage(
   } else {
     name = absl::StrJoin(parts, ".");
   }
-  int64_t obj_id = factory_.NextId(SourceRangeFromToken(ctx->op));
+  int64_t obj_id = factory_.NextId(SourceRangeFromParserRuleContext(ctx));
   std::vector<StructExprField> fields;
   if (ctx->entries) {
     fields = visitFields(ctx->entries);
@@ -1191,7 +1191,7 @@ std::any ParserVisitor::visitNested(CelParser::NestedContext* ctx) {
 }
 
 std::any ParserVisitor::visitCreateList(CelParser::CreateListContext* ctx) {
-  int64_t list_id = factory_.NextId(SourceRangeFromToken(ctx->op));
+  int64_t list_id = factory_.NextId(SourceRangeFromParserRuleContext(ctx));
   auto elems = visitList(ctx->elems);
   return ExprToAny(factory_.NewList(list_id, std::move(elems)));
 }
@@ -1229,7 +1229,7 @@ std::vector<Expr> ParserVisitor::visitList(CelParser::ExprListContext* ctx) {
 }
 
 std::any ParserVisitor::visitCreateMap(CelParser::CreateMapContext* ctx) {
-  int64_t struct_id = factory_.NextId(SourceRangeFromToken(ctx->op));
+  int64_t struct_id = factory_.NextId(SourceRangeFromParserRuleContext(ctx));
   std::vector<MapExprEntry> entries;
   if (ctx->entries) {
     entries = visitEntries(ctx->entries);
diff --git a/parser/parser_test.cc b/parser/parser_test.cc
index 3865476ee..66b7a4a5a 100644
--- a/parser/parser_test.cc
+++ b/parser/parser_test.cc
@@ -1519,6 +1519,29 @@ TEST_P(ExpressionTest, Parse) {
   }
 }
 
+TEST(ExpressionTest, CompositeExpressionOffsets) {
+  ParserOptions options;
+  std::vector<Macro> macros = Macro::AllMacros();
+
+  std::string list_expr = "[1, 2]";
+  auto list_result = EnrichedParse(list_expr, macros, "<input>", options);
+  ASSERT_THAT(list_result, IsOk());
+  auto list_offsets = list_result->enriched_source_info().offsets();
+  EXPECT_EQ(list_offsets.at(1), std::make_pair(0, 5));
+
+  std::string map_expr = "{'a': 1}";
+  auto map_result = EnrichedParse(map_expr, macros, "<input>", options);
+  ASSERT_THAT(map_result, IsOk());
+  auto map_offsets = map_result->enriched_source_info().offsets();
+  EXPECT_EQ(map_offsets.at(1), std::make_pair(0, 7));
+
+  std::string msg_expr = "Msg{f: 1}";
+  auto msg_result = EnrichedParse(msg_expr, macros, "<input>", options);
+  ASSERT_THAT(msg_result, IsOk());
+  auto msg_offsets = msg_result->enriched_source_info().offsets();
+  EXPECT_EQ(msg_offsets.at(1), std::make_pair(0, 8));
+}
+
 TEST(ExpressionTest, TsanOom) {
   Parse(
       "[[a([[???[a[[??[a([[[[[[[[[[[[[[[[[[[[[[[[[[[[[[[[[[[[[[[[[[[[[[[[[[[[[["