bpo-32372: Move __debug__ optimization to the AST level. (#4925)

3 files changed, 47 insertions, 52 deletions

diff --git a/Include/compile.h b/Include/compile.h
index 2e9accf..edfd8bb 100644
--- a/Include/compile.h
+++ b/Include/compile.h
@@ -75,7 +75,7 @@ PyAPI_FUNC(PyObject*) _Py_Mangle(PyObject *p, PyObject *name);
 #define PY_INVALID_STACK_EFFECT INT_MAX
 PyAPI_FUNC(int) PyCompile_OpcodeStackEffect(int opcode, int oparg);
 
-PyAPI_FUNC(int) _PyAST_Optimize(struct _mod *, PyArena *arena);
+PyAPI_FUNC(int) _PyAST_Optimize(struct _mod *, PyArena *arena, int optimize);
 
 #ifdef __cplusplus
 }
diff --git a/Python/ast_opt.c b/Python/ast_opt.c
index 2f659d0..f98a28a 100644
--- a/Python/ast_opt.c
+++ b/Python/ast_opt.c
@@ -75,7 +75,7 @@ unary_not(PyObject *v)
 }
 
 static int
-fold_unaryop(expr_ty node, PyArena *arena)
+fold_unaryop(expr_ty node, PyArena *arena, int optimize)
 {
     expr_ty arg = node->v.UnaryOp.operand;
 
@@ -252,7 +252,7 @@ safe_mod(PyObject *v, PyObject *w)
 }
 
 static int
-fold_binop(expr_ty node, PyArena *arena)
+fold_binop(expr_ty node, PyArena *arena, int optimize)
 {
     expr_ty lhs, rhs;
     lhs = node->v.BinOp.left;
@@ -334,7 +334,7 @@ make_const_tuple(asdl_seq *elts)
 }
 
 static int
-fold_tuple(expr_ty node, PyArena *arena)
+fold_tuple(expr_ty node, PyArena *arena, int optimize)
 {
     PyObject *newval;
 
@@ -346,7 +346,7 @@ fold_tuple(expr_ty node, PyArena *arena)
 }
 
 static int
