Add support for estimating original source positions.

common/source.cc

@@ -19,6 +19,7 @@
 #include <cstdint>
 #include <limits>
 #include <memory>
+#include <optional>
 #include <string>
 #include <tuple>
 #include <utility>

policy/BUILD

@@ -84,6 +84,7 @@ cc_library(
         ":cel_policy_parser",
         "//common:source",
         "//internal:status_macros",
+        "//policy/internal:yaml_string_element_scanner",
         "@com_google_absl//absl/status",
         "@com_google_absl//absl/status:statusor",
         "@com_google_absl//absl/strings",
@@ -190,6 +191,7 @@ cc_test(
         ":compiler",
         ":yaml_policy_parser",
         "//common:ast",
+        "//common:ast_proto",
         "//common:decl",
         "//common:navigable_ast",
         "//common:source",
@@ -211,6 +213,7 @@ cc_test(
         "//runtime:runtime_builder",
         "//runtime:runtime_options",
         "//runtime:standard_runtime_builder_factory",
+        "//tools:cel_unparser",
         "@com_google_absl//absl/status",
         "@com_google_absl//absl/status:status_matchers",
         "@com_google_absl//absl/status:statusor",

policy/cel_policy.cc

@@ -59,34 +59,64 @@ std::string IndentBlock(absl::string_view text) {
 
 void CelPolicySource::NoteSourcePosition(CelPolicyElementId id,
                                          SourcePosition position) {
-  source_positions_[id] = position;
+  source_info_[id].position = position;
+}
+
+void CelPolicySource::NoteSourceRange(CelPolicyElementId id,
+                                      std::optional<SourceRange> range,
+                                      bool quoted) {
+  ElementSourceInfo& info = source_info_[id];
+  info.range = range;
+  info.quoted = quoted;
+  if (range.has_value() && info.position == -1) {
+    info.position = range->begin;
+  }
 }
 
 std::optional<SourcePosition> CelPolicySource::GetSourcePosition(
     CelPolicyElementId id) const {
-  auto it = source_positions_.find(id);
-  if (it == source_positions_.end()) {
+  auto it = source_info_.find(id);
+  if (it == source_info_.end() || it->second.position == -1) {
     return std::nullopt;
   }
-  return it->second;
+  return it->second.position;
+}
+
+std::optional<SourceRange> CelPolicySource::GetSourceRange(
+    CelPolicyElementId id) const {
+  auto it = source_info_.find(id);
+  if (it == source_info_.end() || !it->second.range.has_value()) {
+    return std::nullopt;
+  }
+  return it->second.range;
+}
+
+std::optional<bool> CelPolicySource::IsQuoted(CelPolicyElementId id) const {
+  auto it = source_info_.find(id);
+  if (it == source_info_.end() || !it->second.range.has_value()) {
+    return std::nullopt;
+  }
+  return it->second.quoted;
 }
 
 std::optional<SourceLocation> CelPolicySource::GetSourceLocation(
     CelPolicyElementId id) const {
-  auto it = source_positions_.find(id);
-  if (it == source_positions_.end()) {
+  auto it = source_info_.find(id);
+  if (it == source_info_.end() || it->second.position == -1) {
     return std::nullopt;
   }
-  return policy_source_->GetLocation(it->second);
+  return policy_source_->GetLocation(it->second.position);
 }
 
 std::string CelPolicySource::DebugString() const {
   std::string result;
 
   // Sort the source elements in descending order of position
   std::vector<std::pair<CelPolicyElementId, SourcePosition>> sorted_positions;
-  for (const auto& pair : source_positions_) {
-    sorted_positions.push_back(pair);
+  for (const auto& [id, info] : source_info_) {
+    if (info.position != -1) {
+      sorted_positions.push_back({id, info.position});
+    }
   }
   std::sort(sorted_positions.begin(), sorted_positions.end(),
             [](const auto& a, const auto& b) {

policy/cel_policy.h

@@ -34,6 +34,12 @@ namespace cel {
 
 using CelPolicyElementId = int32_t;
 
+struct ElementSourceInfo {
+  SourcePosition position = -1;
+  std::optional<SourceRange> range;
+  bool quoted = false;
+};
+
 class CelPolicySource {
  public:
   explicit CelPolicySource(cel::SourcePtr policy_source)
@@ -43,15 +49,22 @@ class CelPolicySource {
 
   void NoteSourcePosition(CelPolicyElementId id, SourcePosition position);
 
+  void NoteSourceRange(CelPolicyElementId id, std::optional<SourceRange> range,
+                       bool quoted);
+
   std::optional<SourcePosition> GetSourcePosition(CelPolicyElementId id) const;
 
+  std::optional<SourceRange> GetSourceRange(CelPolicyElementId id) const;
+
+  std::optional<bool> IsQuoted(CelPolicyElementId id) const;
+
   std::optional<SourceLocation> GetSourceLocation(CelPolicyElementId id) const;
 
   std::string DebugString() const;
 
  private:
   cel::SourcePtr policy_source_;
-  absl::flat_hash_map<CelPolicyElementId, SourcePosition> source_positions_;
+  absl::flat_hash_map<CelPolicyElementId, ElementSourceInfo> source_info_;
 };
 
 class ValueString {

policy/cel_policy_parse_result.cc

@@ -52,8 +52,12 @@ std::string CelPolicyIssue::ToDisplayString(
   std::string snippet;
   if (source != nullptr) {
     if (relative_position_) {
+      std::optional<SourceRange> range = source->GetSourceRange(element_id_);
       std::optional<SourcePosition> base =
           source->GetSourcePosition(element_id_);
+      if (range.has_value()) {
+        base = range->begin;
+      }
       if (element_id_ == -1) {
         base.emplace(0);
       }

policy/compiler.cc

@@ -19,6 +19,7 @@
 #include <memory>
 #include <optional>
 #include <string>
+#include <type_traits>
 #include <utility>
 #include <variant>
 #include <vector>
@@ -157,11 +158,19 @@ class IntermediateCompiledPolicy {
   void set_semantics(RuleSemantics semantics) { semantics_ = semantics; }
   RuleSemantics semantics() const { return semantics_; }
 
+  void set_policy_source(const CelPolicySource* absl_nullable src) {
+    policy_source_ = src;
+  }
+  const CelPolicySource* absl_nullable policy_source() const {
+    return policy_source_;
+  }
+
  private:
   std::string name_;
   std::string display_name_;
   std::string description_;
   RuleSemantics semantics_ = RuleSemantics::kFirstMatch;
+  const CelPolicySource* absl_nullable policy_source_ = nullptr;
 
   CompiledRule root_rule_;
 };
@@ -315,6 +324,33 @@ class PolicyCompiler {
     return src_->content()->description();
   }
 
+  absl::StatusOr<ValidationResult> CompileExpression(CelPolicyElementId id,
+                                                     absl::string_view val,
+                                                     const Compiler* env) {
+    std::unique_ptr<cel::Source> source;
+    if (src_ != nullptr && src_->content() != nullptr) {
+      std::optional<SourceRange> range;
+      range = src_->GetSourceRange(id);
+      bool use_subrange = !(src_->IsQuoted(id).value_or(true));
+      if (range.has_value() && use_subrange) {
+        source = std::make_unique<SourceSubrange>(*src_->content(), *range);
+      }
+    }
+
+    if (source == nullptr) {
+      // For quoted strings, the source should be generated from the interpreted
+      // YAML value.
+      CEL_ASSIGN_OR_RETURN(
+          source, cel::NewSource(val, std::string(GetSourceDescription())));
+    }
+    auto result = env->Compile(*source, &arena_);
+    if (!result.ok()) {
+      return result;
+    }
+    result->SetSource(std::move(source));
+    return result;
+  }
+
   void AdaptTypeCheckIssues(CelPolicyElementId id, const ValidationResult& r) {
     const Source* source = r.GetSource();
 
@@ -336,8 +372,8 @@ class PolicyCompiler {
       const cel::OutputBlock& output_block, const Compiler* env) {
     CompiledOutputBlock output;
     CEL_ASSIGN_OR_RETURN(auto output_validation,
-                         env->Compile(output_block.output().value(),
-                                      GetSourceDescription(), &arena_));
+                         CompileExpression(output_block.output().id(),
+                                           output_block.output().value(), env));
     AdaptTypeCheckIssues(output_block.output().id(), output_validation);
 
     cel::Type result_type = DynType();
@@ -352,9 +388,10 @@ class PolicyCompiler {
       }
     }
     if (output_block.explanation().has_value()) {
-      CEL_ASSIGN_OR_RETURN(auto explanation_validation,
-                           env->Compile(output_block.explanation()->value(),
-                                        GetSourceDescription(), &arena_));
+      CEL_ASSIGN_OR_RETURN(
+          auto explanation_validation,
+          CompileExpression(output_block.explanation()->id(),
+                            output_block.explanation()->value(), env));
       AdaptTypeCheckIssues(output_block.explanation()->id(),
                            explanation_validation);
       if (explanation_validation.IsValid()) {
@@ -378,8 +415,8 @@ class PolicyCompiler {
     c_match.id = match.id();
     if (match.condition().has_value()) {
       CEL_ASSIGN_OR_RETURN(auto validation,
-                           env->Compile(match.condition()->value(),
-                                        GetSourceDescription(), &arena_));
+                           CompileExpression(match.condition()->id(),
+                                             match.condition()->value(), env));
       AdaptTypeCheckIssues(match.condition()->id(), validation);
       if (validation.IsValid()) {
         CEL_ASSIGN_OR_RETURN(auto ast, validation.ReleaseAst());
@@ -422,9 +459,10 @@ class PolicyCompiler {
         continue;
       }
       std::string ident = absl::StrCat("variables.", name);
-      CEL_ASSIGN_OR_RETURN(auto validation,
-                           env->Compile(variable.expression().value(),
-                                        GetSourceDescription(), &arena_));
+      CEL_ASSIGN_OR_RETURN(
+          auto validation,
+          CompileExpression(variable.expression().id(),
+                            variable.expression().value(), env));
       AdaptTypeCheckIssues(variable.expression().id(), validation);
       if (!validation.IsValid()) {
         continue;
@@ -480,6 +518,7 @@ class PolicyCompiler {
   absl::Status CompilePolicy(const CelPolicy& policy,
                              IntermediateCompiledPolicy* out) {
     src_ = policy.source();
+    out->set_policy_source(src_);
     out->set_semantics(RuleSemantics::kFirstMatch);
     out->set_name(policy.name().value());
     out->set_display_name(
@@ -513,6 +552,21 @@ class FirstMatchComposer {
   std::unique_ptr<cel::Ast> ReleaseAst() { return std::move(ast_); }
 
  private:
+  SourcePosition GetAstOffset(CelPolicyElementId id) const {
+    if (icp_.policy_source() == nullptr) {
+      return 0;
+    }
+    if (auto range = icp_.policy_source()->GetSourceRange(id);
+        range.has_value()) {
+      return range->begin;
+    }
+    if (auto pos = icp_.policy_source()->GetSourcePosition(id);
+        pos.has_value()) {
+      return *pos;
+    }
+    return 0;
+  }
+
   using VariableScope = absl::flat_hash_map<std::string, int>;
 
   std::optional<int> ResolvePolicyVariable(absl::string_view reference);
@@ -733,7 +787,8 @@ absl::StatusOr<bool> FirstMatchComposer::ComposeRule(const CompiledRule& rule,
       MapVariables(condition);
       factory_.StartCopyContext();
       auto copy = factory_.Copy(condition.root_expr());
-      auto source_info = factory_.RemapSourceInfo(condition.source_info());
+      auto source_info = factory_.RemapSourceInfo(
+          condition.source_info(), GetAstOffset(match.condition->id));
       factory_.MergeSourceInfo(source_info);
       *insertion_point = factory_.NewCall("_?_:_", std::move(copy));
       insertion_point->mutable_call_expr().mutable_args().push_back(
@@ -792,7 +847,8 @@ absl::StatusOr<bool> FirstMatchComposer::ComposeProduction(
   MapVariables(ast);
   factory_.StartCopyContext();
   Expr to_insert = factory_.Copy(ast.root_expr());
-  auto source_info = factory_.RemapSourceInfo(ast.source_info());
+  auto source_info =
+      factory_.RemapSourceInfo(ast.source_info(), GetAstOffset(output_ast.id));
   factory_.MergeSourceInfo(source_info);
   insertion_expr = std::move(to_insert);
 
@@ -832,8 +888,9 @@ void FirstMatchComposer::ComposeRuleVariables(const CompiledRule& rule,
     MapVariables(ast);
     factory_.StartCopyContext();
     auto insertion = factory_.Copy(ast.root_expr());
-    // TODO(b/506179116): apply the position offsets here.
-    auto info = factory_.RemapSourceInfo(ast.source_info());
+    auto info = factory_.RemapSourceInfo(ast.source_info(),
+                                         GetAstOffset(variable.ast.id));
+    factory_.MergeSourceInfo(info);
     ABSL_DCHECK(init.has_list_expr());
     int index = init.mutable_list_expr().elements().size();
     init.mutable_list_expr().mutable_elements().push_back(