[codex] Expose token byte offsets

src/lexer.rs

@@ -353,25 +353,30 @@ where
                 None
             }
         });
-        let source_text = if text.is_none() && stop != usize::MAX {
+        let source_interval = if text.is_none() && stop != usize::MAX && self.token_start <= stop {
             self.input
                 .text_source_interval(TextInterval::new(self.token_start, stop))
-                .and_then(|(input, start_byte, stop_byte)| {
-                    Some(crate::token::TokenSourceText {
-                        input,
-                        start_byte: u32::try_from(start_byte).ok()?,
-                        stop_byte: u32::try_from(stop_byte).ok()?,
-                    })
-                })
         } else {
             None
         };
+        let source_text = source_interval
+            .as_ref()
+            .and_then(|(input, start_byte, stop_byte)| {
+                Some(crate::token::TokenSourceText {
+                    input: Rc::clone(input),
+                    start_byte: u32::try_from(*start_byte).ok()?,
+                    stop_byte: u32::try_from(*stop_byte).ok()?,
+                })
+            });
+        let source_byte_span = source_text
+            .as_ref()
+            .map(|source_text| (source_text.start_byte, source_text.stop_byte));
         let text = text.or_else(|| {
             source_text
                 .is_none()
                 .then(|| self.input.text(TextInterval::new(self.token_start, stop)))
         });
-        self.factory.create(TokenSpec {
+        let mut token = self.factory.create(TokenSpec {
             token_type,
             channel,
             start: self.token_start,
@@ -381,7 +386,13 @@ where
             text,
             source_text,
             source_name: self.input.source_name(),
-        })
+        });
+        if let Some((start_byte, stop_byte)) =
+            source_byte_span.or_else(|| self.token_byte_span(stop))
+        {
+            token = token.with_byte_span(start_byte, stop_byte);
+        }
+        token
     }
 
     /// Returns the current token text from the token start through the input
@@ -428,12 +439,44 @@ where
 
     /// Builds the synthetic EOF token at the current input cursor.
     pub fn eof_token(&self) -> CommonToken {
-        CommonToken::eof(
+        let token = CommonToken::eof(
             self.input.source_name(),
             self.input.index(),
             self.line,
             self.column,
-        )
+        );
+        match self.eof_byte_offset() {
+            Some(byte_offset) => token.with_byte_span(byte_offset, byte_offset),
+            None => token,
+        }
+    }
+
+    fn eof_byte_offset(&self) -> Option<u32> {
+        self.byte_offset_at(self.input.index())
+    }
+
+    fn token_byte_span(&self, stop: usize) -> Option<(u32, u32)> {
+        if stop != usize::MAX && self.token_start <= stop {
+            let (_, start_byte, stop_byte) = self
+                .input
+                .text_source_interval(TextInterval::new(self.token_start, stop))?;
+            return Some((
+                u32::try_from(start_byte).ok()?,
+                u32::try_from(stop_byte).ok()?,
+            ));
+        }
+        let byte_offset = self.byte_offset_at(self.token_start)?;
+        Some((byte_offset, byte_offset))
+    }
+
+    fn byte_offset_at(&self, index: usize) -> Option<u32> {
+        let byte_offset = if index == 0 {
+            0
+        } else {
+            let previous = TextInterval::new(index - 1, index - 1);
+            self.input.text_source_interval(previous)?.2
+        };
+        u32::try_from(byte_offset).ok()
     }
 }
 
