Skip to content

gh-123958: move assert optimization from codegen to ast_opt #124143

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Closed
iritkatriel wants to merge 4 commits into from
Closed

gh-123958: move assert optimization from codegen to ast_opt #124143

Show file tree
Hide file tree
Changes from 1 commit
Commits
Show all changes
4 commits
Select commit Hold shift + click to select a range
8243017
gh-123958: move assert optimization from codegen to ast_opt
iritkatriel Sep 16, 2024
a465a77
and news
iritkatriel Sep 16, 2024
e8f935e
Merge branch 'main' into assert
Eclips4 Sep 16, 2024
d862d1f
Apply suggestions from code review
iritkatriel Sep 17, 2024
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
Next Next commit
gh-123958: move assert optimization from codegen to ast_opt
  • Loading branch information
@iritkatriel
iritkatriel committed Sep 16, 2024
commit 82430174ac1ce2e2aa9d320140f405d39feb06a9
3 changes: 2 additions & 1 deletion Include/internal/pycore_compile.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -37,7 +37,8 @@ extern int _PyAST_Optimize(
struct _mod *,
struct _arena *arena,
int optimize,
int ff_features);
int ff_features,
PyObject *filename);


typedef struct {
Expand Down
29 changes: 23 additions & 6 deletions Lib/test/test_builtin.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -443,12 +443,6 @@ async def arange(n):
'''a = [x async for x in (x async for x in arange(5))][1]''',
'''a, = [1 for x in {x async for x in arange(1)}]''',
'''a = [await asyncio.sleep(0, x) async for x in arange(2)][1]''',
# gh-121637: Make sure we correctly handle the case where the
# async code is optimized away
'''assert not await asyncio.sleep(0); a = 1''',
'''assert [x async for x in arange(1)]; a = 1''',
'''assert {x async for x in arange(1)}; a = 1''',
'''assert {x: x async for x in arange(1)}; a = 1''',
'''
if (a := 1) and __debug__:
async with asyncio.Lock() as l:
Expand Down Expand Up @@ -491,6 +485,29 @@ async def arange(n):
finally:
asyncio.set_event_loop_policy(policy)

def test_top_level_async_in_assert_compiles(self):
# gh-121637: do the right thing when async code is optimized away
modes = ('single', 'exec')
optimizations = (0, 1, 2)
code_samples = [
'''assert not await asyncio.sleep(0)''',
'''assert [x async for x in arange(1)]''',
'''assert {x async for x in arange(1)}''',
'''assert {x: x async for x in arange(1)}''',
]
for mode, code_sample, optimize in product(modes, code_samples, optimizations):
with self.subTest(mode=mode, code_sample=code_sample, optimize=optimize):
source = dedent(code_sample)
if (optimize):
Comment thread
iritkatriel marked this conversation as resolved.
Outdated
co = compile(source, '?', mode, optimize=optimize)
co_expected = compile('pass', '?', mode, optimize=optimize)
Comment thread
iritkatriel marked this conversation as resolved.
Outdated
self.assertEqual(list(co.co_lines()), list(co_expected.co_lines()))
Copy link
Copy Markdown
Member

@JelleZijlstra JelleZijlstra Sep 17, 2024

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This looks like we no longer raise an error in optimized mode if there is an invalid await. That seems wrong and doesn't align with what we decided on in #121637.

Copy link
Copy Markdown
Member Author

@iritkatriel iritkatriel Sep 17, 2024

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

If we want to stick with this then we probably can't optimise the assertions in ast_opt.

else:
with self.assertRaises(SyntaxError,
msg=f"source={source} mode={mode}"):
compile(source, '?', mode, optimize=optimize)


def test_compile_top_level_await_invalid_cases(self):
# helper function just to check we can run top=level async-for
async def arange(n):
Expand Down
33 changes: 32 additions & 1 deletion Python/ast_opt.c
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -10,6 +10,7 @@
typedef struct {
int optimize;
int ff_features;
PyObject *filename;

int recursion_depth; /* current recursion depth */
int recursion_limit; /* recursion limit */
Expand Down Expand Up @@ -51,6 +52,12 @@ make_const(expr_ty node, PyObject *val, PyArena *arena)
return 1;
}

static void
make_pass(stmt_ty node)
{
node->kind = Pass_kind;
}

#define COPY_NODE(TO, FROM) (memcpy((TO), (FROM), sizeof(struct _expr)))

static int
Expand Down Expand Up @@ -1005,6 +1012,29 @@ astfold_stmt(stmt_ty node_, PyArena *ctx_, _PyASTOptimizeState *state)
case Assert_kind:
CALL(astfold_expr, expr_ty, node_->v.Assert.test);
CALL_OPT(astfold_expr, expr_ty, node_->v.Assert.msg);
/* Always emit a warning if the test is a non-zero length tuple */
if ((node_->v.Assert.test->kind == Tuple_kind &&
asdl_seq_LEN(node_->v.Assert.test->v.Tuple.elts) > 0) ||
(node_->v.Assert.test->kind == Constant_kind &&
PyTuple_Check(node_->v.Assert.test->v.Constant.value) &&
PyTuple_Size(node_->v.Assert.test->v.Constant.value) > 0))
{
PyObject *msg = PyUnicode_FromString("assertion is always true, "
"perhaps remove parentheses?");
if (msg == NULL) {
return 0;
}
int ret = _PyErr_EmitSyntaxWarning(msg, state->filename,
node_->lineno, node_->col_offset + 1,
node_->end_lineno, node_->end_col_offset + 1);
Py_DECREF(msg);
if (ret < 0) {
return 0;
}
}
if (state->optimize) {
make_pass(node_);
Comment on lines +1035 to +1036
Copy link
Copy Markdown
Member

@Eclips4 Eclips4 Sep 17, 2024 •
edited
Loading

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Technically speaking, this is not the removal of the assert statement; it's replacement of assert with pass statement. Is it possible to completely remove the assert node? Or will it break some existing code?

Copy link
Copy Markdown
Member Author

@iritkatriel iritkatriel Sep 17, 2024

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Pass becomes a NOP, which gets removed later. This is how we get rid of code. However, I did forget to wipe out the location information (a NOP is not removed if it's the only instruction from its line in the source code).

}
break;
case Expr_kind:
CALL(astfold_expr, expr_ty, node_->v.Expr.value);
Expand Down Expand Up @@ -1125,14 +1155,15 @@ astfold_type_param(type_param_ty node_, PyArena *ctx_, _PyASTOptimizeState *stat
#undef CALL_SEQ

int
_PyAST_Optimize(mod_ty mod, PyArena *arena, int optimize, int ff_features)
_PyAST_Optimize(mod_ty mod, PyArena *arena, int optimize, int ff_features, PyObject *filename)
{
PyThreadState *tstate;
int starting_recursion_depth;

_PyASTOptimizeState state;
state.optimize = optimize;
state.ff_features = ff_features;
state.filename = filename;

/* Setup recursion depth check counters */
tstate = _PyThreadState_GET();
Expand Down
15 changes: 1 addition & 14 deletions Python/codegen.c
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -2792,20 +2792,7 @@ codegen_from_import(compiler *c, stmt_ty s)
static int
codegen_assert(compiler *c, stmt_ty s)
{
/* Always emit a warning if the test is a non-zero length tuple */
if ((s->v.Assert.test->kind == Tuple_kind &&
asdl_seq_LEN(s->v.Assert.test->v.Tuple.elts) > 0) ||
(s->v.Assert.test->kind == Constant_kind &&
PyTuple_Check(s->v.Assert.test->v.Constant.value) &&
PyTuple_Size(s->v.Assert.test->v.Constant.value) > 0))
{
RETURN_IF_ERROR(
_PyCompile_Warn(c, LOC(s), "assertion is always true, "
"perhaps remove parentheses?"));
}
if (OPTIMIZATION_LEVEL(c)) {
return SUCCESS;
}
assert(!OPTIMIZATION_LEVEL(c));
Copy link
Copy Markdown
Member Author

@iritkatriel iritkatriel Sep 16, 2024

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Left this here to convince ourselves of correctness. I'll remove it before committing.

NEW_JUMP_TARGET_LABEL(c, end);
RETURN_IF_ERROR(codegen_jump_if(c, LOC(s), s->v.Assert.test, end, 1));
ADDOP_I(c, LOC(s), LOAD_COMMON_CONSTANT, CONSTANT_ASSERTIONERROR);
Expand Down
4 changes: 2 additions & 2 deletions Python/compile.c
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -124,7 +124,7 @@ compiler_setup(compiler *c, mod_ty mod, PyObject *filename,
c->c_optimize = (optimize == -1) ? _Py_GetConfig()->optimization_level : optimize;
c->c_save_nested_seqs = false;

if (!_PyAST_Optimize(mod, arena, c->c_optimize, merged)) {
if (!_PyAST_Optimize(mod, arena, c->c_optimize, merged, filename)) {
return ERROR;
}
c->c_st = _PySymtable_Build(mod, filename, &c->c_future);
Expand Down Expand Up @@ -1381,7 +1381,7 @@ _PyCompile_AstOptimize(mod_ty mod, PyObject *filename, PyCompilerFlags *cf,
if (optimize == -1) {
optimize = _Py_GetConfig()->optimization_level;
}
if (!_PyAST_Optimize(mod, arena, optimize, flags)) {
if (!_PyAST_Optimize(mod, arena, optimize, flags, filename)) {
return -1;
}
return 0;
Expand Down