fix: count token spans in chars so non-ASCII text doesn't shift or crash error labels

CHANGELOG.md

@@ -29,6 +29,9 @@
 
 **Fixes**:
 
+- Error labels after non-ASCII text now point at the right column, rather than
+  drifting right or crashing `prqlc` when a label ran past the end of the query.
+  (@prql-bot, #6378)
 - `date.to_text` errors now name the format specifier the target dialect
   rejected — `format specifier %P is not supported for Postgres` rather than
   `PRQL doesn't support this format specifier`, whose span covers the whole

prqlc/prqlc-parser/src/lexer/mod.rs

@@ -68,7 +68,7 @@ pub fn lex_source_recovery(source: &str, source_id: u16) -> (Option<Vec<Token>>,
     let result = lexer().parse(source).into_result();
 
     match result {
-        Ok(tokens) => (Some(insert_start(tokens.to_vec())), vec![]),
+        Ok(tokens) => (Some(insert_start(to_char_spans(source, tokens))), vec![]),
         Err(errors) => {
             // Convert chumsky Simple errors to our Error type
             let errors = errors
@@ -86,7 +86,7 @@ pub fn lex_source(source: &str) -> Result<Tokens, Vec<E>> {
     let result = lexer().parse(source).into_result();
 
     match result {
-        Ok(tokens) => Ok(Tokens(insert_start(tokens.to_vec()))),
+        Ok(tokens) => Ok(Tokens(insert_start(to_char_spans(source, tokens)))),
         Err(errors) => {
             // Convert chumsky Simple errors to our Error type
             let errors = errors
@@ -99,6 +99,27 @@ pub fn lex_source(source: &str) -> Result<Tokens, Vec<E>> {
     }
 }
 
+/// Convert token spans from the byte offsets chumsky produces over `&str` to
+/// the char offsets the rest of the compiler (and ariadne) expects.
+fn to_char_spans(source: &str, mut tokens: Vec<Token>) -> Vec<Token> {
+    if source.is_ascii() {
+        return tokens;
+    }
+
+    // Char offset of every byte offset that falls on a char boundary, built
+    // once so the conversion stays linear in the source length.
+    let mut char_offsets = vec![0; source.len() + 1];
+    for (char_idx, (byte_idx, _)) in source.char_indices().enumerate() {
+        char_offsets[byte_idx] = char_idx;
+    }
+    char_offsets[source.len()] = source.chars().count();
+
+    for token in &mut tokens {
+        token.span = char_offsets[token.span.start]..char_offsets[token.span.end];
+    }
+    tokens
+}
+
 /// Insert a start token so later stages can treat the start of a file like a newline
 fn insert_start(tokens: Vec<Token>) -> Vec<Token> {
     std::iter::once(Token {

prqlc/prqlc-parser/src/lexer/test.rs

@@ -551,3 +551,23 @@ fn recovery_returns_no_tokens_on_error() {
     ]
     "#);
 }
+
+#[test]
+fn test_lex_source_non_ascii_spans() {
+    // Spans count chars, not bytes, so a multi-byte char doesn't shift the
+    // spans of the tokens after it.
+    assert_debug_snapshot!(lex_source("# é\nx 'ü' y"), @r#"
+    Ok(
+        Tokens(
+            [
+                0..0: Start,
+                0..3: Comment(" é"),
+                3..4: NewLine,
+                4..5: Ident("x"),
+                6..9: Literal(String("ü")),
+                10..11: Ident("y"),
+            ],
+        ),
+    )
+    "#);
+}

prqlc/prqlc-parser/src/parser/interpolation.rs

@@ -11,14 +11,18 @@ pub(crate) fn parse(string: String, span_base: Span) -> Result<Vec<InterpolateIt
 
     let (output, errors) = res.into_output_errors();
 
+    // chumsky's spans over `&str` are byte offsets, but `span_base` and the
+    // rest of the compiler count chars.
+    let to_char = |byte_idx: usize| string[..byte_idx].chars().count();
+
     if !errors.is_empty() {
         return Err(errors
             .into_iter()
             .map(|e| {
                 // Adjust span to be relative to span_base
                 let span = Span {
-                    start: span_base.start + e.span().start,
-                    end: span_base.start + e.span().end,
+                    start: span_base.start + to_char(e.span().start),
+                    end: span_base.start + to_char(e.span().end),
                     source_id: span_base.source_id,
                 };
 
@@ -76,8 +80,8 @@ pub(crate) fn parse(string: String, span_base: Span) -> Result<Vec<InterpolateIt
             InterpolateItem::Expr { expr, format } => {
                 let adjusted_expr = Box::new(Expr {
                     span: expr.span.map(|s| Span {
-                        start: span_base.start + s.start,
-                        end: span_base.start + s.end,
+                        start: span_base.start + to_char(s.start),
+                        end: span_base.start + to_char(s.end),
                         source_id: span_base.source_id,
                     }),
                     ..(*expr)

prqlc/prqlc-parser/src/parser/mod.rs

@@ -35,14 +35,14 @@ pub fn parse_lr_to_pr(source_id: u16, lr: Vec<lr::Token>) -> (Option<Vec<pr::Stm
         })
         .collect();
 
-    // Use built-in Input impl for &[Token], then map_span to convert token indices to byte spans
+    // Use built-in Input impl for &[Token], then map_span to convert token indices to source spans
     let input = semantic_tokens
         .as_slice()
         .map_span(|simple_span: SimpleSpan| {
             let start_idx = simple_span.start();
             let end_idx = simple_span.end();
 
-            // Convert token indices to byte offsets in the source file
+            // Convert token indices to char offsets in the source file
             let start = semantic_tokens
                 .get(start_idx)
                 .map(|t| t.span.start)

prqlc/prqlc-parser/src/test.rs

@@ -1624,9 +1624,9 @@ fn test_unicode() {
             args:
               - Ident:
                   - tète
-                span: "0:5-10"
-          span: "0:0-10"
-      span: "0:0-10"
+                span: "0:5-9"
+          span: "0:0-9"
+      span: "0:0-9"
     "#);
 }
 

prqlc/prqlc/tests/integration/error_messages.rs

@@ -806,3 +806,51 @@ fn unknown_named_arg() {
     ───╯
     ");
 }
+
+#[test]
+fn test_error_after_non_ascii() {
+    // Multi-byte chars before an error mustn't shift its label, or push the
+    // span past the end of the source.
+    assert_snapshot!(compile(r#"
+    from t
+    derive {x = "éééééééééé"}
+    select {x, foo.bar.baz}
+    "#).unwrap_err(), @"
+    Error:
+       ╭─[ :4:16 ]
+       │
+     4 │     select {x, foo.bar.baz}
+       │                ─────┬─────
+       │                     ╰─────── Unknown name `foo.bar.baz`
+       │
+       │ Help: available columns: x
+    ───╯
+    ");
+
+    assert_snapshot!(compile(r#"
+    # café
+    from t
+    select {foo.bar.baz}
+    "#).unwrap_err(), @"
+    Error:
+       ╭─[ :4:13 ]
+       │
+     4 │     select {foo.bar.baz}
+       │             ─────┬─────
+       │                  ╰─────── Unknown name `foo.bar.baz`
+    ───╯
+    ");
+
+    assert_snapshot!(compile(r#"
+    from t
+    derive {x = f"éééé{foo.}"}
+    "#).unwrap_err(), @r#"
+    Error:
+       ╭─[ :3:28 ]
+       │
+     3 │     derive {x = f"éééé{foo.}"}
+       │                            ┬
+       │                            ╰── expected interpolated string or interp:backticks, but found "}"
+    ───╯
+    "#);
+}