gh-153569: Keep tokenizer source scans and copies on offsets and spans

Author: pablogsal

Parser/lexer/lexer.c

@@ -13,12 +13,6 @@
    tokenizing. */
 static const char* type_comment_prefix = "# type: ";
 
-static inline int
-contains_null_bytes(const char* str, size_t size)
-{
-    return memchr(str, 0, size) != NULL;
-}
-
 int
 _PyLexer_refill(struct tok_state *tok)
 {
@@ -40,8 +34,8 @@ _PyLexer_refill(struct tok_state *tok)
         return 0;
     }
     tok->line_start = tok->cur;
-    if (contains_null_bytes(_PyTok_SourcePointer(&tok->source, tok->line_start),
-                            tok->inp - tok->line_start)) {
+    _PyTok_Span line = {tok->line_start, tok->inp};
+    if (_PyTok_SourceFindByte(&tok->source, line, '\0') >= 0) {
         _PyTokenizer_syntaxerror(tok, "source code cannot contain null bytes");
         tok->cur = tok->inp;
         return 0;
@@ -57,8 +51,7 @@ _PyLexer_backup(struct tok_state *tok, int c)
         if (--tok->cur < tok->buf_offset) {
             Py_FatalError("tokenizer beginning of buffer");
         }
-        const char *cur = _PyTok_SourcePointer(&tok->source, tok->cur);
-        if ((int)(unsigned char)*cur != Py_CHARMASK(c)) {
+        if (_PyTok_SourceByte(&tok->source, tok->cur) != Py_CHARMASK(c)) {
             Py_FatalError("tok_backup: wrong character");
         }
     }
@@ -175,10 +168,6 @@ _PyLexer_get_normal(struct tok_state *tok, ftstring_state *current, struct token
     /* Skip comment, unless it's a type comment */
     if (c == '#') {
 
-        const char* p = NULL;
-        const char *prefix, *type_start;
-        int current_starting_col_offset;
-
         while (c != EOF && c != '\n' && c != '\r') {
             c = tok_nextc(tok);
         }
@@ -195,23 +184,20 @@ _PyLexer_get_normal(struct tok_state *tok, ftstring_state *current, struct token
             }
         }
 
-        if (tok->tok_extra_tokens) {
-            p = _PyTok_SourcePointer(&tok->source, tok->start);
-        }
-
+        _PyTok_Off comment = tok->start;
         if (tok->type_comments) {
-            p = _PyTok_SourcePointer(&tok->source, tok->start);
-            current_starting_col_offset = tok->start_loc.byte_col;
-            prefix = type_comment_prefix;
-            while (*prefix && p < _PyTok_SourcePointer(&tok->source, tok->cur)) {
+            const char *prefix = type_comment_prefix;
+            while (*prefix && comment < tok->cur) {
                 if (*prefix == ' ') {
-                    while (*p == ' ' || *p == '\t') {
-                        p++;
-                        current_starting_col_offset++;
+                    while (comment < tok->cur) {
+                        int ch = _PyTok_SourceByte(&tok->source, comment);
+                        if (ch != ' ' && ch != '\t') {
+                            break;
+                        }
+                        comment++;
                     }
-                } else if (*prefix == *p) {
-                    p++;
-                    current_starting_col_offset++;
+                } else if (*prefix == _PyTok_SourceByte(&tok->source, comment)) {
+                    comment++;
                 } else {
                     break;
                 }
@@ -221,36 +207,26 @@ _PyLexer_get_normal(struct tok_state *tok, ftstring_state *current, struct token
 
             /* This is a type comment if we matched all of type_comment_prefix. */
             if (!*prefix) {
-                int is_type_ignore = 1;
-                // +6 in order to skip the word 'ignore'
-                const char *ignore_end = p + 6;
-                const int ignore_end_col_offset = current_starting_col_offset + 6;
                 tok_backup(tok, c);  /* don't eat the newline or EOF */
-
-                type_start = p;
-
                 /* A TYPE_IGNORE is "type: ignore" followed by the end of the token
                  * or anything ASCII and non-alphanumeric. */
-                is_type_ignore = (
-                    _PyTok_SourcePointer(&tok->source, tok->cur) >= ignore_end
-                    && memcmp(p, "ignore", 6) == 0
-                    && !(_PyTok_SourcePointer(&tok->source, tok->cur) > ignore_end
-                         && ((unsigned char)ignore_end[0] >= 128 || Py_ISALNUM(ignore_end[0]))));
+                int is_type_ignore = tok->cur - comment >= 6 &&
+                    memcmp(_PyTok_SourcePointer(&tok->source, comment),
+                           "ignore", 6) == 0;
+                if (is_type_ignore && comment + 6 < tok->cur) {
+                    int ch = _PyTok_SourceByte(&tok->source, comment + 6);
+                    is_type_ignore = ch < 128 && !Py_ISALNUM(ch);
+                }
 
                 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_start = comment + (is_type_ignore ? 6 : 0);
                 p_end = tok->cur;
-                if (is_type_ignore) {
-                    p_start = _PyTok_SourceOffset(&tok->source, ignore_end);
-
-                    /* If this type ignore is the only thing on the line, consume the newline also. */
-                    if (blankline) {
-                        tok_nextc(tok);
-                        tok->layout.at_bol = 1;
-                    }
-                } else {
-                    p_start = _PyTok_SourceOffset(&tok->source, type_start);
+                int start_col_offset = tok->start_loc.byte_col +
+                    (int)(p_start - tok->start);
+                /* Consume the newline after a standalone type ignore. */
+                if (is_type_ignore && blankline) {
+                    tok_nextc(tok);
+                    tok->layout.at_bol = 1;
                 }
                 _PyLexer_token_setup(tok, token, type, p_start, p_end);
                 token->start_loc = (_PyTok_Loc){tok->lineno, start_col_offset};
@@ -261,7 +237,7 @@ _PyLexer_get_normal(struct tok_state *tok, ftstring_state *current, struct token
         }
         if (tok->tok_extra_tokens) {
             tok_backup(tok, c);  /* don't eat the newline or EOF */
-            p_start = _PyTok_SourceOffset(&tok->source, p);
+            p_start = comment;
             p_end = tok->cur;
             tok->layout.comment_newline = blankline;
             return MAKE_TOKEN(COMMENT);

Parser/lexer/string.c

@@ -76,13 +76,10 @@ finish_ftstring_expr(struct tok_state *tok, ftstring_state *state,
     if (!(state->debug_expr || tstring_interpolation) || token->metadata) {
         return 0;
     }
-    Py_ssize_t expr_len;
-    const char *expr = _PyTok_SourceSpanView(
-        &tok->source, state->expr_span, &expr_len);
     tokenizer_comments *comments = state->comments;
     PyObject *res;
     if (comments != NULL && comments->count > 0) {
-        Py_ssize_t stripped_size = expr_len;
+        Py_ssize_t stripped_size = state->expr_span.end - state->expr_span.start;
         Py_ssize_t comment_count = 0;
         for (Py_ssize_t i = 0; i < comments->count; i++) {
             _PyTok_Span comment = comments->spans[i];
@@ -103,24 +100,26 @@ finish_ftstring_expr(struct tok_state *tok, ftstring_state *state,
         }
         _PyTok_Off copied_to = state->expr_span.start;
         Py_ssize_t stripped_len = 0;
-        for (Py_ssize_t i = 0; i < comment_count; i++) {
-            _PyTok_Span comment = comments->spans[i];
-            Py_ssize_t length = comment.start - copied_to;
-            memcpy(stripped + stripped_len,
-                   expr + copied_to - state->expr_span.start,
-                   (size_t)length);
+        for (Py_ssize_t i = 0; i <= comment_count; i++) {
+            _PyTok_Span span = {
+                copied_to,
+                i < comment_count ? comments->spans[i].start : state->expr_span.end,
+            };
+            Py_ssize_t length;
+            const char *text = _PyTok_SourceSpanView(&tok->source, span, &length);
+            memcpy(stripped + stripped_len, text, (size_t)length);
             stripped_len += length;
-            copied_to = comment.end;
+            if (i < comment_count) {
+                copied_to = comments->spans[i].end;
+            }
         }
-        Py_ssize_t length = state->expr_span.end - copied_to;
-        memcpy(stripped + stripped_len,
-               expr + copied_to - state->expr_span.start,
-               (size_t)length);
-        stripped_len += length;
         res = PyUnicode_DecodeUTF8(stripped, stripped_len, NULL);
         PyMem_Free(stripped);
     }
     else {
+        Py_ssize_t expr_len;
+        const char *expr = _PyTok_SourceSpanView(
+            &tok->source, state->expr_span, &expr_len);
         res = PyUnicode_DecodeUTF8(expr, expr_len, NULL);
     }
 

Parser/tokenizer/api.c

@@ -107,22 +107,28 @@ _PyTokenizer_ImplyDedents(struct tok_state *tok)
 int
 _PyTokenizer_HasTrailingStatement(const struct tok_state *tok)
 {
-    const char *cur = _PyTok_SourcePointer(&tok->source, tok->cur);
-    char c = *cur;
-    for (;;) {
-        while (c == ' ' || c == '\t' || c == '\n' || c == '\014') {
-            c = *++cur;
-        }
-        if (!c) {
+    _PyTok_Off cur = tok->cur;
+    _PyTok_Off end = tok->source.base_offset + tok->source.len;
+    while (cur < end) {
+        int c = _PyTok_SourceByte(&tok->source, cur++);
+        if (c == '\0') {
             return 0;
         }
+        if (c == ' ' || c == '\t' || c == '\n' || c == '\014') {
+            continue;
+        }
         if (c != '#') {
             return 1;
         }
-        while (c && c != '\n') {
-            c = *++cur;
+        while (cur < end) {
+            c = _PyTok_SourceByte(&tok->source, cur);
+            if (c == '\0' || c == '\n') {
+                break;
+            }
+            cur++;
         }
     }
+    return 0;
 }
 
 int

Parser/tokenizer/reader.c

@@ -175,13 +175,11 @@ next_prepared(struct tok_state *tok, _PyTok_Chunk *chunk)
         return _PYTOK_READ_EOF;
     }
     _PyTok_Span tail = {tok->inp, tok->source.base_offset + tok->source.len};
-    Py_ssize_t remaining;
-    const char *start = _PyTok_SourceSpanView(&tok->source, tail, &remaining);
-    const char *newline = memchr(start, '\n', remaining);
-    chunk->data = (char *)start;
-    chunk->len = newline != NULL ? newline - start + 1 : remaining;
+    _PyTok_Off newline = _PyTok_SourceFindByte(&tok->source, tail, '\n');
+    _PyTok_Span line = {tail.start, newline >= 0 ? newline + 1 : tail.end};
+    chunk->data = (char *)_PyTok_SourceSpanView(&tok->source, line, &chunk->len);
     chunk->ownership = _PYTOK_CHUNK_BORROWED;
-    chunk->implicit_newline = chunk->len == remaining &&
+    chunk->implicit_newline = line.end == tail.end &&
         tok->reader->prepared_final_newline_is_implicit;
     return _PYTOK_READ_LINE;
 }
@@ -667,11 +665,10 @@ _PyTok_ReaderUnderflow(struct tok_state *tok)
         tok->inp = source_start + chunk.len;
     }
     else {
-        _PyTok_Off source_start = _PyTok_SourceOffset(&tok->source, chunk.data);
         if (tok->start < 0 && _PyLexer_CurrentFTString(tok) == NULL) {
-            tok->buf_offset = source_start;
+            tok->buf_offset = tok->inp;
         }
-        tok->inp = source_start + chunk.len;
+        tok->inp += chunk.len;
     }
     tok->implicit_newline = chunk.implicit_newline;
 

Parser/tokenizer/source.h

@@ -21,14 +21,6 @@ _PyTok_SourceData(const _PyTok_SourceText *source)
 
 /* Convert positions within the retained source window. Pointers and views
    are borrowed; append, discard, and clear invalidate them. */
-static inline _PyTok_Off
-_PyTok_SourceOffset(const _PyTok_SourceText *source, const char *position)
-{
-    const char *base = _PyTok_SourceData(source);
-    assert(position >= base && position <= base + source->len);
-    return source->base_offset + (position - base);
-}
-
 static inline const char *
 _PyTok_SourcePointer(const _PyTok_SourceText *source, _PyTok_Off offset)
 {
@@ -56,6 +48,16 @@ _PyTok_SourceSpanView(const _PyTok_SourceText *source, _PyTok_Span span,
     return _PyTok_SourcePointer(source, span.start);
 }
 
+/* Return the first matching offset within span, or -1 if absent. */
+static inline _PyTok_Off
+_PyTok_SourceFindByte(const _PyTok_SourceText *source, _PyTok_Span span, int byte)
+{
+    Py_ssize_t length;
+    const char *data = _PyTok_SourceSpanView(source, span, &length);
+    const char *found = memchr(data, byte, length);
+    return found != NULL ? span.start + (found - data) : -1;
+}
+
 PyAPI_FUNC(void) _PyTok_SourceInit(_PyTok_SourceText *);
 /* Clear invalidates all spans and views for the source. */
 PyAPI_FUNC(void) _PyTok_SourceClear(_PyTok_SourceText *);