-fold_subscr(expr_ty node, PyArena *arena)
+fold_subscr(expr_ty node, PyArena *arena, int optimize)
 {
     PyObject *newval;
     expr_ty arg, idx;
@@ -374,7 +374,7 @@ fold_subscr(expr_ty node, PyArena *arena)
    in "for" loop and comprehensions.
 */
 static int
-fold_iter(expr_ty arg, PyArena *arena)
+fold_iter(expr_ty arg, PyArena *arena, int optimize)
 {
     PyObject *newval;
     if (arg->kind == List_kind) {
@@ -393,7 +393,7 @@ fold_iter(expr_ty arg, PyArena *arena)
 }
 
 static int
-fold_compare(expr_ty node, PyArena *arena)
+fold_compare(expr_ty node, PyArena *arena, int optimize)
 {
     asdl_int_seq *ops;
     asdl_seq *args;
@@ -407,29 +407,29 @@ fold_compare(expr_ty node, PyArena *arena)
     i = asdl_seq_LEN(ops) - 1;
     int op = asdl_seq_GET(ops, i);
     if (op == In || op == NotIn) {
-        if (!fold_iter((expr_ty)asdl_seq_GET(args, i), arena)) {
+        if (!fold_iter((expr_ty)asdl_seq_GET(args, i), arena, optimize)) {
             return 0;
         }
     }
     return 1;
 }
 
-static int astfold_mod(mod_ty node_, PyArena* ctx_);
-static int astfold_stmt(stmt_ty node_, PyArena* ctx_);
-static int astfold_expr(expr_ty node_, PyArena* ctx_);
-static int astfold_arguments(arguments_ty node_, PyArena* ctx_);
-static int astfold_comprehension(comprehension_ty node_, PyArena* ctx_);
-static int astfold_keyword(keyword_ty node_, PyArena* ctx_);
-static int astfold_slice(slice_ty node_, PyArena* ctx_);
-static int astfold_arg(arg_ty node_, PyArena* ctx_);
-static int astfold_withitem(withitem_ty node_, PyArena* ctx_);
-static int astfold_excepthandler(excepthandler_ty node_, PyArena* ctx_);
+static int astfold_mod(mod_ty node_, PyArena *ctx_, int optimize_);
+static int astfold_stmt(stmt_ty node_, PyArena *ctx_, int optimize_);
+static int astfold_expr(expr_ty node_, PyArena *ctx_, int optimize_);
+static int astfold_arguments(arguments_ty node_, PyArena *ctx_, int optimize_);
+static int astfold_comprehension(comprehension_ty node_, PyArena *ctx_, int optimize_);
+static int astfold_keyword(keyword_ty node_, PyArena *ctx_, int optimize_);
+static int astfold_slice(slice_ty node_, PyArena *ctx_, int optimize_);
+static int astfold_arg(arg_ty node_, PyArena *ctx_, int optimize_);
+static int astfold_withitem(withitem_ty node_, PyArena *ctx_, int optimize_);
+static int astfold_excepthandler(excepthandler_ty node_, PyArena *ctx_, int optimize_);
 #define CALL(FUNC, TYPE, ARG) \
-    if (!FUNC((ARG), ctx_)) \
+    if (!FUNC((ARG), ctx_, optimize_)) \
         return 0;
 
 #define CALL_OPT(FUNC, TYPE, ARG) \
-    if ((ARG) != NULL && !FUNC((ARG), ctx_)) \
+    if ((ARG) != NULL && !FUNC((ARG), ctx_, optimize_)) \
         return 0;
 
 #define CALL_SEQ(FUNC, TYPE, ARG) { \
@@ -437,7 +437,7 @@ static int astfold_excepthandler(excepthandler_ty node_, PyArena* ctx_);
     asdl_seq *seq = (ARG); /* avoid variable capture */ \
     for (i = 0; i < asdl_seq_LEN(seq); i++) { \
         TYPE elt = (TYPE)asdl_seq_GET(seq, i); \
-        if (elt != NULL && !FUNC(elt, ctx_)) \
+        if (elt != NULL && !FUNC(elt, ctx_, optimize_)) \
             return 0; \
     } \
 }
@@ -447,13 +447,13 @@ static int astfold_excepthandler(excepthandler_ty node_, PyArena* ctx_);
     asdl_int_seq *seq = (ARG); /* avoid variable capture */ \
     for (i = 0; i < asdl_seq_LEN(seq); i++) { \
         TYPE elt = (TYPE)asdl_seq_GET(seq, i); \
-        if (!FUNC(elt, ctx_)) \
+        if (!FUNC(elt, ctx_, optimize_)) \
             return 0; \
     } \
 }
 
 static int
-astfold_mod(mod_ty node_, PyArena* ctx_)
+astfold_mod(mod_ty node_, PyArena *ctx_, int optimize_)
 {
     switch (node_->kind) {
     case Module_kind:
@@ -475,7 +475,7 @@ astfold_mod(mod_ty node_, PyArena* ctx_)
 }
 
 static int
-astfold_expr(expr_ty node_, PyArena* ctx_)
+astfold_expr(expr_ty node_, PyArena *ctx_, int optimize_)
 {
     switch (node_->kind) {
     case BoolOp_kind:
@@ -567,6 +567,11 @@ astfold_expr(expr_ty node_, PyArena* ctx_)
         CALL_SEQ(astfold_expr, expr_ty, node_->v.Tuple.elts);
         CALL(fold_tuple, expr_ty, node_);
         break;
+    case Name_kind:
+        if (_PyUnicode_EqualToASCIIString(node_->v.Name.id, "__debug__")) {
+            return make_const(node_, PyBool_FromLong(!optimize_), ctx_);
+        }
+        break;
     default:
         break;
     }
@@ -574,7 +579,7 @@ astfold_expr(expr_ty node_, PyArena* ctx_)
 }
 
 static int
-astfold_slice(slice_ty node_, PyArena* ctx_)
+astfold_slice(slice_ty node_, PyArena *ctx_, int optimize_)
 {
     switch (node_->kind) {
     case Slice_kind:
@@ -595,14 +600,14 @@ astfold_slice(slice_ty node_, PyArena* ctx_)
 }
 
 static int
-astfold_keyword(keyword_ty node_, PyArena* ctx_)
+astfold_keyword(keyword_ty node_, PyArena *ctx_, int optimize_)
 {
     CALL(astfold_expr, expr_ty, node_->value);
     return 1;
 }
 
 static int
-astfold_comprehension(comprehension_ty node_, PyArena* ctx_)
+astfold_comprehension(comprehension_ty node_, PyArena *ctx_, int optimize_)
 {
     CALL(astfold_expr, expr_ty, node_->target);
     CALL(astfold_expr, expr_ty, node_->iter);
@@ -613,7 +618,7 @@ astfold_comprehension(comprehension_ty node_, PyArena* ctx_)
 }
 
 static int
-astfold_arguments(arguments_ty node_, PyArena* ctx_)
+astfold_arguments(arguments_ty node_, PyArena *ctx_, int optimize_)
 {
     CALL_SEQ(astfold_arg, arg_ty, node_->args);
     CALL_OPT(astfold_arg, arg_ty, node_->vararg);
@@ -625,14 +630,14 @@ astfold_arguments(arguments_ty node_, PyArena* ctx_)
 }
 
 static int
-astfold_arg(arg_ty node_, PyArena* ctx_)
+astfold_arg(arg_ty node_, PyArena *ctx_, int optimize_)
 {
     CALL_OPT(astfold_expr, expr_ty, node_->annotation);
     return 1;
 }
 
 static int
-astfold_stmt(stmt_ty node_, PyArena* ctx_)
+astfold_stmt(stmt_ty node_, PyArena *ctx_, int optimize_)
 {
     switch (node_->kind) {
     case FunctionDef_kind:
@@ -728,7 +733,7 @@ astfold_stmt(stmt_ty node_, PyArena* ctx_)
 }
 
 static int
-astfold_excepthandler(excepthandler_ty node_, PyArena* ctx_)
+astfold_excepthandler(excepthandler_ty node_, PyArena *ctx_, int optimize_)
 {
     switch (node_->kind) {
     case ExceptHandler_kind:
@@ -742,7 +747,7 @@ astfold_excepthandler(excepthandler_ty node_, PyArena* ctx_)
 }
 
 static int
-astfold_withitem(withitem_ty node_, PyArena* ctx_)
+astfold_withitem(withitem_ty node_, PyArena *ctx_, int optimize_)
 {
     CALL(astfold_expr, expr_ty, node_->context_expr);
     CALL_OPT(astfold_expr, expr_ty, node_->optional_vars);
@@ -755,9 +760,9 @@ astfold_withitem(withitem_ty node_, PyArena* ctx_)
 #undef CALL_INT_SEQ
 
 int
-_PyAST_Optimize(mod_ty mod, PyArena *arena)
+_PyAST_Optimize(mod_ty mod, PyArena *arena, int optimize)
 {
-    int ret = astfold_mod(mod, arena);
+    int ret = astfold_mod(mod, arena, optimize);
     assert(ret || PyErr_Occurred());
     return ret;
 }
diff --git a/Python/compile.c b/Python/compile.c
index 392de40..e5745d1 100644
--- a/Python/compile.c
+++ b/Python/compile.c
@@ -191,7 +191,7 @@ static void compiler_pop_fblock(struct compiler *, enum fblocktype,
 static int compiler_in_loop(struct compiler *);
 
 static int inplace_binop(struct compiler *, operator_ty);
-static int expr_constant(struct compiler *, expr_ty);
+static int expr_constant(expr_ty);
 
 static int compiler_with(struct compiler *, stmt_ty, int);
 static int compiler_async_with(struct compiler *, stmt_ty, int);
@@ -331,7 +331,7 @@ PyAST_CompileObject(mod_ty mod, PyObject *filename, PyCompilerFlags *flags,
     c.c_optimize = (optimize == -1) ? Py_OptimizeFlag : optimize;
     c.c_nestlevel = 0;
 
-    if (!_PyAST_Optimize(mod, arena)) {
+    if (!_PyAST_Optimize(mod, arena, c.c_optimize)) {
         goto finally;
     }
 
@@ -1332,15 +1332,13 @@ is_const(expr_ty e)
     case Ellipsis_kind:
     case NameConstant_kind:
         return 1;
-    case Name_kind:
-        return _PyUnicode_EqualToASCIIString(e->v.Name.id, "__debug__");
     default:
         return 0;
     }
 }
 
 static PyObject *
-get_const_value(struct compiler *c, expr_ty e)
+get_const_value(expr_ty e)
 {
     switch (e->kind) {
     case Constant_kind:
@@ -1355,9 +1353,6 @@ get_const_value(struct compiler *c, expr_ty e)
         return Py_Ellipsis;
     case NameConstant_kind:
         return e->v.NameConstant.value;
-    case Name_kind:
-        assert(_PyUnicode_EqualToASCIIString(e->v.Name.id, "__debug__"));
-        return c->c_optimize ? Py_False : Py_True;
     default:
         Py_UNREACHABLE();
     }
@@ -2217,7 +2212,7 @@ compiler_if(struct compiler *c, stmt_ty s)
     if (end == NULL)
         return 0;
 
-    constant = expr_constant(c, s->v.If.test);
+    constant = expr_constant(s->v.If.test);
     /* constant = 0: "if 0"
      * constant = 1: "if 1", "if 2", ...
      * constant = -1: rest */
@@ -2363,7 +2358,7 @@ static int
 compiler_while(struct compiler *c, stmt_ty s)
 {
     basicblock *loop, *orelse, *end, *anchor = NULL;
-    int constant = expr_constant(c, s->v.While.test);
+    int constant = expr_constant(s->v.While.test);
 
     if (constant == 0) {
         if (s->v.While.orelse)
@@ -3098,11 +3093,6 @@ compiler_nameop(struct compiler *c, identifier name, expr_context_ty ctx)
            !_PyUnicode_EqualToASCIIString(name, "True") &&
            !_PyUnicode_EqualToASCIIString(name, "False"));
 
-    if (ctx == Load && _PyUnicode_EqualToASCIIString(name, "__debug__")) {
-        ADDOP_O(c, LOAD_CONST, c->c_optimize ? Py_False : Py_True, consts);
-        return 1;
-    }
-
     mangled = _Py_Mangle(c->u->u_private, name);
     if (!mangled)
         return 0;
@@ -3370,7 +3360,7 @@ compiler_subdict(struct compiler *c, expr_ty e, Py_ssize_t begin, Py_ssize_t end
             return 0;
         }
         for (i = begin; i < end; i++) {
-            key = get_const_value(c, (expr_ty)asdl_seq_GET(e->v.Dict.keys, i));
+            key = get_const_value((expr_ty)asdl_seq_GET(e->v.Dict.keys, i));
             Py_INCREF(key);
             PyTuple_SET_ITEM(keys, i - begin, key);
         }
@@ -4140,10 +4130,10 @@ compiler_visit_keyword(struct compiler *c, keyword_ty k)
  */
 
 static int
-expr_constant(struct compiler *c, expr_ty e)
+expr_constant(expr_ty e)
 {
     if (is_const(e)) {
-        return PyObject_IsTrue(get_const_value(c, e));
+        return PyObject_IsTrue(get_const_value(e));
     }
     return -1;
 }