Fix the optimizer's understanding of exponentiation

Author: brandtbucher

Python/optimizer_bytecodes.c

@@ -167,24 +167,81 @@ 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)) {
+            res = sym_new_unknown(ctx);
+            goto binary_op_done;
+        }
+        if (oparg == NB_POWER || oparg == NB_INPLACE_POWER) {
+            // This one's fun: the *type* of the result depends on the *values*
+            // being exponentiated. But exponents with one constant part are
+            // reasonably common, so it's probably worth trying to be precise:
+            PyObject *lhs_const = sym_get_const(left);
+            PyObject *rhs_const = sym_get_const(right);
+            if (lhs_int && rhs_int) {
+                if (rhs_const == NULL) {
+                    // Unknown RHS means either int or float:
+                    res = sym_new_unknown(ctx);
+                    goto binary_op_done;
+                }
+                if (!_PyLong_IsNegative((PyLongObject *)rhs_const)) {
+                    // Non-negative RHS means int:
+                    res = sym_new_type(ctx, &PyLong_Type);
+                    goto binary_op_done;
+                }
+                // Negative RHS uses float_pow...
             }
-            else {
-                /* For any other op combining ints/floats the result is a float */
+            // Negative LHS *and* non-integral RHS means complex. So we need to
+            // disprove at least one to prove a float result:
+            if (rhs_int) {
+                // Integral RHS means float:
                 res = sym_new_type(ctx, &PyFloat_Type);
+                goto binary_op_done;
+            }
+            if (rhs_const) {
+                double rhs_double = PyFloat_AS_DOUBLE(rhs_const);
+                if (rhs_double == floor(rhs_double)) {
+                    // Integral RHS means float:
+                    res = sym_new_type(ctx, &PyFloat_Type);
+                    goto binary_op_done;
+                }
+            }
+            if (lhs_const) {
+                if (lhs_int) {
+                    if (!_PyLong_IsNegative((PyLongObject *)lhs_const)) {
+                        // Non-negative LHS means float:
+                        res = sym_new_type(ctx, &PyFloat_Type);
+                        goto binary_op_done;
+                    }
+                }
+                else if (0.0 <= PyFloat_AS_DOUBLE(lhs_const)) {
+                    // Non-negative LHS means float:
+                    res = sym_new_type(ctx, &PyFloat_Type);
+                    goto binary_op_done;
+                }
+                if (rhs_const) {
+                    // If we have two constants and failed to disprove that it's
+                    // complex, then it's complex:
+                    res = sym_new_type(ctx, &PyComplex_Type);
+                    goto binary_op_done;
+                }
             }
+            // Couldn't prove anything. It's either float or complex:
+            res = sym_new_unknown(ctx);
+        }
+        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);
         }
+binary_op_done:
     }
 
     op(_BINARY_OP_ADD_INT, (left, right -- res)) {

Python/optimizer_cases.c.h

@@ -596,6 +596,7 @@ def has_error_without_pop(op: parser.InstDef) -> bool:
     "_PyLong_CompactValue",
     "_PyLong_DigitCount",
     "_PyLong_IsCompact",
+    "_PyLong_IsNegative",
     "_PyLong_IsNonNegativeCompact",
     "_PyLong_IsZero",
     "_PyLong_Multiply",
@@ -634,6 +635,7 @@ def has_error_without_pop(op: parser.InstDef) -> bool:
     "advance_backoff_counter",
     "assert",
     "backoff_counter_triggers",
+    "floor",
     "initial_temperature_backoff_counter",
     "maybe_lltrace_resume_frame",
     "restart_backoff_counter",