GH-127809: Fix the JIT's understanding of ** (GH-127844)

This commit is contained in:
Brandt Bucher 2025-01-07 17:25:48 -08:00 committed by GitHub
parent e08b28235a
commit 65ae3d5a73
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
8 changed files with 200 additions and 27 deletions

View file

@ -167,23 +167,56 @@ dummy_func(void) {
}
op(_BINARY_OP, (left, right -- res)) {
PyTypeObject *ltype = sym_get_type(left);
PyTypeObject *rtype = sym_get_type(right);
if (ltype != NULL && (ltype == &PyLong_Type || ltype == &PyFloat_Type) &&
rtype != NULL && (rtype == &PyLong_Type || rtype == &PyFloat_Type))
{
if (oparg != NB_TRUE_DIVIDE && oparg != NB_INPLACE_TRUE_DIVIDE &&
ltype == &PyLong_Type && rtype == &PyLong_Type) {
/* If both inputs are ints and the op is not division the result is an int */
res = sym_new_type(ctx, &PyLong_Type);
bool lhs_int = sym_matches_type(left, &PyLong_Type);
bool rhs_int = sym_matches_type(right, &PyLong_Type);
bool lhs_float = sym_matches_type(left, &PyFloat_Type);
bool rhs_float = sym_matches_type(right, &PyFloat_Type);
if (!((lhs_int || lhs_float) && (rhs_int || rhs_float))) {
// There's something other than an int or float involved:
res = sym_new_unknown(ctx);
}
else if (oparg == NB_POWER || oparg == NB_INPLACE_POWER) {
// This one's fun... the *type* of the result depends on the
// *values* being exponentiated. However, exponents with one
// constant part are reasonably common, so it's probably worth
// trying to infer some simple cases:
// - A: 1 ** 1 -> 1 (int ** int -> int)
// - B: 1 ** -1 -> 1.0 (int ** int -> float)
// - C: 1.0 ** 1 -> 1.0 (float ** int -> float)
// - D: 1 ** 1.0 -> 1.0 (int ** float -> float)
// - E: -1 ** 0.5 ~> 1j (int ** float -> complex)
// - F: 1.0 ** 1.0 -> 1.0 (float ** float -> float)
// - G: -1.0 ** 0.5 ~> 1j (float ** float -> complex)
if (rhs_float) {
// Case D, E, F, or G... can't know without the sign of the LHS
// or whether the RHS is whole, which isn't worth the effort:
res = sym_new_unknown(ctx);
}
else {
/* For any other op combining ints/floats the result is a float */
else if (lhs_float) {
// Case C:
res = sym_new_type(ctx, &PyFloat_Type);
}
else if (!sym_is_const(right)) {
// Case A or B... can't know without the sign of the RHS:
res = sym_new_unknown(ctx);
}
else if (_PyLong_IsNegative((PyLongObject *)sym_get_const(right))) {
// Case B:
res = sym_new_type(ctx, &PyFloat_Type);
}
else {
// Case A:
res = sym_new_type(ctx, &PyLong_Type);
}
}
else if (oparg == NB_TRUE_DIVIDE || oparg == NB_INPLACE_TRUE_DIVIDE) {
res = sym_new_type(ctx, &PyFloat_Type);
}
else if (lhs_int && rhs_int) {
res = sym_new_type(ctx, &PyLong_Type);
}
else {
res = sym_new_unknown(ctx);
res = sym_new_type(ctx, &PyFloat_Type);
}
}