Skip to content

gh-128706: Evaluate constant comparisons in fold_compare #128705

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
wants to merge 9 commits into from
Closed

gh-128706: Evaluate constant comparisons in fold_compare #128705

Show file tree
Hide file tree
Changes from 5 commits
Commits
Show all changes
9 commits
Select commit Hold shift + click to select a range
de026a9
Evaluate constant comparions in `fold_compare`
WolframAlph Jan 10, 2025
22bcb10
do not fold is/isnot
WolframAlph Jan 10, 2025
713fde5
fix failing tests due to unexpected comparion folding
WolframAlph Jan 10, 2025
285f266
add tests for comparison folding
WolframAlph Jan 10, 2025
c3abf4a
add news entry
WolframAlph Jan 10, 2025
5cb7cd8
address comments
WolframAlph Jan 11, 2025
65093fb
address comments
WolframAlph Jan 11, 2025
891bff5
add whats new entry
WolframAlph Jan 11, 2025
4f5f1b5
address coments
WolframAlph Jan 11, 2025
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
65 changes: 58 additions & 7 deletions Lib/test/test_ast/test_ast.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -3180,7 +3180,8 @@ def create_unaryop(operand):
self.assert_ast(result_code, non_optimized_target, optimized_target)

def test_folding_not(self):
code = "not (1 %s (1,))"
# use list as left-hand side to avoid folding constant expression to True/False
code = "not ([] %s (1,))"
operators = {
"in": ast.In(),
"is": ast.Is(),
Expand All @@ -3192,7 +3193,7 @@ def test_folding_not(self):

def create_notop(operand):
return ast.UnaryOp(op=ast.Not(), operand=ast.Compare(
left=ast.Constant(value=1),
left=ast.List(),
ops=[operators[operand]],
comparators=[ast.Tuple(elts=[ast.Constant(value=1)])]
))
Expand All @@ -3201,7 +3202,7 @@ def create_notop(operand):
result_code = code % op
non_optimized_target = self.wrap_expr(create_notop(op))
optimized_target = self.wrap_expr(
ast.Compare(left=ast.Constant(1), ops=[opt_operators[op]], comparators=[ast.Constant(value=(1,))])
ast.Compare(left=ast.List(), ops=[opt_operators[op]], comparators=[ast.Constant(value=(1,))])
)

with self.subTest(
Expand Down Expand Up @@ -3239,8 +3240,58 @@ def test_folding_tuple(self):

self.assert_ast(code, non_optimized_target, optimized_target)

def test_folding_comparator(self):
code = "1 %s %s1%s"
def test_folding_compare(self):
true = self.wrap_expr(ast.Constant(value=True))
false = self.wrap_expr(ast.Constant(value=False))

folded_cases = (
("3 > 2 > 1", (ast.Constant(3), [ast.Gt(), ast.Gt()], [ast.Constant(value=2), ast.Constant(value=1)]), true),
("3 > 4 > 1", (ast.Constant(3), [ast.Gt(), ast.Gt()], [ast.Constant(value=4), ast.Constant(value=1)]), false),
("3 >= 3 >= 1", (ast.Constant(3), [ast.GtE(), ast.GtE()], [ast.Constant(value=3), ast.Constant(value=1)]), true),
("3 >= 4 >= 1", (ast.Constant(3), [ast.GtE(), ast.GtE()], [ast.Constant(value=4), ast.Constant(value=1)]), false),
("1 < 2 < 3", (ast.Constant(1), [ast.Lt(), ast.Lt()], [ast.Constant(value=2), ast.Constant(value=3)]), true),
("1 < 0 < 3", (ast.Constant(1), [ast.Lt(), ast.Lt()], [ast.Constant(value=0), ast.Constant(value=3)]), false),
("1 <= 2 <= 3", (ast.Constant(1), [ast.LtE(), ast.LtE()], [ast.Constant(value=2), ast.Constant(value=3)]), true),
("1 <= 0 <= 3", (ast.Constant(1), [ast.LtE(), ast.LtE()], [ast.Constant(value=0), ast.Constant(value=3)]), false),
("1 == 1.0 == True", (ast.Constant(1), [ast.Eq(), ast.Eq()], [ast.Constant(value=1.0), ast.Constant(value=True)]), true),
("1 == 2 == True", (ast.Constant(1), [ast.Eq(), ast.Eq()], [ast.Constant(value=2), ast.Constant(value=True)]), false),
("1 != 2 != 3", (ast.Constant(1), [ast.NotEq(), ast.NotEq()], [ast.Constant(value=2), ast.Constant(value=3)]), true),
("1 != 1 != 3", (ast.Constant(1), [ast.NotEq(), ast.NotEq()], [ast.Constant(value=1), ast.Constant(value=3)]), false),
("1 in [1, 2]", (ast.Constant(1), [ast.In()], [ast.List(elts=[ast.Constant(1), ast.Constant(2)])]), true),
("1 in [2, 2]", (ast.Constant(1), [ast.In()], [ast.List(elts=[ast.Constant(2), ast.Constant(2)])]), false),
("1 not in [1, 2]", (ast.Constant(1), [ast.NotIn()], [ast.List(elts=[ast.Constant(1), ast.Constant(2)])]), false),
("1 not in [2, 2]", (ast.Constant(1), [ast.NotIn()], [ast.List(elts=[ast.Constant(2), ast.Constant(2)])]), true),
)

for code, original, folded in folded_cases:
left, ops, comparators = original
unfolded = self.wrap_expr(ast.Compare(left=left, ops=ops, comparators=comparators))
self.assert_ast(code=code, non_optimized_target=unfolded, optimized_target=folded)

# these should stay as they were
unfolded_cases = (
("3 > 2 > []", ast.Compare(left=ast.Constant(3), ops=[ast.Gt(), ast.Gt()], comparators=[ast.Constant(2), ast.List()])),
("1 > [] > 0", ast.Compare(left=ast.Constant(1), ops=[ast.Gt(), ast.Gt()], comparators=[ast.List(), ast.Constant(0)])),
("1 >= [] >= 0", ast.Compare(left=ast.Constant(1), ops=[ast.GtE(), ast.GtE()], comparators=[ast.List(), ast.Constant(0)])),
("1 < [] < 0", ast.Compare(left=ast.Constant(1), ops=[ast.Lt(), ast.Lt()], comparators=[ast.List(), ast.Constant(0)])),
("1 <= [] <= 0", ast.Compare(left=ast.Constant(1), ops=[ast.LtE(), ast.LtE()], comparators=[ast.List(), ast.Constant(0)])),
("1 == [] == 0", ast.Compare(left=ast.Constant(1), ops=[ast.Eq(), ast.Eq()], comparators=[ast.List(), ast.Constant(0)])),
("1 != [] != 0", ast.Compare(left=ast.Constant(1), ops=[ast.NotEq(), ast.NotEq()], comparators=[ast.List(), ast.Constant(0)])),
("1 is 1", ast.Compare(left=ast.Constant(1), ops=[ast.Is()], comparators=[ast.Constant(1)])),
("1 is not 1", ast.Compare(left=ast.Constant(1), ops=[ast.IsNot()], comparators=[ast.Constant(1)])),
# invalid also should stay as they were
("1 in 1", ast.Compare(left=ast.Constant(1), ops=[ast.In()], comparators=[ast.Constant(1)])),
("1 not in 1", ast.Compare(left=ast.Constant(1), ops=[ast.NotIn()], comparators=[ast.Constant(1)])),
)

for code, expected in unfolded_cases:
self.assertTrue(ast.compare(ast.parse(code), self.wrap_expr(expected)))

def test_folding_comparator_list_set_subst(self):
"""Test substitution of list/set with tuple/frozenset in expressions like "1 in [1]" or "1 in {1}" """

# use list as left-hand side to avoid folding constant comparison expression to True/False
code = "[] %s %s1%s"
operators = [("in", ast.In()), ("not in", ast.NotIn())]
braces = [
("[", "]", ast.List, (1,)),
Expand All @@ -3249,11 +3300,11 @@ def test_folding_comparator(self):
for left, right, non_optimized_comparator, optimized_comparator in braces:
for op, node in operators:
non_optimized_target = self.wrap_expr(ast.Compare(
left=ast.Constant(1), ops=[node],
left=ast.List(), ops=[node],
comparators=[non_optimized_comparator(elts=[ast.Constant(1)])]
))
optimized_target = self.wrap_expr(ast.Compare(
left=ast.Constant(1), ops=[node],
left=ast.List(), ops=[node],
comparators=[ast.Constant(value=optimized_comparator)]
))
self.assert_ast(code % (op, left, right), non_optimized_target, optimized_target)
Expand Down
1 change: 1 addition & 0 deletions Misc/NEWS.d/next/Core_and_Builtins/2025-01-10-18-42-57.gh-issue-128706.4T2XSt.rst
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -0,0 +1 @@
Add constant folding for constant comparisons
84 changes: 84 additions & 0 deletions Python/ast_opt.c
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -639,6 +639,90 @@ fold_compare(expr_ty node, PyArena *arena, _PyASTOptimizeState *state)
return 0;
}
}

static const int richcompare_table[] = {
[Eq] = Py_EQ,
[NotEq] = Py_NE,
[Gt] = Py_GT,
[Lt] = Py_LT,
[GtE] = Py_GE,
[LtE] = Py_LE,
};

if (node->v.Compare.left->kind == Constant_kind) {

PyObject *lhs = node->v.Compare.left->v.Constant.value;

for (int i=0; i < asdl_seq_LEN(args); i++) {

expr_ty curr_expr = (expr_ty)asdl_seq_GET(args, i);

if (curr_expr->kind != Constant_kind) {
/* try to fold only if every comparator is constant */
goto exit;
}

int op = asdl_seq_GET(ops, i);

if (op == Is || op == IsNot) {
/*
Do not fold expression for now if "is"/"is not" is present.
It breaks expected syntax warnings. For example:
>>> 1 is 1
<python-input-0>:1: SyntaxWarning: "is" with 'int' literal. Did you mean "=="?
*/
goto exit;
}

PyObject *rhs = curr_expr->v.Constant.value;

switch (op) {
case Eq: case NotEq:
case Gt: case Lt:
case GtE: case LtE:
{
Copy link
Member

Choose a reason for hiding this comment

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

Suggested change
case GtE: case LtE:
{
case GtE:
case LtE: {

Copy link
Member

Choose a reason for hiding this comment

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

And have each case on each line even though they can be paired.

int res = PyObject_RichCompareBool(lhs, rhs, richcompare_table[op]);
if (res < 0) {
/* error */
Copy link
Member

Choose a reason for hiding this comment

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

By convention res < 0 indicates an error so this comment is not really needed IMO.

Copy link
Contributor Author

Choose a reason for hiding this comment

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

Had it for myself while developing. I can remove

if (PyErr_Occurred()) {
return make_const(node, NULL, arena);
}
return 0;
}
if (!res) {
/* shortcut, whole expression is False */
return make_const(node, Py_False, arena);
}
break;
}
case In:
case NotIn:
{
int res = PySequence_Contains(rhs, lhs);
if (res < 0) {
/* error */
if (PyErr_Occurred()) {
return make_const(node, NULL, arena);
}
return 0;
}
if (op == NotIn) {
res = !res;
}
if (!res) {
/* shortcut, whole expression is False */
return make_const(node, Py_False, arena);
}
break;
}
}
lhs = rhs;
}
/* whole expression is True */
return make_const(node, Py_True, arena);
}

exit:
return 1;
}

Expand Down
Loading