gh-153569: return token spans and unify decoded source storage

Lib/test/test_capi/test_tokenizer.py

@@ -9,6 +9,9 @@ class TokenizerTests(unittest.TestCase):
     def test_source(self):
         _testinternalcapi.test_tokenizer_source()
 
+    def test_source_discard(self):
+        _testinternalcapi.test_tokenizer_source_discard()
+
     def test_cursor(self):
         _testinternalcapi.test_tokenizer_cursor()
 

Lib/test/test_repl.py

@@ -198,6 +198,22 @@ def test_lexer_buffer_realloc_with_null_start(self):
         self.assertEqual(p.returncode, 0)
         self.assertIn(long_value, output)
 
+    @cpython_only
+    def test_multiline_fstring_source_reallocation(self):
+        long_line = " " * 9000 + "+ 2"
+        user_input = (
+            'value = f"""{(\n'
+            '1\n'
+            f'{long_line}\n'
+            ')}"""\n'
+            'print(value)\n'
+        )
+        p = spawn_repl()
+        p.stdin.write(user_input)
+        output = kill_python(p)
+        self.assertEqual(p.returncode, 0)
+        self.assertIn(">>> 3\n>>> ", output)
+
     def test_close_stdin(self):
         user_input = dedent('''
             import os

Lib/test/test_tokenize.py

@@ -2427,6 +2427,30 @@ def test_stop_iteration_skips_encoded_readline_codec_lookup(self):
             (token.ENDMARKER, "", (1, 0), (1, 0), ""),
         )
 
+    def test_fstring_offsets_survive_buffer_reallocation(self):
+        for prefix in ("f", "t"):
+            for extra_tokens in (False, True):
+                with self.subTest(prefix=prefix, extra_tokens=extra_tokens):
+                    physical_lines = [
+                        prefix + '"""\n',
+                        "{(\n",
+                        " " * 9000 + "1\n",
+                        ")=:>{2}}\n",
+                        '"""\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=extra_tokens)
+                    tokens = list(tokenize._generate_tokens_from_c_tokenizer(
+                        chunks.__next__, extra_tokens=extra_tokens))
+                    self.assertEqual(tokens, expected)
+
     def test_extra_tokens_relaxes_lexer_errors(self):
         cases = [
             (

Modules/_testinternalcapi/tokenizer.c

@@ -195,6 +195,28 @@ test_tokenizer_source(PyObject *Py_UNUSED(module),
         goto error;
     }
 
+    _PyTok_SourceDiscard(&source);
+    if (check(_PyTok_SourceAppendLine(&source, "a\n", 2, 0) == 4,
+              "wrong retained source offset") < 0 ||
+            _PyTok_SourceLine(&source, 1, &line) < 0 ||
+            check(line.start == 4 && line.end == 6,
+                  "wrong retained source line") < 0 ||
+            _PyTok_SourceLocation(
+                &source, 4, _PYTOK_AFFINITY_LEFT, &loc) < 0 ||
+            check(loc.lineno == 1 && loc.byte_col == 0,
+                  "wrong retained source location") < 0) {
+        goto error;
+    }
+    view = _PyTok_SourceSpanView(
+        &source, _PyTok_SpanFromBounds(4, 5), &view_len);
+    if (check(view != NULL && view_len == 1 && view[0] == 'a',
+              "wrong retained source span") < 0 ||
+            check_system_error(_PyTok_SourceSpanView(
+                &source, _PyTok_SpanFromBounds(0, 1), &view_len) == NULL,
+                "accepted discarded source span") < 0) {
+        goto error;
+    }
+
     _PyTok_SourceClear(&source);
     Py_RETURN_NONE;
 
@@ -296,6 +318,88 @@ test_tokenizer_cursor(PyObject *Py_UNUSED(module),
     }
 #endif
 
+    _PyTok_Off base = source.len;
+    _PyTok_SourceDiscard(&source);
+    if (_PyTok_SourceAppendLine(&source, "ab\n", 3, 0) < 0 ||
+            _PyTok_SourceAppendLine(&source, "cd", 2, 0) < 0) {
+        goto error;
+    }
+    _PyTok_CursorInit(&cursor, &source);
+    if (_PyTok_CursorSetLine(&cursor, 1) < 0 ||
+            check(cursor.pos == base && _PyTok_CursorPeek(&cursor, 1) == 'b',
+                  "wrong retained cursor line") < 0 ||
+            _PyTok_CursorSetLine(&cursor, 2) < 0 ||
+            check(_PyTok_CursorAdvance(&cursor) == 'c',
+                  "wrong retained cursor byte") < 0 ||
+            _PyTok_CursorSetOffset(&cursor, base + 5) < 0 ||
+            check(cursor.lineno == 2 && _PyTok_CursorAdvance(&cursor) == EOF,
+                  "wrong retained cursor EOF") < 0) {
+        goto error;
+    }
+
+    _PyTok_SourceClear(&source);
+    Py_RETURN_NONE;
+
+error:
+    _PyTok_SourceClear(&source);
+    return NULL;
+}
+
+static PyObject *
+test_tokenizer_source_discard(PyObject *Py_UNUSED(module),
+                             PyObject *Py_UNUSED(args))
+{
+    _PyTok_SourceText source;
+    _PyTok_SourceInit(&source);
+    for (int i = 0; i < 260; i++) {
+        if (_PyTok_SourceAppendLine(&source, "x\n", 2, 1) < 0) {
+            goto error;
+        }
+    }
+    char *bytes = source.bytes;
+    _PyTok_Off capacity = source.cap;
+    _PyTok_SourceDiscard(&source);
+    if (check(source.base_offset == 520 && source.len == 0 &&
+                  source.nlines == 0 && source.bytes == bytes &&
+                  source.cap == capacity && source.bytes[0] == '\0',
+              "discard did not preserve source allocation") < 0) {
+        goto error;
+    }
+    for (int i = 0; i < 260; i++) {
+        if (check(_PyTok_SourceAppendLine(&source, "y\n", 2, 0) == 520 + 2 * i,
+                  "wrong source offset after discard") < 0 ||
+                check(!_PyTok_SourceLineIsImplicit(&source, i + 1),
+                      "discard preserved implicit newline flag") < 0) {
+            goto error;
+        }
+    }
+    if (check(source.bytes == bytes && source.cap == capacity,
+              "discarded allocation was not reused") < 0) {
+        goto error;
+    }
+    _PyTok_SourceDiscard(&source);
+    if (check(_PyTok_SourceAppendLine(&source, "tail", 4, 0) == 1040,
+              "wrong source offset after repeated discard") < 0) {
+        goto error;
+    }
+    _PyTok_SourceDiscard(&source);
+    if (check(_PyTok_SourceAppendLine(&source, "z\n", 2, 0) == 1044,
+              "cannot append after discarding unterminated line") < 0) {
+        goto error;
+    }
+    _PyTok_SourceDiscard(&source);
+    source.base_offset = PY_SSIZE_T_MAX - 1;
+    if (check(_PyTok_SourceAppendLine(&source, "z\n", 2, 0) < 0 &&
+                  PyErr_ExceptionMatches(PyExc_MemoryError),
+              "accepted overflowing logical source offset") < 0) {
+        goto error;
+    }
+    PyErr_Clear();
+    if (check(source.len == 0 && source.nlines == 0 &&
+                  source.base_offset == PY_SSIZE_T_MAX - 1,
+              "overflow changed retained source") < 0) {
+        goto error;
+    }
     _PyTok_SourceClear(&source);
     Py_RETURN_NONE;
 
@@ -307,6 +411,7 @@ test_tokenizer_cursor(PyObject *Py_UNUSED(module),
 static PyMethodDef test_methods[] = {
     {"test_tokenizer_source", test_tokenizer_source, METH_NOARGS},
     {"test_tokenizer_cursor", test_tokenizer_cursor, METH_NOARGS},
+    {"test_tokenizer_source_discard", test_tokenizer_source_discard, METH_NOARGS},
     {NULL},
 };
 

Parser/lexer/buffer.c

@@ -1,62 +1,46 @@
 #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_SaveBufferPointers(struct tok_state *tok, const char *base,
+                            _PyLexer_BufferPointers *pointers)
 {
-    int index;
-    tokenizer_mode *mode;
-
-    for (index = tok->tok_mode_stack_index; index >= 0; --index) {
-        mode = &(tok->tok_mode_stack[index]);
+    pointers->buf_from_base = tok->buf - base;
+    pointers->cur_from_buf = tok->cur - tok->buf;
+    pointers->inp_from_buf = tok->inp - tok->buf;
+    pointers->start_from_buf = tok->start == NULL
+        ? -1 : tok->start - tok->buf;
+    pointers->line_start_from_buf = tok->line_start == NULL
+        ? -1 : tok->line_start - tok->buf;
+    pointers->multi_line_start_from_buf = 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_RestoreBufferPointers(struct tok_state *tok, char *base,
+                               const _PyLexer_BufferPointers *pointers)
 {
-    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 + pointers->buf_from_base;
+    tok->cur = tok->buf + pointers->cur_from_buf;
+    tok->inp = tok->buf + pointers->inp_from_buf;
+    tok->start = pointers->start_from_buf < 0
+        ? NULL : tok->buf + pointers->start_from_buf;
+    tok->line_start = pointers->line_start_from_buf < 0
+        ? NULL : tok->buf + pointers->line_start_from_buf;
+    tok->multi_line_start = pointers->multi_line_start_from_buf < 0
+        ? NULL : tok->buf + pointers->multi_line_start_from_buf;
+    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_from_base;
+    Py_ssize_t cur_from_buf;
+    Py_ssize_t inp_from_buf;
+    Py_ssize_t start_from_buf;
+    Py_ssize_t line_start_from_buf;
+    Py_ssize_t multi_line_start_from_buf;
+} _PyLexer_BufferPointers;
+
+void _PyLexer_SaveBufferPointers(
+    struct tok_state *, const char *, _PyLexer_BufferPointers *);
+void _PyLexer_RestoreBufferPointers(
+    struct tok_state *, char *, const _PyLexer_BufferPointers *);
 
 #endif

Parser/lexer/lexer.c

@@ -12,8 +12,6 @@
 
 
 #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))
 
 /* Spaces in this constant are treated as "zero or more spaces or tabs" when
    tokenizing. */
@@ -360,21 +358,25 @@ _PyLexer_get_normal_mode(struct tok_state *tok, tokenizer_mode* current_tok, str
                     && !(tok->cur > ignore_end
                          && ((unsigned char)ignore_end[0] >= 128 || Py_ISALNUM(ignore_end[0]))));
 
+                int type = is_type_ignore ? TYPE_IGNORE : TYPE_COMMENT;
+                int start_col_offset = is_type_ignore
+                    ? ignore_end_col_offset : current_starting_col_offset;
+                p_end = tok->cur;
                 if (is_type_ignore) {
                     p_start = ignore_end;
-                    p_end = tok->cur;
 
                     /* If this type ignore is the only thing on the line, consume the newline also. */
                     if (blankline) {
                         tok_nextc(tok);
                         tok->atbol = 1;
                     }
-                    return MAKE_TYPE_COMMENT_TOKEN(TYPE_IGNORE, ignore_end_col_offset, tok->col_offset);
                 } else {
                     p_start = type_start;
-                    p_end = tok->cur;
-                    return MAKE_TYPE_COMMENT_TOKEN(TYPE_COMMENT, current_starting_col_offset, tok->col_offset);
                 }
+                _PyLexer_token_setup(tok, token, type, p_start, p_end);
+                token->start_loc = (_PyTok_Loc){tok->lineno, start_col_offset};
+                token->end_loc = (_PyTok_Loc){tok->lineno, tok->col_offset};
+                return type;
             }
         }
         if (tok->tok_extra_tokens) {

Parser/lexer/lexer.h

@@ -7,4 +7,25 @@ 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 (token->span.start < 0) {
+        assert(token->span.start == -1 && token->span.end == -1);
+        *length = 0;
+        return "";
+    }
+    assert(_PyTok_SpanIsValid(token->span));
+    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