[Pratt Parser] Add support for enriched source info

parser/BUILD

@@ -164,6 +164,9 @@ cc_library(
     hdrs = [
         "source_factory.h",
     ],
+    deps = [
+        "@com_google_absl//absl/container:flat_hash_map",
+    ],
 )
 
 cc_library(

parser/internal/BUILD

@@ -117,6 +117,7 @@ cc_library(
         "//parser:macro_registry",
         "//parser:options",
         "//parser:parser_interface",
+        "//parser:source_factory",
         "@com_google_absl//absl/algorithm:container",
         "@com_google_absl//absl/base:nullability",
         "@com_google_absl//absl/cleanup",

parser/internal/pratt_parser.cc

@@ -44,6 +44,7 @@
 #include "parser/macro_registry.h"
 #include "parser/options.h"
 #include "parser/parser_interface.h"
+#include "parser/source_factory.h"
 
 namespace cel::parser_internal {
 
@@ -173,15 +174,18 @@ absl::StatusOr<std::unique_ptr<cel::Source>> PrattParserImpl::PrepareSourceImpl(
 
 absl::StatusOr<std::unique_ptr<cel::Ast>> PrattParseImpl(
     const cel::Source& source, const cel::MacroRegistry& registry,
-    const ParserOptions& options, std::vector<cel::ParseIssue>* parse_issues) {
+    const ParserOptions& options, std::vector<cel::ParseIssue>* parse_issues,
+    cel::EnrichedSourceInfo* enriched_source_info) {
   if (source.content().size() > options.expression_size_codepoint_limit) {
     return absl::InvalidArgumentError(absl::StrFormat(
         "expression size exceeds codepoint limit. input size: %zu, limit: %d",
         source.content().size(), options.expression_size_codepoint_limit));
   }
   std::vector<cel::ParseIssue> issues;
   AstFactory factory(&registry);
-  PrattParserWorker<cel::Expr> worker(source, options, &issues, factory);
+  PrattParserWorker<cel::Expr> worker(
+      source, options, &issues, factory,
+      /*track_node_ranges=*/enriched_source_info != nullptr);
   Expr expr = worker.Parse();
   if (worker.is_recursion_limit_exceeded()) {
     return absl::CancelledError(
@@ -203,6 +207,10 @@ absl::StatusOr<std::unique_ptr<cel::Ast>> PrattParseImpl(
     return absl::InvalidArgumentError(err_msg);
   }
 
+  if (enriched_source_info != nullptr) {
+    *enriched_source_info = cel::EnrichedSourceInfo(worker.GetNodeRanges());
+  }
+
   cel::SourceInfo source_info;
   source_info.set_location(std::string(source.description()));
   for (const auto& [id, pos] : worker.GetNodePositions()) {

parser/internal/pratt_parser.h

@@ -25,13 +25,18 @@
 #include "absl/container/flat_hash_set.h"
 #include "absl/status/status.h"
 #include "absl/status/statusor.h"
+#include "absl/strings/string_view.h"
 #include "common/ast.h"
 #include "common/source.h"
 #include "parser/macro.h"
 #include "parser/macro_registry.h"
 #include "parser/options.h"
 #include "parser/parser_interface.h"
 
+namespace cel {
+class EnrichedSourceInfo;
+}  // namespace cel
+
 namespace cel::parser_internal {
 
 // PrattParserImpl implements the Pratt parsing algorithm for CEL expressions.
@@ -73,7 +78,8 @@ class PrattParserImpl final : public cel::Parser {
 absl::StatusOr<std::unique_ptr<cel::Ast>> PrattParseImpl(
     const cel::Source& source, const cel::MacroRegistry& registry,
     const ParserOptions& options,
-    std::vector<cel::ParseIssue>* parse_issues = nullptr);
+    std::vector<cel::ParseIssue>* parse_issues = nullptr,
+    cel::EnrichedSourceInfo* enriched_source_info = nullptr);
 
 class PrattParserBuilderImpl final : public cel::ParserBuilder {
  public:

parser/internal/pratt_parser_worker.cc

@@ -19,6 +19,7 @@
 #include <vector>
 
 #include "absl/base/nullability.h"
+#include "absl/base/optimization.h"
 #include "absl/strings/str_cat.h"
 #include "absl/strings/str_format.h"
 #include "absl/strings/string_view.h"
@@ -101,11 +102,13 @@ const BinaryOpInfo& GetBinaryOpInfo(TokenType type) {
 
 ParserWorker::ParserWorker(
     const cel::Source& source, const cel::ParserOptions& options,
-    std::vector<cel::ParseIssue>* absl_nullable parse_issues)
+    std::vector<cel::ParseIssue>* absl_nullable parse_issues,
+    bool track_node_ranges)
     : source_(source),
       options_(options),
       lexer_(source_),
-      parse_issues_(parse_issues) {}
+      parse_issues_(parse_issues),
+      track_node_ranges_(track_node_ranges) {}
 
 void ParserWorker::InitTokenStream() {
   current_token_ = Token{.type = TokenType::kError, .start = 0, .end = 0};
@@ -204,14 +207,15 @@ int64_t ParserWorker::NextId(int32_t position) {
   }
   if (position >= 0) {
     positions_.insert({id, position});
+    if (ABSL_PREDICT_FALSE(track_node_ranges_)) {
+      node_ranges_.insert({id, {position, position}});
+    }
   }
   return id;
 }
 
 int64_t ParserWorker::NextId() { return NextId(-1); }
 
-bool ParserWorker::NodeLimitExceeded() { return node_limit_exceeded_; }
-
 int64_t ParserWorker::CopyId(int64_t id) {
   if (id == 0) {
     return 0;
@@ -220,11 +224,20 @@ int64_t ParserWorker::CopyId(int64_t id) {
   if (auto it = positions_.find(id); it != positions_.end()) {
     pos = it->second;
   }
-  return NextId(pos);
+  int64_t new_id = NextId(pos);
+  if (ABSL_PREDICT_FALSE(track_node_ranges_)) {
+    if (auto it = node_ranges_.find(id); it != node_ranges_.end()) {
+      node_ranges_[new_id] = it->second;
+    }
+  }
+  return new_id;
 }
 
 void ParserWorker::EraseId(int64_t id) {
   positions_.erase(id);
+  if (ABSL_PREDICT_FALSE(track_node_ranges_)) {
+    node_ranges_.erase(id);
+  }
   if (next_id_ == id + 1) {
     --next_id_;
   }

parser/internal/pratt_parser_worker.h

@@ -51,11 +51,16 @@ namespace cel::parser_internal {
 class ParserWorker {
  public:
   ParserWorker(const cel::Source& source, const cel::ParserOptions& options,
-               std::vector<cel::ParseIssue>* absl_nullable parse_issues);
+               std::vector<cel::ParseIssue>* absl_nullable parse_issues,
+               bool track_node_ranges = false);
 
   const absl::flat_hash_map<int64_t, int32_t>& GetNodePositions() const {
     return positions_;
   }
+  const absl::flat_hash_map<int64_t, std::pair<int32_t, int32_t>>&
+  GetNodeRanges() const {
+    return node_ranges_;
+  }
   absl::Span<const int32_t> GetLineOffsets() const {
     return source_.line_offsets();
   }
@@ -75,11 +80,25 @@ class ParserWorker {
 
   // ID and Position tracking
   int64_t NextId(int32_t position);
-  int64_t NextId(const Token& token) { return NextId(token.start); }
+  int64_t NextId(const Token& token) {
+    int64_t id = NextId(token.start);
+    if (ABSL_PREDICT_FALSE(track_node_ranges_)) {
+      if (token.start >= 0 && token.end > token.start) {
+        node_ranges_[id] = {token.start, token.end - 1};
+      }
+    }
+    return id;
+  }
   int64_t NextId();
-  bool NodeLimitExceeded();
   int64_t CopyId(int64_t id);
   void EraseId(int64_t id);
+  void SetNodeRange(int64_t id, int32_t begin, int32_t end) {
+    if (ABSL_PREDICT_FALSE(track_node_ranges_)) {
+      if (id != 0 && begin >= 0 && end >= begin) {
+        node_ranges_[id] = {begin, end};
+      }
+    }
+  }
 
   // Error reporting and recovery
   bool is_recovery_limit_exceeded() const {
@@ -100,10 +119,12 @@ class ParserWorker {
   int64_t next_id_ = 1;
   bool node_limit_exceeded_ = false;
   absl::flat_hash_map<int64_t, int32_t> positions_;
+  absl::flat_hash_map<int64_t, std::pair<int32_t, int32_t>> node_ranges_;
   std::vector<cel::ParseIssue>* absl_nullable parse_issues_;
   int error_count_ = 0;
   bool lexer_error_reported_ = false;
   bool recursion_limit_exceeded_ = false;
+  bool track_node_ranges_ = false;
 };
 
 struct BinaryOpInfo {
@@ -131,8 +152,10 @@ class PrattParserWorker : public ParserWorker {
   explicit PrattParserWorker(
       const cel::Source& source, const cel::ParserOptions& options,
       std::vector<cel::ParseIssue>* absl_nullable parse_issues,
-      AstFactoryInterface<ExprNode>& ast_factory)
-      : ParserWorker(source, options, parse_issues), ast_factory_(ast_factory) {
+      AstFactoryInterface<ExprNode>& ast_factory,
+      bool track_node_ranges = false)
+      : ParserWorker(source, options, parse_issues, track_node_ranges),
+        ast_factory_(ast_factory) {
     this->InitTokenStream();
   }
 
@@ -696,7 +719,9 @@ ExprNode PrattParserWorker<ExprNode>::ParseList() {
       break;
     }
   }
-  Expect(TokenType::kRightBracket, "expected ']'");
+  if (Expect(TokenType::kRightBracket, "expected ']'")) {
+    SetNodeRange(list_id, open_tok.start, current_token_.end - 1);
+  }
   return builder.Build();
 }
 
@@ -732,7 +757,9 @@ ExprNode PrattParserWorker<ExprNode>::ParseMap() {
       break;
     }
   }
-  Expect(TokenType::kRightBrace, "expected '}'");
+  if (Expect(TokenType::kRightBrace, "expected '}'")) {
+    SetNodeRange(map_id, open_tok.start, current_token_.end - 1);
+  }
   return builder.Build();
 }
 
@@ -774,7 +801,14 @@ ExprNode PrattParserWorker<ExprNode>::ParseStruct(
       break;
     }
   }
-  Expect(TokenType::kRightBrace, "expected '}'");
+  if (Expect(TokenType::kRightBrace, "expected '}'")) {
+    int32_t start_pos = open_tok.start;
+    auto it = positions_.find(obj_id);
+    if (it != positions_.end()) {
+      start_pos = it->second;
+    }
+    SetNodeRange(obj_id, start_pos, current_token_.end - 1);
+  }
   return builder.Build();
 }
 
@@ -1052,7 +1086,7 @@ std::optional<ExprNode> PrattParserWorker<ExprNode>::TryExpandMacro(
   if (!expander) {
     return std::nullopt;
   }
-  if (NodeLimitExceeded()) {
+  if (node_limit_exceeded_) {
     ReportError(expr_id,
                 "could not expand macro: expression node limit exceeded");
     return std::nullopt;

parser/parser.cc

@@ -23,7 +23,6 @@
 #include <functional>
 #include <iterator>
 #include <limits>
-#include <map>
 #include <memory>
 #include <optional>
 #include <string>
@@ -1445,7 +1444,8 @@ cel::SourceInfo ParserVisitor::GetSourceInfo() {
 }
 
 EnrichedSourceInfo ParserVisitor::enriched_source_info() const {
-  std::map<int64_t, std::pair<int32_t, int32_t>> offsets;
+  absl::flat_hash_map<int64_t, std::pair<int32_t, int32_t>> offsets;
+  offsets.reserve(factory_.positions().size());
   for (const auto& positions : factory_.positions()) {
     offsets.insert(
         std::pair{positions.first,
@@ -1951,14 +1951,17 @@ absl::StatusOr<VerboseParsedExpr> EnrichedParse(
     const ParserOptions& options) {
   ParsedExpr parsed_expr;
   if (options.enable_pratt_parser) {
-    CEL_ASSIGN_OR_RETURN(
-        std::unique_ptr<cel::Ast> ast,
-        cel::parser_internal::PrattParseImpl(source, registry, options));
+    EnrichedSourceInfo enriched_source_info;
+    CEL_ASSIGN_OR_RETURN(std::unique_ptr<cel::Ast> ast,
+                         cel::parser_internal::PrattParseImpl(
+                             source, registry, options,
+                             /*parse_issues=*/nullptr, &enriched_source_info));
     CEL_RETURN_IF_ERROR(cel::ast_internal::ExprToProto(
         ast->root_expr(), parsed_expr.mutable_expr()));
     CEL_RETURN_IF_ERROR(cel::ast_internal::SourceInfoToProto(
         ast->source_info(), parsed_expr.mutable_source_info()));
-    return VerboseParsedExpr(std::move(parsed_expr), EnrichedSourceInfo());
+    return VerboseParsedExpr(std::move(parsed_expr),
+                             std::move(enriched_source_info));
   }
   CEL_ASSIGN_OR_RETURN(ParseResult parse_result,
                        ParseImpl(source, registry, options));

parser/parser_test.cc

@@ -1650,8 +1650,13 @@ class LocationAdorner : public cel::ExpressionAdorner {
 
 std::string ConvertEnrichedSourceInfoToString(
     const EnrichedSourceInfo& enriched_source_info) {
+  std::vector<std::pair<int64_t, std::pair<int32_t, int32_t>>> sorted_offsets(
+      enriched_source_info.offsets().begin(),
+      enriched_source_info.offsets().end());
+  absl::c_sort(sorted_offsets);
   std::vector<std::string> offsets;
-  for (const auto& offset : enriched_source_info.offsets()) {
+  offsets.reserve(sorted_offsets.size());
+  for (const auto& offset : sorted_offsets) {
     offsets.push_back(absl::StrFormat(
         "[%d,%d,%d]", offset.first, offset.second.first, offset.second.second));
   }
@@ -1747,24 +1752,23 @@ TEST_P(ExpressionTest, Parse) {
   }
 }
 
-TEST(ExpressionTest, CompositeExpressionOffsets) {
-  ParserOptions options;
+TEST_P(ExpressionTest, CompositeExpressionOffsets) {
   std::vector<Macro> macros = Macro::AllMacros();
 
   std::string list_expr = "[1, 2]";
-  auto list_result = EnrichedParse(list_expr, macros, "<input>", options);
+  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);
+  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);
+  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));

parser/source_factory.h

@@ -16,15 +16,16 @@
 #define THIRD_PARTY_CEL_CPP_PARSER_SOURCE_FACTORY_H_
 
 #include <cstdint>
-#include <map>
 #include <utility>
 
-namespace google::api::expr::parser {
+#include "absl/container/flat_hash_map.h"
+
+namespace cel {
 
 class EnrichedSourceInfo {
  public:
   explicit EnrichedSourceInfo(
-      std::map<int64_t, std::pair<int32_t, int32_t>> offsets)
+      absl::flat_hash_map<int64_t, std::pair<int32_t, int32_t>> offsets)
       : offsets_(std::move(offsets)) {}
 
   EnrichedSourceInfo() = default;
@@ -33,15 +34,22 @@ class EnrichedSourceInfo {
   EnrichedSourceInfo(EnrichedSourceInfo&& other) = default;
   EnrichedSourceInfo& operator=(EnrichedSourceInfo&& other) = default;
 
-  const std::map<int64_t, std::pair<int32_t, int32_t>>& offsets() const {
+  const absl::flat_hash_map<int64_t, std::pair<int32_t, int32_t>>& offsets()
+      const {
     return offsets_;
   }
 
  private:
   // A map between node_id and pair of start position and end position
-  std::map<int64_t, std::pair<int32_t, int32_t>> offsets_;
+  absl::flat_hash_map<int64_t, std::pair<int32_t, int32_t>> offsets_;
 };
 
+}  // namespace cel
+
+namespace google::api::expr::parser {
+
+using EnrichedSourceInfo = ::cel::EnrichedSourceInfo;
+
 }  // namespace google::api::expr::parser
 
 #endif  // THIRD_PARTY_CEL_CPP_PARSER_SOURCE_FACTORY_H_