gh-153569: return tokenizer tokens as source spans

Author: pablogsal

Parser/lexer/lexer.c

@@ -12,8 +12,21 @@
 
 
 #define MAKE_TOKEN(token_type) _PyLexer_token_setup(tok, token, token_type, p_start, p_end)
-#define MAKE_TYPE_COMMENT_TOKEN(token_type, col_offset, end_col_offset) (\
-                _PyLexer_type_comment_token_setup(tok, token, token_type, col_offset, end_col_offset, p_start, p_end))
+
+static int
+type_comment_token_setup(struct tok_state *tok, struct token *token, int type,
+                         int col_offset, int end_col_offset,
+                         const char *start, const char *end)
+{
+    _PyLexer_token_setup(tok, token, type, start, end);
+    token->start_loc = (_PyTok_Loc){tok->lineno, col_offset};
+    token->end_loc = (_PyTok_Loc){tok->lineno, end_col_offset};
+    return type;
+}
+
+#define MAKE_TYPE_COMMENT_TOKEN(token_type, col_offset, end_col_offset) \
+    type_comment_token_setup(tok, token, token_type, col_offset, \
+                             end_col_offset, p_start, p_end)
 
 /* Spaces in this constant are treated as "zero or more spaces or tabs" when
    tokenizing. */

Parser/lexer/lexer.h

@@ -7,4 +7,23 @@ int _PyLexer_update_ftstring_expr(struct tok_state *tok, char cur);
 
 int _PyTokenizer_Get(struct tok_state *, struct token *);
 
+/* The view points into the current input window. The next
+   _PyTokenizer_Get() call may discard it. */
+static inline const char *
+_PyToken_TextView(const struct tok_state *tok, const struct token *token,
+                  Py_ssize_t *length)
+{
+    assert(length != NULL);
+    if (!_PyTok_SpanIsValid(token->span)) {
+        *length = 0;
+        return "";
+    }
+    assert(tok->buf != NULL);
+    assert(tok->inp >= tok->buf);
+    assert(token->span.start >= tok->buf_offset);
+    assert(token->span.end - tok->buf_offset <= tok->inp - tok->buf);
+    *length = token->span.end - token->span.start;
+    return tok->buf + (token->span.start - tok->buf_offset);
+}
+
 #endif

Parser/lexer/state.c

@@ -51,7 +51,6 @@ _PyTokenizer_tok_new(void)
     tok->comment_newline = 0;
     tok->implicit_newline = 0;
     _PyTok_SourceInit(&tok->source);
-    _PyTok_CursorInit(&tok->reader_cursor, &tok->source);
     tok->reader = NULL;
     tok->tok_mode_stack[0] = (tokenizer_mode){.kind =TOK_REGULAR_MODE, .quote='\0', .quote_size = 0, .in_debug=0};
     tok->tok_mode_stack_index = 0;