@@ -606,3 +649,77 @@ where
 fn lexer_dfa_edge_label(symbol: i32) -> Option<String> {
     char::from_u32(symbol.cast_unsigned()).map(|ch| format!("'{ch}'"))
 }
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+    use crate::char_stream::InputStream;
+    use crate::recognizer::RecognizerData;
+    use crate::token::{DEFAULT_CHANNEL, Token};
+    use crate::vocabulary::Vocabulary;
+
+    #[test]
+    fn eof_token_uses_utf8_byte_offset_after_non_ascii_input() {
+        let data = RecognizerData::new(
+            "T",
+            Vocabulary::new(
+                std::iter::empty::<Option<&str>>(),
+                std::iter::empty::<Option<&str>>(),
+                std::iter::empty::<Option<&str>>(),
+            ),
+        );
+        let mut lexer = BaseLexer::new(InputStream::new("β"), data);
+        lexer.consume_char();
+
+        let token = lexer.eof_token();
+
+        assert_eq!(token.start(), 1);
+        assert_eq!(token.stop(), 0);
+        assert_eq!(token.text(), Some("<EOF>"));
+        assert_eq!(token.byte_span(), 2..2);
+    }
+
+    #[test]
+    fn eof_rule_token_uses_utf8_byte_offset_after_non_ascii_input() {
+        let data = RecognizerData::new(
+            "T",
+            Vocabulary::new(
+                std::iter::empty::<Option<&str>>(),
+                std::iter::empty::<Option<&str>>(),
+                std::iter::empty::<Option<&str>>(),
+            ),
+        );
+        let mut lexer = BaseLexer::new(InputStream::new("β"), data);
+        lexer.consume_char();
+        lexer.begin_token();
+
+        let token = lexer.emit_with_stop(1, DEFAULT_CHANNEL, 0, Some("<EOF>".to_owned()));
+
+        assert_eq!(token.start(), 1);
+        assert_eq!(token.stop(), 0);
+        assert_eq!(token.text(), Some("<EOF>"));
+        assert_eq!(token.byte_span(), 2..2);
+    }
+
+    #[test]
+    fn emit_implicit_text_uses_utf8_byte_span_for_non_ascii_input() {
+        let data = RecognizerData::new(
+            "T",
+            Vocabulary::new(
+                std::iter::empty::<Option<&str>>(),
+                std::iter::empty::<Option<&str>>(),
+                std::iter::empty::<Option<&str>>(),
+            ),
+        );
+        let mut lexer = BaseLexer::new(InputStream::new("β"), data);
+        lexer.begin_token();
+        lexer.consume_char();
+
+        let token = lexer.emit(1, DEFAULT_CHANNEL, None);
+
+        assert_eq!(token.start(), 0);
+        assert_eq!(token.stop(), 0);
+        assert_eq!(token.text(), Some("β"));
+        assert_eq!(token.byte_span(), 0..2);
+    }
+}

src/token.rs

@@ -1,5 +1,6 @@
 use crate::char_stream::TextInterval;
 use std::fmt;
+use std::ops::Range;
 use std::rc::Rc;
 
 pub const TOKEN_EOF: i32 = -1;
@@ -37,17 +38,47 @@ impl From<i32> for TokenChannel {
 pub trait Token: fmt::Debug {
     fn token_type(&self) -> i32;
     fn channel(&self) -> i32;
+    /// Zero-based absolute start index measured in Unicode scalar values.
     fn start(&self) -> usize;
+    /// Zero-based absolute inclusive stop index measured in Unicode scalar
+    /// values.
     fn stop(&self) -> usize;
     fn token_index(&self) -> isize;
+    /// One-based source line where the token starts.
     fn line(&self) -> usize;
+    /// Zero-based source column where the token starts, measured in Unicode
+    /// scalar values from the start of `line`.
     fn column(&self) -> usize;
     fn text(&self) -> Option<&str>;
     fn source_name(&self) -> &str;
 
     fn interval(&self) -> TextInterval {
         TextInterval::new(self.start(), self.stop())
     }
+
+    /// Zero-based absolute start offset measured in UTF-8 bytes.
+    ///
+    /// The default implementation treats the character index as a byte offset,
+    /// which is exact for ASCII and preserves compatibility for token
+    /// implementations that do not expose source byte bounds.
+    fn start_byte(&self) -> usize {
+        self.start()
+    }
+
+    /// Zero-based exclusive end offset measured in UTF-8 bytes.
+    ///
+    /// Unlike [`Self::stop`], this is exclusive so
+    /// `token.start_byte()..token.stop_byte()` can slice the original UTF-8
+    /// source when the token carries source byte bounds. The default
+    /// implementation treats character indices as byte offsets.
+    fn stop_byte(&self) -> usize {
+        default_stop_byte(self.start(), self.stop())
+    }
+
+    /// Zero-based UTF-8 byte span for the token text.
+    fn byte_span(&self) -> Range<usize> {
+        self.start_byte()..self.stop_byte()
+    }
 }
 
 #[derive(Clone, Debug, Eq, PartialEq)]
@@ -60,9 +91,16 @@ pub struct CommonToken {
     line: usize,
     column: usize,
     text: Option<TokenText>,
+    byte_span: Option<TokenByteSpan>,
     source_name: Rc<str>,
 }
 
+#[derive(Clone, Copy, Debug, Eq, PartialEq)]
+struct TokenByteSpan {
+    start_byte: u32,
+    stop_byte: u32,
+}
+
 #[derive(Clone, Debug, Eq, PartialEq)]
 enum TokenText {
     Explicit(Rc<str>),
@@ -117,6 +155,7 @@ impl CommonToken {
             line: 1,
             column: 0,
             text: None,
+            byte_span: None,
             source_name: Rc::from(""),
         }
     }
@@ -131,6 +170,7 @@ impl CommonToken {
             line,
             column,
             text: Some(TokenText::Explicit(Rc::from("<EOF>"))),
+            byte_span: None,
             source_name: source_name.into(),
         }
     }
