From 167e797bdd017b074c1f4fe4a80598d2fa1e1613 Mon Sep 17 00:00:00 2001
From: Antoine Pietri <seirl@google.com>
Date: Thu, 11 Jun 2026 00:28:02 -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: 930338742
---
 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 6c6434319..a858337a4 100644
--- a/parser/parser.cc
+++ b/parser/parser.cc
@@ -1104,7 +1104,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);
@@ -1206,7 +1206,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)));
 }
@@ -1244,7 +1244,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 33c52b1d2..1add80f84 100644
--- a/parser/parser_test.cc
+++ b/parser/parser_test.cc
@@ -1520,6 +1520,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([[[[[[[[[[[[[[[[[[[[[[[[[[[[[[[[[[[[[[[[[[[[[[[[[[[[[["