@@ -101,41 +100,46 @@ _PyToken_Free(struct token *token) {
 
 void
 _PyToken_Init(struct token *token) {
+#ifdef Py_DEBUG
+    token->span = (_PyTok_Span){-1, -1};
+    token->start_loc = (_PyTok_Loc){-1, -1};
+    token->end_loc = (_PyTok_Loc){-1, -1};
+#endif
     token->metadata = NULL;
 }
 
-int
-_PyLexer_type_comment_token_setup(struct tok_state *tok, struct token *token, int type, int col_offset,
-                         int end_col_offset, const char *start, const char *end)
+static inline _PyTok_Span
+buffer_span(const struct tok_state *tok, const char *start, const char *end)
 {
-    token->level = tok->level;
-    token->lineno = token->end_lineno = tok->lineno;
-    token->col_offset = col_offset;
-    token->end_col_offset = end_col_offset;
-    token->start = start;
-    token->end = end;
-    return type;
+    if (start == NULL) {
+        assert(end == NULL);
+        return (_PyTok_Span){-1, -1};
+    }
+    assert(end != NULL);
+    const char *base = tok->buf;
+    assert(base != NULL);
+    assert(tok->inp >= base);
+    Py_ssize_t start_offset = start - base;
+    Py_ssize_t end_offset = end - base;
+    assert(start_offset >= 0 && start_offset <= end_offset);
+    assert(end_offset <= tok->inp - base);
+    assert(tok->buf_offset <= PY_SSIZE_T_MAX - end_offset);
+    return _PyTok_SpanFromBounds(
+        tok->buf_offset + start_offset, tok->buf_offset + end_offset);
 }
 
 int
 _PyLexer_token_setup(struct tok_state *tok, struct token *token, int type, const char *start, const char *end)
 {
-    assert((start == NULL && end == NULL) || (start != NULL && end != NULL));
     token->level = tok->level;
-    if (ISSTRINGLIT(type)) {
-        token->lineno = tok->first_lineno;
-    }
-    else {
-        token->lineno = tok->lineno;
-    }
-    token->end_lineno = tok->lineno;
-    token->col_offset = token->end_col_offset = -1;
-    token->start = start;
-    token->end = end;
+    token->span = buffer_span(tok, start, end);
+    int lineno = ISSTRINGLIT(type) ? tok->first_lineno : tok->lineno;
+    token->start_loc = (_PyTok_Loc){lineno, -1};
+    token->end_loc = (_PyTok_Loc){tok->lineno, -1};
 
     if (start != NULL && end != NULL) {
-        token->col_offset = tok->starting_col_offset;
-        token->end_col_offset = tok->col_offset;
+        token->start_loc.byte_col = tok->starting_col_offset;
+        token->end_loc.byte_col = tok->col_offset;
     }
     return type;
 }

Parser/lexer/state.h

@@ -2,7 +2,7 @@
 #define _PY_LEXER_H_
 
 #include "object.h"
-#include "../tokenizer/cursor.h"
+#include "../tokenizer/source.h"
 
 #define MAXINDENT 100       /* Max indentation level */
 #define MAXLEVEL 200        /* Max parentheses level */
@@ -23,8 +23,9 @@ enum interactive_underflow_t {
 
 struct token {
     int level;
-    int lineno, col_offset, end_lineno, end_col_offset;
-    const char *start, *end;
+    _PyTok_Span span;
+    _PyTok_Loc start_loc;
+    _PyTok_Loc end_loc;
     PyObject *metadata;
 };
 
@@ -69,7 +70,7 @@ typedef struct _tokenizer_mode {
 struct tok_state {
     /* Input state; buf <= cur <= inp */
     /* NB an entire line is held in the buffer */
-    char *buf;          /* Input buffer, or NULL; malloc'ed if fp != NULL or readline != NULL */
+    char *buf;          /* Owned for file/readline input; source-backed otherwise. */
     char *cur;          /* Next character in buffer */
     char *inp;          /* End of data in buffer */
     _PyTok_Off buf_offset; /* Logical offset of buf[0]. */
@@ -110,7 +111,6 @@ struct tok_state {
     char* str;          /* Source string being tokenized (if tokenizing from a string)*/
 
     _PyTok_SourceText source;
-    _PyTok_Cursor reader_cursor;
     struct _PyTok_Reader *reader;
 
     int type_comments;      /* Whether to look for type comments */
@@ -129,8 +129,6 @@ struct tok_state {
 #endif
 };
 
-int _PyLexer_type_comment_token_setup(struct tok_state *tok, struct token *token, int type, int col_offset,
-                         int end_col_offset, const char *start, const char *end);
 int _PyLexer_token_setup(struct tok_state *tok, struct token *token, int type, const char *start, const char *end);
 
 struct tok_state *_PyTokenizer_tok_new(void);

Parser/pegen.c

@@ -171,18 +171,17 @@ growable_comment_array_deallocate(growable_comment_array *arr) {
 }
 
 static int
-_get_keyword_or_name_type(Parser *p, struct token *new_token)
+_get_keyword_or_name_type(Parser *p, const char *text, Py_ssize_t length)
 {
-    Py_ssize_t name_len = new_token->end_col_offset - new_token->col_offset;
-    assert(name_len > 0);
+    assert(length > 0);
 
-    if (name_len >= p->n_keyword_lists ||
-        p->keywords[name_len] == NULL ||
-        p->keywords[name_len]->type == -1) {
+    if (length >= p->n_keyword_lists ||
+        p->keywords[length] == NULL ||
+        p->keywords[length]->type == -1) {
         return NAME;
     }
-    for (KeywordToken *k = p->keywords[name_len]; k != NULL && k->type != -1; k++) {
-        if (strncmp(k->str, new_token->start, (size_t)name_len) == 0) {
+    for (KeywordToken *k = p->keywords[length]; k != NULL && k->type != -1; k++) {
+        if (memcmp(k->str, text, (size_t)length) == 0) {
             return k->type;
         }
     }
@@ -193,8 +192,11 @@ static int
 initialize_token(Parser *p, Token *parser_token, struct token *new_token, int token_type) {
     assert(parser_token != NULL);
 
-    parser_token->type = (token_type == NAME) ? _get_keyword_or_name_type(p, new_token) : token_type;
-    parser_token->bytes = PyBytes_FromStringAndSize(new_token->start, new_token->end - new_token->start);
+    Py_ssize_t length;
+    const char *text = _PyToken_TextView(p->tok, new_token, &length);
+    parser_token->type = token_type == NAME
+        ? _get_keyword_or_name_type(p, text, length) : token_type;
+    parser_token->bytes = PyBytes_FromStringAndSize(text, length);
     if (parser_token->bytes == NULL) {
         return -1;
     }
@@ -214,12 +216,14 @@ initialize_token(Parser *p, Token *parser_token, struct token *new_token, int to
     }
 
     parser_token->level = new_token->level;
-    parser_token->lineno = new_token->lineno;
-    parser_token->col_offset = p->tok->lineno == p->starting_lineno ? p->starting_col_offset + new_token->col_offset
-                                                                    : new_token->col_offset;
-    parser_token->end_lineno = new_token->end_lineno;
-    parser_token->end_col_offset = p->tok->lineno == p->starting_lineno ? p->starting_col_offset + new_token->end_col_offset
-                                                                 : new_token->end_col_offset;
+    parser_token->lineno = new_token->start_loc.lineno;
+    parser_token->col_offset = p->tok->lineno == p->starting_lineno
+        ? p->starting_col_offset + new_token->start_loc.byte_col
+        : new_token->start_loc.byte_col;
+    parser_token->end_lineno = new_token->end_loc.lineno;
+    parser_token->end_col_offset = p->tok->lineno == p->starting_lineno
+        ? p->starting_col_offset + new_token->end_loc.byte_col
+        : new_token->end_loc.byte_col;
 
     p->fill += 1;
 
@@ -261,13 +265,14 @@ _PyPegen_fill_token(Parser *p)
 
     // Record and skip '# type: ignore' comments
     while (type == TYPE_IGNORE) {
-        Py_ssize_t len = new_token.end_col_offset - new_token.col_offset;
+        Py_ssize_t len;
+        const char *text = _PyToken_TextView(p->tok, &new_token, &len);
         char *tag = PyMem_Malloc((size_t)len + 1);
         if (tag == NULL) {
             PyErr_NoMemory();
             goto error;
         }
-        strncpy(tag, new_token.start, (size_t)len);
+        memcpy(tag, text, (size_t)len);
         tag[len] = '\0';
         // Ownership of tag passes to the growable array
         if (!growable_comment_array_add(&p->type_ignore_comments, p->tok->lineno, tag)) {

Parser/tokenizer/source.h

@@ -5,7 +5,8 @@
 
 typedef Py_ssize_t _PyTok_Off;
 
-/* Half-open byte offsets into a _PyTok_SourceText. */
+/* Spans use half-open logical byte offsets into decoded input. Their backing
+   storage may retain only the current input window. */
 typedef struct {
     _PyTok_Off start;
     _PyTok_Off end;

Python/Python-tokenize.c

@@ -203,14 +203,15 @@ _get_current_line(tokenizeriterobject *it, const char *line_start, Py_ssize_t si
 }
 
 static void
-_get_col_offsets(tokenizeriterobject *it, struct token token, const char *line_start,
-                 PyObject *line, int line_changed, Py_ssize_t lineno, Py_ssize_t end_lineno,
+_get_col_offsets(tokenizeriterobject *it, const char *token_start,
+                 const char *token_end, const char *line_start, PyObject *line,
+                 int line_changed, Py_ssize_t lineno, Py_ssize_t end_lineno,
                  Py_ssize_t *col_offset, Py_ssize_t *end_col_offset)
 {
     _Py_CRITICAL_SECTION_ASSERT_OBJECT_LOCKED(it);
     Py_ssize_t byte_offset = -1;
-    if (token.start != NULL && token.start >= line_start) {
-        byte_offset = token.start - line_start;
+    if (token_start != NULL && token_start >= line_start) {
+        byte_offset = token_start - line_start;
         if (line_changed) {
             *col_offset = _PyPegen_byte_offset_to_character_offset_line(line, 0, byte_offset);
             it->byte_col_offset_diff = byte_offset - *col_offset;
@@ -220,15 +221,13 @@ _get_col_offsets(tokenizeriterobject *it, struct token token, const char *line_s
         }
     }
 
-    if (token.end != NULL && token.end >= it->tok->line_start) {
-        Py_ssize_t end_byte_offset = token.end - it->tok->line_start;
+    if (token_end != NULL && token_end >= it->tok->line_start) {
+        Py_ssize_t end_byte_offset = token_end - it->tok->line_start;
         if (lineno == end_lineno) {
-            // If the whole token is at the same line, we can just use the token.start
-            // buffer for figuring out the new column offset, since using line is not
-            // performant for very long lines.
+            // Avoid rescanning the prefix of a very long line.
             Py_ssize_t token_col_offset = _PyPegen_byte_offset_to_character_offset_line(line, byte_offset, end_byte_offset);
             *end_col_offset = *col_offset + token_col_offset;
-            it->byte_col_offset_diff += token.end - token.start - token_col_offset;
+            it->byte_col_offset_diff += token_end - token_start - token_col_offset;
         }
         else {
             *end_col_offset = _PyPegen_byte_offset_to_character_offset_raw(it->tok->line_start, end_byte_offset);
@@ -263,12 +262,18 @@ tokenizeriter_next(PyObject *op)
         it->done = 1;
         goto exit;
     }
-    PyObject *str = NULL;
-    if (token.start == NULL || token.end == NULL) {
+    const char *token_start = NULL;
+    const char *token_end = NULL;
+    PyObject *str;
+    if (!_PyTok_SpanIsValid(token.span)) {
         str = Py_GetConstant(Py_CONSTANT_EMPTY_STR);
     }
     else {
-        str = PyUnicode_FromStringAndSize(token.start, token.end - token.start);
+        Py_ssize_t token_length;
+        token_start = _PyToken_TextView(
+            it->tok, &token, &token_length);
+        token_end = token_start + token_length;
+        str = PyUnicode_FromStringAndSize(token_start, token_length);
     }
     if (str == NULL) {
         goto exit;
@@ -297,11 +302,11 @@ tokenizeriter_next(PyObject *op)
         goto exit;
     }
 
-    Py_ssize_t lineno = ISSTRINGLIT(type) ? it->tok->first_lineno : it->tok->lineno;
-    Py_ssize_t end_lineno = it->tok->lineno;
+    Py_ssize_t lineno = token.start_loc.lineno;
+    Py_ssize_t end_lineno = token.end_loc.lineno;
     Py_ssize_t col_offset = -1;
     Py_ssize_t end_col_offset = -1;
-    _get_col_offsets(it, token, line_start, line, line_changed,
+    _get_col_offsets(it, token_start, token_end, line_start, line, line_changed,
                      lineno, end_lineno, &col_offset, &end_col_offset);
 
     if (it->tok->tok_extra_tokens) {
@@ -317,7 +322,8 @@ tokenizeriter_next(PyObject *op)
         else if (type == NEWLINE) {
             Py_DECREF(str);
             if (!it->tok->implicit_newline) {
-                if (it->tok->start[0] == '\r') {
+                assert(token_start != NULL);
+                if (token_start[0] == '\r') {
                     str = PyUnicode_FromString("\r\n");
                 } else {
                     str = PyUnicode_FromString("\n");