@@ -153,6 +193,23 @@ impl CommonToken {
             start_byte,
             stop_byte,
         });
+        self.byte_span = Some(TokenByteSpan {
+            start_byte,
+            stop_byte,
+        });
+        self
+    }
+
+    #[must_use]
+    pub(crate) fn with_byte_span(mut self, start_byte: u32, stop_byte: u32) -> Self {
+        debug_assert!(
+            start_byte <= stop_byte,
+            "invalid token byte span: start={start_byte}, stop={stop_byte}"
+        );
+        self.byte_span = Some(TokenByteSpan {
+            start_byte,
+            stop_byte,
+        });
         self
     }
 
@@ -185,6 +242,16 @@ impl CommonToken {
     pub const fn set_token_index(&mut self, token_index: isize) {
         self.token_index = token_index;
     }
+
+    const fn source_byte_span(&self) -> Option<Range<usize>> {
+        match self.byte_span {
+            Some(TokenByteSpan {
+                start_byte,
+                stop_byte,
+            }) => Some(start_byte as usize..stop_byte as usize),
+            None => None,
+        }
+    }
 }
 
 impl Token for CommonToken {
@@ -223,6 +290,16 @@ impl Token for CommonToken {
     fn source_name(&self) -> &str {
         self.source_name.as_ref()
     }
+
+    fn start_byte(&self) -> usize {
+        self.source_byte_span()
+            .map_or(self.start, |byte_span| byte_span.start)
+    }
+
+    fn stop_byte(&self) -> usize {
+        self.source_byte_span()
+            .map_or_else(|| default_stop_byte(self.start, self.stop), |span| span.end)
+    }
 }
 
 impl fmt::Display for CommonToken {
@@ -257,6 +334,13 @@ fn display_token_boundary(value: usize) -> String {
     }
 }
 
+const fn default_stop_byte(start: usize, stop: usize) -> usize {
+    match stop.checked_add(1) {
+        Some(end) if end >= start => end,
+        Some(_) | None => start,
+    }
+}
+
 /// Escapes token text the way ANTLR's token display format expects.
 ///
 /// Debug escaping is close but not identical: ANTLR leaves ordinary
@@ -379,4 +463,37 @@ mod tests {
         let token = CommonToken::eof("", 0, 1, 0);
         assert_eq!(token.to_string(), "[@-1,0:-1='<EOF>',<-1>,1:0]");
     }
+
+    #[test]
+    fn source_backed_token_exposes_utf8_byte_span() {
+        let source: Rc<str> = Rc::from("éβz");
+        let token = CommonToken::new(1)
+            .with_span(1, 1)
+            .with_source_text(source, 2, 4);
+
+        assert_eq!(token.start(), 1);
+        assert_eq!(token.stop(), 1);
+        assert_eq!(token.start_byte(), 2);
+        assert_eq!(token.stop_byte(), 4);
+        assert_eq!(token.byte_span(), 2..4);
+        assert_eq!(token.text(), Some("β"));
+    }
+
+    #[test]
+    fn explicit_text_byte_span_falls_back_to_character_span() {
+        let token = CommonToken::new(1).with_text("β").with_span(3, 3);
+
+        assert_eq!(token.byte_span(), 3..4);
+    }
+
+    #[test]
+    fn explicit_text_can_carry_utf8_byte_span() {
+        let token = CommonToken::new(TOKEN_EOF)
+            .with_text("<EOF>")
+            .with_span(1, 0)
+            .with_byte_span(2, 2);
+
+        assert_eq!(token.text(), Some("<EOF>"));
+        assert_eq!(token.byte_span(), 2..2);
+    }
 }