gh-153569: return tokenizer tokens as source spans

Lib/test/test_tokenize.py

@@ -2427,6 +2427,31 @@ def test_stop_iteration_skips_encoded_readline_codec_lookup(self):
             (token.ENDMARKER, "", (1, 0), (1, 0), ""),
         )
 
+    def test_fstring_offsets_survive_buffer_reallocation(self):
+        padding = " " * 9000
+        expression_line = ")=:>{2}}\n"
+        physical_lines = [
+            'f"""\n',
+            "{(\n",
+            padding + "1\n",
+            expression_line,
+            '"""\n',
+        ]
+        source = "".join(physical_lines)
+        chunks = iter([
+            "".join(physical_lines[:2]),
+            "".join(physical_lines[2:4]),
+            physical_lines[4],
+            "",
+        ])
+
+        expected = self._get_tokens(source, extra_tokens=True)
+        tokens = list(tokenize._generate_tokens_from_c_tokenizer(
+            chunks.__next__,
+            extra_tokens=True,
+        ))
+        self.assertEqual(tokens, expected)
+
     def test_extra_tokens_relaxes_lexer_errors(self):
         cases = [
             (

Parser/lexer/buffer.c

@@ -1,62 +1,45 @@
 #include "Python.h"
-#include "errcode.h"
-
+#include "buffer.h"
 #include "state.h"
 
-/* Traverse and remember all f-string buffers, in order to be able to restore
-   them after reallocating tok->buf */
 void
-_PyLexer_remember_fstring_buffers(struct tok_state *tok)
+_PyLexer_SnapshotBuffer(struct tok_state *tok, const char *base,
+                        _PyLexer_BufferSnapshot *snapshot)
 {
-    int index;
-    tokenizer_mode *mode;
-
-    for (index = tok->tok_mode_stack_index; index >= 0; --index) {
-        mode = &(tok->tok_mode_stack[index]);
+    snapshot->buf = tok->buf - base;
+    snapshot->cur = tok->cur - tok->buf;
+    snapshot->inp = tok->inp - tok->buf;
+    snapshot->start = tok->start == NULL ? -1 : tok->start - tok->buf;
+    snapshot->line_start = tok->line_start == NULL
+        ? -1 : tok->line_start - tok->buf;
+    snapshot->multi_line_start = tok->multi_line_start == NULL
+        ? -1 : tok->multi_line_start - tok->buf;
+    for (int index = tok->tok_mode_stack_index; index >= 0; --index) {
+        tokenizer_mode *mode = &tok->tok_mode_stack[index];
         mode->start_offset = mode->start == NULL ? -1 : mode->start - tok->buf;
-        mode->multi_line_start_offset = mode->multi_line_start == NULL ? -1 : mode->multi_line_start - tok->buf;
+        mode->multi_line_start_offset = mode->multi_line_start == NULL
+            ? -1 : mode->multi_line_start - tok->buf;
     }
 }
 
-/* Traverse and restore all f-string buffers after reallocating tok->buf */
 void
-_PyLexer_restore_fstring_buffers(struct tok_state *tok)
-{
-    int index;
-    tokenizer_mode *mode;
-
-    for (index = tok->tok_mode_stack_index; index >= 0; --index) {
-        mode = &(tok->tok_mode_stack[index]);
-        mode->start = mode->start_offset < 0 ? NULL : tok->buf + mode->start_offset;
-        mode->multi_line_start = mode->multi_line_start_offset < 0 ? NULL : tok->buf + mode->multi_line_start_offset;
-    }
-}
-
-int
-_PyLexer_tok_reserve_buf(struct tok_state *tok, Py_ssize_t size)
+_PyLexer_RestoreBuffer(struct tok_state *tok, char *base,
+                       const _PyLexer_BufferSnapshot *snapshot)
 {
-    Py_ssize_t cur = tok->cur - tok->buf;
-    Py_ssize_t oldsize = tok->inp - tok->buf;
-    Py_ssize_t newsize = oldsize + Py_MAX(size, oldsize >> 1);
-    if (newsize > tok->end - tok->buf) {
-        char *newbuf = tok->buf;
-        Py_ssize_t start = tok->start == NULL ? -1 : tok->start - tok->buf;
-        Py_ssize_t line_start = tok->start == NULL ? -1 : tok->line_start - tok->buf;
-        Py_ssize_t multi_line_start = tok->multi_line_start - tok->buf;
-        _PyLexer_remember_fstring_buffers(tok);
-        newbuf = (char *)PyMem_Realloc(newbuf, newsize);
-        if (newbuf == NULL) {
-            tok->done = E_NOMEM;
-            return 0;
-        }
-        tok->buf = newbuf;
-        tok->cur = tok->buf + cur;
-        tok->inp = tok->buf + oldsize;
-        tok->end = tok->buf + newsize;
-        tok->start = start < 0 ? NULL : tok->buf + start;
-        tok->line_start = line_start < 0 ? NULL : tok->buf + line_start;
-        tok->multi_line_start = multi_line_start < 0 ? NULL : tok->buf + multi_line_start;
-        _PyLexer_restore_fstring_buffers(tok);
+    tok->buf = base + snapshot->buf;
+    tok->cur = tok->buf + snapshot->cur;
+    tok->inp = tok->buf + snapshot->inp;
+    tok->start = snapshot->start < 0
+        ? NULL : tok->buf + snapshot->start;
+    tok->line_start = snapshot->line_start < 0
+        ? NULL : tok->buf + snapshot->line_start;
+    tok->multi_line_start = snapshot->multi_line_start < 0
+        ? NULL : tok->buf + snapshot->multi_line_start;
+    for (int index = tok->tok_mode_stack_index; index >= 0; --index) {
+        tokenizer_mode *mode = &tok->tok_mode_stack[index];
+        mode->start = mode->start_offset < 0
+            ? NULL : tok->buf + mode->start_offset;
+        mode->multi_line_start = mode->multi_line_start_offset < 0
+            ? NULL : tok->buf + mode->multi_line_start_offset;
     }
-    return 1;
 }

Parser/lexer/buffer.h

@@ -3,8 +3,20 @@
 
 #include "pyport.h"
 
-void _PyLexer_remember_fstring_buffers(struct tok_state *tok);
-void _PyLexer_restore_fstring_buffers(struct tok_state *tok);
-int _PyLexer_tok_reserve_buf(struct tok_state *tok, Py_ssize_t size);
+struct tok_state;
+
+typedef struct {
+    Py_ssize_t buf;
+    Py_ssize_t cur;
+    Py_ssize_t inp;
+    Py_ssize_t start;
+    Py_ssize_t line_start;
+    Py_ssize_t multi_line_start;
+} _PyLexer_BufferSnapshot;
+
+void _PyLexer_SnapshotBuffer(
+    struct tok_state *, const char *, _PyLexer_BufferSnapshot *);
+void _PyLexer_RestoreBuffer(
+    struct tok_state *, char *, const _PyLexer_BufferSnapshot *);
 
 #endif

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

@@ -26,7 +26,6 @@ _PyTokenizer_tok_new(void)
     tok->interactive_src_start = NULL;
     tok->interactive_src_end = NULL;
     tok->start = NULL;
-    tok->end = NULL;
     tok->done = E_OK;
     tok->fp = NULL;
     tok->tabsize = TABSIZE;
@@ -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

@@ -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;
 };
 
@@ -67,15 +68,15 @@ typedef struct _tokenizer_mode {
 
 /* Tokenizer state */
 struct tok_state {
-    /* Input state; buf <= cur <= inp <= end */
+    /* 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]. */
     int fp_interactive; /* If the file descriptor is interactive */
     char *interactive_src_start; /* The start of the source parsed so far in interactive mode */
     char *interactive_src_end; /* The end of the source parsed so far in interactive mode */
-    const char *end;    /* End of input buffer if buf != NULL */
     const char *start;  /* Start of current token if not NULL */
     int done;           /* E_OK normally, E_EOF at EOF, otherwise error code */
     /* NB If done != E_OK, cur must be == inp!!! */
@@ -128,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/lexer/string.c

@@ -125,7 +125,8 @@ _PyLexer_update_ftstring_expr(struct tok_state *tok, char cur)
 {
     assert(tok->cur != NULL);
 
-    Py_ssize_t size = strlen(tok->cur);
+    Py_ssize_t size = cur == 0
+        ? tok->inp - tok->cur : (Py_ssize_t)strlen(tok->cur);
     tokenizer_mode *tok_mode = TOK_GET_MODE(tok);
 
     switch (cur) {
@@ -142,7 +143,8 @@ _PyLexer_update_ftstring_expr(struct tok_state *tok, char cur)
                 goto error;
             }
             tok_mode->last_expr_buffer = new_buffer;
-            strncpy(tok_mode->last_expr_buffer + tok_mode->last_expr_size, tok->cur, size);
+            memcpy(tok_mode->last_expr_buffer + tok_mode->last_expr_size,
+                   tok->cur, size);
             tok_mode->last_expr_size += size;
             break;
         case '{':
@@ -155,7 +157,7 @@ _PyLexer_update_ftstring_expr(struct tok_state *tok, char cur)
             }
             tok_mode->last_expr_size = size;
             tok_mode->last_expr_end = -1;
-            strncpy(tok_mode->last_expr_buffer, tok->cur, size);
+            memcpy(tok_mode->last_expr_buffer, tok->cur, size);
             break;
         case '}':
         case '!':