gh-159672: Retain tokens for syntax error diagnostics

Author: pablogsal

Parser/pegen.c

@@ -327,6 +327,28 @@ _resize_tokens_array(Parser *p) {
     return 0;
 }
 
+int
+_PyPegen_append_token(Parser *p, struct token *token)
+{
+    int type = token->type;
+    // Single input inserts a final newline and may still need dedents. Apply
+    // the same normalization when the error pass fills the replay cache.
+    if (p->start_rule == Py_single_input && type == ENDMARKER && p->parsing_started) {
+        type = NEWLINE;
+        p->parsing_started = 0;
+        if (!(p->flags & PyPARSE_DONT_IMPLY_DEDENT)) {
+            _PyTokenizer_ImplyDedents(p->tok);
+        }
+    }
+    else {
+        p->parsing_started = 1;
+    }
+    if (p->fill == p->size && _resize_tokens_array(p) < 0) {
+        return -1;
+    }
+    return initialize_token(p, p->tokens[p->fill], token, type);
+}
+
 int
 _PyPegen_fill_token(Parser *p)
 {
@@ -363,26 +385,7 @@ _PyPegen_fill_token(Parser *p)
         p->input_exhausted = 1;
     }
 
-    // If we have reached the end and we are in single input mode we need to insert a newline and reset the parsing
-    if (p->start_rule == Py_single_input && type == ENDMARKER && p->parsing_started) {
-        type = NEWLINE; /* Add an extra newline */
-        p->parsing_started = 0;
-
-        if (!(p->flags & PyPARSE_DONT_IMPLY_DEDENT)) {
-            _PyTokenizer_ImplyDedents(p->tok);
-        }
-    }
-    else {
-        p->parsing_started = 1;
-    }
-
-    // Check if we are at the limit of the token array capacity and resize if needed
-    if ((p->fill == p->size) && (_resize_tokens_array(p) != 0)) {
-        goto error;
-    }
-
-    Token *t = p->tokens[p->fill];
-    int result = initialize_token(p, t, &new_token, type);
+    int result = _PyPegen_append_token(p, &new_token);
     _PyToken_Free(&new_token);
     return result;
 error:

Parser/pegen.h

@@ -29,6 +29,9 @@
 
 #define CURRENT_POS (-5)
 
+// Error-only token replay is bounded independently of the source length.
+#define MAX_KEYWORD_TYPO_TOKENS 4096
+
 typedef struct _memo {
     int type;
     void *node;
@@ -193,6 +196,7 @@ expr_ty _PyPegen_soft_keyword_token(Parser *p);
 expr_ty _PyPegen_fstring_middle_token(Parser* p);
 Token *_PyPegen_get_last_nonnwhitespace_token(Parser *);
 int _PyPegen_fill_token(Parser *p);
+int _PyPegen_append_token(Parser *p, struct token *token);
 // Return -1 on failure, leaving the buffer and capacity unchanged.
 int _PyPegen_grow_loop_buffer(void ***buffer, Py_ssize_t *capacity);
 expr_ty _PyPegen_name_token(Parser *p);

Parser/pegen_errors.c

@@ -154,11 +154,24 @@ _PyPegen_tokenize_full_source_to_check_for_errors(Parser *p) {
     Py_ssize_t current_err_line = current_token->lineno;
 
     int ret = 0;
+    int cache_tokens = 1;
     struct token new_token;
     _PyToken_Init(&new_token);
 
     for (;;) {
         _PyTokenizer_Get(p->tok, &new_token);
+        // Retain the tokens this error check already reads for bounded typo
+        // trials. Never treat a truncated cache as complete input.
+        if (cache_tokens && p->fill < MAX_KEYWORD_TYPO_TOKENS &&
+            (p->fill == 0 || p->tokens[p->fill - 1]->type != ENDMARKER) &&
+            new_token.type != ERRORTOKEN && new_token.type != TYPE_IGNORE &&
+            _PyPegen_append_token(p, &new_token) < 0)
+        {
+            // Caching is optional; its allocation failures must not replace
+            // the syntax error we are already reporting.
+            PyErr_Clear();
+            cache_tokens = 0;
+        }
         switch (new_token.type) {
             case ERRORTOKEN: {
                 if (PyErr_Occurred()) {
@@ -177,6 +190,14 @@ _PyPegen_tokenize_full_source_to_check_for_errors(Parser *p) {
                 break;
             }
             case ENDMARKER:
+                if (cache_tokens && p->fill < MAX_KEYWORD_TYPO_TOKENS &&
+                    p->start_rule == Py_single_input && p->fill > 0 &&
+                    p->tokens[p->fill - 1]->type == NEWLINE)
+                {
+                    // append_token inserted single input's implicit newline;
+                    // retain any pending dedents and the actual end marker.
+                    continue;
+                }
                 break;
             default:
                 continue;