From 2bc0624495aba898885b04f12e839714d79b447d Mon Sep 17 00:00:00 2001
From: Michael Smith <spinda@users.noreply.github.com>
Date: Mon, 25 Sep 2023 14:37:14 -0400
Subject: [PATCH] Add `ParseError::primary_source_span` (#324)

Co-authored-by: Craig Disselkoen <cdiss@amazon.com>
---
 cedar-policy-core/src/parser/err.rs | 51 +++++++++++++++++++++++------
 cedar-policy/CHANGELOG.md           |  1 +
 2 files changed, 42 insertions(+), 10 deletions(-)

diff --git a/cedar-policy-core/src/parser/err.rs b/cedar-policy-core/src/parser/err.rs
index 1c0c6ac48..cf5b365aa 100644
--- a/cedar-policy-core/src/parser/err.rs
+++ b/cedar-policy-core/src/parser/err.rs
@@ -23,7 +23,7 @@ use std::ops::{Deref, DerefMut};
 use either::Either;
 use lalrpop_util as lalr;
 use lazy_static::lazy_static;
-use miette::{Diagnostic, LabeledSpan, Severity, SourceCode};
+use miette::{Diagnostic, LabeledSpan, Severity, SourceCode, SourceSpan};
 use smol_str::SmolStr;
 use thiserror::Error;
 
@@ -64,6 +64,22 @@ pub enum ParseError {
     ParseLiteral(#[from] ParseLiteralError),
 }
 
+impl ParseError {
+    /// Extract a primary source span locating the error, if one is available.
+    pub fn primary_source_span(&self) -> Option<SourceSpan> {
+        match self {
+            ParseError::ToCST(to_cst_err) => Some(to_cst_err.primary_source_span()),
+            ParseError::RestrictedExpr(restricted_expr_err) => match restricted_expr_err {
+                RestrictedExprError::InvalidRestrictedExpression { .. } => None,
+                RestrictedExprError::Parse(ParseErrors(parse_errs)) => {
+                    parse_errs.first().and_then(ParseError::primary_source_span)
+                }
+            },
+            ParseError::ToAST(_) | ParseError::ParseLiteral(_) => None,
+        }
+    }
+}
+
 /// Errors in the top-level parse literal entrypoint
 #[derive(Debug, Clone, PartialEq, Error, Eq)]
 pub enum ParseLiteralError {
@@ -358,6 +374,23 @@ pub struct ToCSTError {
 }
 
 impl ToCSTError {
+    /// Extract a primary source span locating the error.
+    pub fn primary_source_span(&self) -> SourceSpan {
+        match &self.err {
+            OwnedRawParseError::InvalidToken { location } => *location..*location,
+            OwnedRawParseError::UnrecognizedEof { location, .. } => *location..*location,
+            OwnedRawParseError::UnrecognizedToken {
+                token: (token_start, _, token_end),
+                ..
+            } => *token_start..*token_end,
+            OwnedRawParseError::ExtraToken {
+                token: (token_start, _, token_end),
+            } => *token_start..*token_end,
+            OwnedRawParseError::User { error } => error.info.0.clone(),
+        }
+        .into()
+    }
+
     pub(crate) fn from_raw_parse_err(err: RawParseError<'_>) -> Self {
         Self {
             err: err.map_token(|token| token.to_string()),
@@ -393,21 +426,19 @@ impl Diagnostic for ToCSTError {
     }
 
     fn labels(&self) -> Option<Box<dyn Iterator<Item = LabeledSpan> + '_>> {
+        let primary_source_span = self.primary_source_span();
         let labeled_span = match &self.err {
-            OwnedRawParseError::InvalidToken { location } => {
-                LabeledSpan::underline(*location..*location)
+            OwnedRawParseError::InvalidToken { .. } => LabeledSpan::underline(primary_source_span),
+            OwnedRawParseError::UnrecognizedEof { expected, .. } => {
+                LabeledSpan::new_with_span(expected_to_string(expected), primary_source_span)
             }
-            OwnedRawParseError::UnrecognizedEof { location, expected } => {
-                LabeledSpan::new_with_span(expected_to_string(expected), *location..*location)
+            OwnedRawParseError::UnrecognizedToken { expected, .. } => {
+                LabeledSpan::new_with_span(expected_to_string(expected), primary_source_span)
             }
-            OwnedRawParseError::UnrecognizedToken {
-                token: (token_start, _, token_end),
-                expected,
-            } => LabeledSpan::new_with_span(expected_to_string(expected), *token_start..*token_end),
             OwnedRawParseError::ExtraToken {
                 token: (token_start, _, token_end),
             } => LabeledSpan::underline(*token_start..*token_end),
-            OwnedRawParseError::User { error } => LabeledSpan::underline(error.info.0.clone()),
+            OwnedRawParseError::User { .. } => LabeledSpan::underline(primary_source_span),
         };
         Some(Box::new(iter::once(labeled_span)))
     }
diff --git a/cedar-policy/CHANGELOG.md b/cedar-policy/CHANGELOG.md
index 9966e7977..7fdd0119c 100644
--- a/cedar-policy/CHANGELOG.md
+++ b/cedar-policy/CHANGELOG.md
@@ -9,6 +9,7 @@
 - Support `Request`s with `Unknown` fields for partial evaluation.
 - Export the `cedar_policy_core::evaluator::{EvaluationError, EvaluationErrorKind}` and
   `cedar_policy_core::authorizer::AuthorizationError` error types.
+- Added an API to `ParseError` to quickly get the primary source span
 
 ### Changed