Skip to content

Commit de026a9

Browse files
committed
Evaluate constant comparions in fold_compare
1 parent baf6571 commit de026a9

File tree

1 file changed

+85
-0
lines changed

1 file changed

+85
-0
lines changed

Python/ast_opt.c

Lines changed: 85 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -639,6 +639,91 @@ fold_compare(expr_ty node, PyArena *arena, _PyASTOptimizeState *state)
639639
return 0;
640640
}
641641
}
642+
643+
static const int richcompare_table[] = {
644+
[Eq] = Py_EQ,
645+
[NotEq] = Py_NE,
646+
[Gt] = Py_GT,
647+
[Lt] = Py_LT,
648+
[GtE] = Py_GE,
649+
[LtE] = Py_LE,
650+
};
651+
652+
if (node->v.Compare.left->kind == Constant_kind) {
653+
PyObject *lhs = node->v.Compare.left->v.Constant.value;
654+
expr_ty curr_expr;
655+
656+
for (int i=0; i < asdl_seq_LEN(args); i++) {
657+
curr_expr = (expr_ty)asdl_seq_GET(args, i);
658+
659+
if (curr_expr->kind != Constant_kind) {
660+
goto exit;
661+
}
662+
663+
PyObject *rhs = curr_expr->v.Constant.value;
664+
int op = asdl_seq_GET(ops, i);
665+
int res;
666+
667+
switch (op) {
668+
case Eq: case NotEq:
669+
case Gt: case Lt:
670+
case GtE: case LtE:
671+
{
672+
res = PyObject_RichCompareBool(lhs, rhs, richcompare_table[op]);
673+
if (res < 0) {
674+
/* error */
675+
if (PyErr_Occurred()) {
676+
return make_const(node, NULL, arena);
677+
}
678+
return 0;
679+
}
680+
if (!res) {
681+
/* shortcut, whole expression is False */
682+
return make_const(node, Py_False, arena);
683+
}
684+
break;
685+
}
686+
case In:
687+
case NotIn:
688+
{
689+
res = PySequence_Contains(rhs, lhs);
690+
if (res < 0) {
691+
/* error */
692+
if (PyErr_Occurred()) {
693+
return make_const(node, NULL, arena);
694+
}
695+
return 0;
696+
}
697+
if (op == NotIn) {
698+
res = !res;
699+
}
700+
if (!res) {
701+
/* shortcut, whole expression is False */
702+
return make_const(node, Py_False, arena);
703+
}
704+
break;
705+
}
706+
case Is:
707+
case IsNot:
708+
{
709+
res = Py_Is(lhs, rhs);
710+
if (op == IsNot) {
711+
res = !res;
712+
}
713+
if (!res) {
714+
/* shortcut, whole expression is False */
715+
return make_const(node, Py_False, arena);
716+
}
717+
break;
718+
}
719+
}
720+
lhs = rhs;
721+
}
722+
/* whole expression is True */
723+
return make_const(node, Py_True, arena);
724+
}
725+
726+
exit:
642727
return 1;
643728
}
644729

0 commit comments

Comments
 (0)