gh-123958: apply docstring removal optimization in ast_opt instead of codegen (#123959)

1 files changed, 21 insertions, 0 deletions

diff --git a/Python/ast_opt.c b/Python/ast_opt.c
index 5a51305..f5b0475 100644
--- a/Python/ast_opt.c
+++ b/Python/ast_opt.c
@@ -674,9 +674,30 @@ static int astfold_type_param(type_param_ty node_, PyArena *ctx_, _PyASTOptimize
 
 
 static int
+stmt_seq_remove_item(asdl_stmt_seq *stmts, Py_ssize_t idx)
+{
+    if (idx >= asdl_seq_LEN(stmts)) {
+        return 0;
+    }
+    for (Py_ssize_t i = idx; i < asdl_seq_LEN(stmts) - 1; i++) {
+        stmt_ty st = (stmt_ty)asdl_seq_GET(stmts, i+1);
+        asdl_seq_SET(stmts, i, st);
+    }
+    stmts->size--;
+    return 1;
+}
+
+static int
 astfold_body(asdl_stmt_seq *stmts, PyArena *ctx_, _PyASTOptimizeState *state)
 {
     int docstring = _PyAST_GetDocString(stmts) != NULL;
+    if (docstring && (state->optimize >= 2)) {
+        /* remove the docstring */
+        if (!stmt_seq_remove_item(stmts, 0)) {
+            return 0;
+        }
+        docstring = 0;
+    }
     CALL_SEQ(astfold_stmt, stmt, stmts);
     if (!docstring && _PyAST_GetDocString(stmts) != NULL) {
         stmt_ty st = (stmt_ty)asdl_seq_GET(stmts, 0);