Draft: improve source mapping in policy compiler.

common/source.cc

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

policy/BUILD

@@ -27,10 +27,12 @@ cc_library(
     ],
     deps = [
         "//common:source",
+        "//policy/internal:alignment_table",
         "@com_google_absl//absl/base:nullability",
         "@com_google_absl//absl/container:flat_hash_map",
         "@com_google_absl//absl/log:absl_check",
         "@com_google_absl//absl/strings",
+        "@com_google_absl//absl/types:span",
     ],
 )
 
@@ -41,6 +43,7 @@ cc_test(
         ":cel_policy",
         "//common:source",
         "//internal:testing",
+        "//policy/internal:alignment_table",
         "@com_google_absl//absl/strings",
     ],
 )
@@ -84,6 +87,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 +194,7 @@ cc_test(
         ":compiler",
         ":yaml_policy_parser",
         "//common:ast",
+        "//common:ast_proto",
         "//common:decl",
         "//common:navigable_ast",
         "//common:source",
@@ -211,6 +216,8 @@ cc_test(
         "//runtime:runtime_builder",
         "//runtime:runtime_options",
         "//runtime:standard_runtime_builder_factory",
+        "//tools:cel_unparser",
+        "@com_google_absl//absl/log:absl_check",
         "@com_google_absl//absl/status",
         "@com_google_absl//absl/status:status_matchers",
         "@com_google_absl//absl/status:statusor",

policy/cel_policy.cc

@@ -27,6 +27,7 @@
 #include "absl/strings/str_join.h"
 #include "absl/strings/str_split.h"
 #include "absl/strings/string_view.h"
+#include "absl/types/span.h"
 #include "common/source.h"
 
 namespace cel {
@@ -59,34 +60,111 @@ 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,
+    std::vector<policy_internal::AlignmentPoint> alignment_table) {
+  ElementSourceInfo& info = source_info_[id];
+  info.range = range;
+  info.quoted = quoted;
+  info.alignment_table = std::move(alignment_table);
+  if (range.has_value() && info.position == -1) {
+    info.position = range->begin;
+  }
+}
+
+void CelPolicySource::NoteAlignmentTable(
+    CelPolicyElementId id,
+    std::vector<policy_internal::AlignmentPoint> alignment_table) {
+  source_info_[id].alignment_table = std::move(alignment_table);
+}
+
+absl::Span<const policy_internal::AlignmentPoint>
+CelPolicySource::GetAlignmentTable(CelPolicyElementId id) const {
+  auto it = source_info_.find(id);
+  if (it == source_info_.end()) {
+    return {};
+  }
+  return absl::MakeConstSpan(it->second.alignment_table);
+}
+
+std::optional<SourcePosition> CelPolicySource::MapPosition(
+    CelPolicyElementId id, SourcePosition relative_position) const {
+  auto it = source_info_.find(id);
+  if (it != source_info_.end() && !it->second.alignment_table.empty()) {
+    const auto& points = it->second.alignment_table;
+    auto pt_it = std::upper_bound(
+        points.begin(), points.end(), relative_position,
+        [](SourcePosition pos, const policy_internal::AlignmentPoint& pt) {
+          return pos < pt.value_position;
+        });
+    if (pt_it == points.begin()) {
+      return points.front().source_position +
+             (relative_position - points.front().value_position);
+    }
+    --pt_it;
+    return pt_it->source_position + (relative_position - pt_it->value_position);
+  }
+  std::optional<SourceRange> range = GetSourceRange(id);
+  std::optional<SourcePosition> base = GetSourcePosition(id);
+  if (range.has_value()) {
+    base = range->begin;
+  }
+  if (id == -1 && !base.has_value()) {
+    base.emplace(0);
+  }
+  if (base.has_value()) {
+    return *base + relative_position;
+  }
+  return std::nullopt;
 }
 
 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.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;
+  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

@@ -28,30 +28,59 @@
 #include "absl/container/flat_hash_map.h"
 #include "absl/log/absl_check.h"
 #include "absl/strings/string_view.h"
+#include "absl/types/span.h"
 #include "common/source.h"
+#include "policy/internal/alignment_table.h"
 
 namespace cel {
 
 using CelPolicyElementId = int32_t;
 
+struct ElementSourceInfo {
+  SourcePosition position = -1;
+  std::optional<SourceRange> range;
+  bool quoted = false;
+  std::vector<policy_internal::AlignmentPoint> alignment_table;
+};
+
 class CelPolicySource {
  public:
   explicit CelPolicySource(cel::SourcePtr policy_source)
       : policy_source_(std::move(policy_source)) {}
 
   const Source* absl_nonnull content() const { return policy_source_.get(); }
 
+  // Note: NoteSource* and NoteAlignmentTable methods are intended for internal
+  // use only and are not supported for clients outside of the policy package.
   void NoteSourcePosition(CelPolicyElementId id, SourcePosition position);
 
+  void NoteSourceRange(
+      CelPolicyElementId id, std::optional<SourceRange> range, bool quoted,
+      std::vector<policy_internal::AlignmentPoint> alignment_table = {});
+
+  void NoteAlignmentTable(
+      CelPolicyElementId id,
+      std::vector<policy_internal::AlignmentPoint> alignment_table);
+
   std::optional<SourcePosition> GetSourcePosition(CelPolicyElementId id) const;
 
+  std::optional<SourceRange> GetSourceRange(CelPolicyElementId id) const;
+
+  std::optional<bool> IsQuoted(CelPolicyElementId id) const;
+
+  absl::Span<const policy_internal::AlignmentPoint> GetAlignmentTable(
+      CelPolicyElementId id) const;
+
+  std::optional<SourcePosition> MapPosition(
+      CelPolicyElementId id, SourcePosition relative_position) 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,15 +52,11 @@ std::string CelPolicyIssue::ToDisplayString(
   std::string snippet;
   if (source != nullptr) {
     if (relative_position_) {
-      std::optional<SourcePosition> base =
-          source->GetSourcePosition(element_id_);
-      if (element_id_ == -1) {
-        base.emplace(0);
-      }
-      if (base) {
-        location = source->content()
-                       ->GetLocation(*base + *relative_position_)
-                       .value_or(SourceLocation{});
+      if (std::optional<SourcePosition> abs_pos =
+              source->MapPosition(element_id_, *relative_position_);
+          abs_pos.has_value()) {
+        location =
+            source->content()->GetLocation(*abs_pos).value_or(SourceLocation{});
       }
     } else {
       location =