gh-153569: keep f-string comments with their mode

Author: pablogsal

Parser/lexer/state.c

@@ -43,6 +43,8 @@ void
 _PyLexer_PopMode(struct tok_state *tok)
 {
     assert(tok->tok_mode_stack_index > 0);
+    tokenizer_mode *mode = _PyLexer_CurrentMode(tok);
+    PyMem_Free(mode->comments);
     tok->tok_mode_stack_index--;
 }
 
@@ -57,15 +59,12 @@ _PyTokenizer_Free(struct tok_state *tok)
     Py_XDECREF(tok->module);
     _PyTok_ReaderFree(tok);
     _PyTok_SourceClear(&tok->source);
+    for (int i = 0; i <= tok->tok_mode_stack_index; i++) {
+        PyMem_Free(tok->tok_mode_stack[i].comments);
+    }
     if (tok->tok_mode_stack != tok->tok_mode_stack_inline) {
         PyMem_Free(tok->tok_mode_stack);
     }
-    tokenizer_comments *comments = tok->ftstring_comments;
-    while (comments != NULL) {
-        tokenizer_comments *previous = comments->previous;
-        PyMem_Free(comments);
-        comments = previous;
-    }
     PyMem_Free(tok);
 }
 

Parser/lexer/state.h

@@ -35,6 +35,12 @@ typedef enum {
 
 #define MAX_EXPR_NESTING 3
 
+typedef struct _tokenizer_comments {
+    Py_ssize_t count;
+    Py_ssize_t capacity;
+    _PyTok_Span spans[];
+} tokenizer_comments;
+
 typedef struct _tokenizer_mode {
     tokenizer_mode_kind kind;
     ftstring_kind ftstring_kind;
@@ -47,6 +53,7 @@ typedef struct _tokenizer_mode {
     _PyTok_Off start;
     _PyTok_Loc start_loc;
     _PyTok_Span expr_span;
+    tokenizer_comments *comments;
 } tokenizer_mode;
 
 static inline int
@@ -68,14 +75,6 @@ _PyLexer_IsRawString(const tokenizer_mode *mode)
            mode->ftstring_kind == RAW_TSTRING;
 }
 
-typedef struct _tokenizer_comments {
-    struct _tokenizer_comments *previous;
-    Py_ssize_t count;
-    Py_ssize_t capacity;
-    int mode_index;
-    _PyTok_Span spans[];
-} tokenizer_comments;
-
 /* Tokenizer state */
 struct tok_state {
     /* Input state; buf <= cur <= inp */
@@ -115,7 +114,6 @@ struct tok_state {
     tokenizer_mode tok_mode_stack_inline[TOKENIZER_MODE_INLINE_CAPACITY];
     int tok_mode_stack_index;
     int tok_mode_stack_capacity;
-    tokenizer_comments *ftstring_comments;
     int tok_extra_tokens;
     int comment_newline;
     int implicit_newline;
@@ -139,12 +137,6 @@ _PyLexer_InsideFString(const struct tok_state *tok)
     return tok->tok_mode_stack_index > 0;
 }
 
-static inline int
-_PyLexer_ModeDepth(const struct tok_state *tok)
-{
-    return tok->tok_mode_stack_index;
-}
-
 static inline char
 _PyLexer_CurrentStringPrefix(const struct tok_state *tok)
 {

Parser/lexer/string.c

@@ -16,15 +16,6 @@ rewind_to_string_start(struct tok_state *tok, const char *start,
     tok->lineno = location.lineno;
 }
 
-static tokenizer_comments *
-current_comments(const struct tok_state *tok)
-{
-    tokenizer_comments *comments = tok->ftstring_comments;
-    return comments != NULL &&
-            comments->mode_index == _PyLexer_ModeDepth(tok)
-        ? comments : NULL;
-}
-
 int
 _PyLexer_record_ftstring_comment(struct tok_state *tok, const char *start,
                                  const char *end)
@@ -34,9 +25,8 @@ _PyLexer_record_ftstring_comment(struct tok_state *tok, const char *start,
         return 0;
     }
     assert(mode->expr_span.start >= 0);
-    tokenizer_comments *comments = current_comments(tok);
+    tokenizer_comments *comments = mode->comments;
     if (comments == NULL || comments->count == comments->capacity) {
-        int create = comments == NULL;
         Py_ssize_t max_capacity = (PY_SSIZE_T_MAX -
             (Py_ssize_t)sizeof(*comments)) /
             (Py_ssize_t)sizeof(*comments->spans);
@@ -53,13 +43,11 @@ _PyLexer_record_ftstring_comment(struct tok_state *tok, const char *start,
             return -1;
         }
         comments = resized;
-        if (create) {
-            comments->previous = tok->ftstring_comments;
+        if (mode->comments == NULL) {
             comments->count = 0;
-            comments->mode_index = _PyLexer_ModeDepth(tok);
         }
         comments->capacity = capacity;
-        tok->ftstring_comments = comments;
+        mode->comments = comments;
     }
     comments->spans[comments->count++] =
         _PyLexer_BufferSpan(tok, start, end);
@@ -79,7 +67,7 @@ _PyLexer_set_ftstring_expr_metadata(struct tok_state *tok, struct token *token)
     Py_ssize_t expr_len;
     const char *expr = _PyLexer_BufferSpanView(
         tok, tok_mode->expr_span, &expr_len);
-    tokenizer_comments *comments = current_comments(tok);
+    tokenizer_comments *comments = tok_mode->comments;
     PyObject *res;
     if (comments != NULL && comments->count > 0) {
         Py_ssize_t stripped_size = expr_len;
@@ -130,9 +118,10 @@ _PyLexer_set_ftstring_expr_metadata(struct tok_state *tok, struct token *token)
 void
 _PyLexer_begin_ftstring_expr(struct tok_state *tok)
 {
-    _PyLexer_CurrentMode(tok)->expr_span = (_PyTok_Span){
+    tokenizer_mode *mode = _PyLexer_CurrentMode(tok);
+    mode->expr_span = (_PyTok_Span){
         _PyLexer_BufferOffset(tok, tok->cur), -1};
-    tokenizer_comments *comments = current_comments(tok);
+    tokenizer_comments *comments = mode->comments;
     if (comments != NULL) {
         comments->count = 0;
     }
@@ -411,11 +400,6 @@ _PyLexer_get_fstring_mode(struct tok_state *tok, tokenizer_mode* current_tok, st
 
     p_start = tok->start;
     p_end = tok->cur;
-    tokenizer_comments *comments = current_comments(tok);
-    if (comments != NULL) {
-        tok->ftstring_comments = comments->previous;
-        PyMem_Free(comments);
-    }
     _PyLexer_PopMode(tok);
     return MAKE_TOKEN(FTSTRING_END(current_tok));