GH-130415: Use boolean guards to narrow types to values in the JIT (GH-130659)

1 files changed, 62 insertions, 53 deletions

diff --git a/Python/optimizer_bytecodes.c b/Python/optimizer_bytecodes.c
index f3625a1..c4e4b28 100644
--- a/Python/optimizer_bytecodes.c
+++ b/Python/optimizer_bytecodes.c
@@ -34,6 +34,7 @@ typedef struct _Py_UOpsAbstractFrame _Py_UOpsAbstractFrame;
 #define sym_tuple_getitem _Py_uop_sym_tuple_getitem
 #define sym_tuple_length _Py_uop_sym_tuple_length
 #define sym_is_immortal _Py_uop_sym_is_immortal
+#define sym_new_truthiness _Py_uop_sym_new_truthiness
 
 extern int
 optimize_to_bool(
@@ -198,11 +199,11 @@ dummy_func(void) {
                 // Case C:
                 res = sym_new_type(ctx, &PyFloat_Type);
             }
-            else if (!sym_is_const(right)) {
+            else if (!sym_is_const(ctx, 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))) {
+            else if (_PyLong_IsNegative((PyLongObject *)sym_get_const(ctx, right))) {
                 // Case B:
                 res = sym_new_type(ctx, &PyFloat_Type);
             }
@@ -223,13 +224,13 @@ dummy_func(void) {
     }
 
     op(_BINARY_OP_ADD_INT, (left, right -- res)) {
-        if (sym_is_const(left) && sym_is_const(right) &&
+        if (sym_is_const(ctx, left) && sym_is_const(ctx, right) &&
             sym_matches_type(left, &PyLong_Type) && sym_matches_type(right, &PyLong_Type))
         {
-            assert(PyLong_CheckExact(sym_get_const(left)));
-            assert(PyLong_CheckExact(sym_get_const(right)));
-            PyObject *temp = _PyLong_Add((PyLongObject *)sym_get_const(left),
-                                         (PyLongObject *)sym_get_const(right));
+            assert(PyLong_CheckExact(sym_get_const(ctx, left)));
+            assert(PyLong_CheckExact(sym_get_const(ctx, right)));
+            PyObject *temp = _PyLong_Add((PyLongObject *)sym_get_const(ctx, left),
+                                         (PyLongObject *)sym_get_const(ctx, right));
             if (temp == NULL) {
                 goto error;
             }
@@ -244,13 +245,13 @@ dummy_func(void) {
     }
 
     op(_BINARY_OP_SUBTRACT_INT, (left, right -- res)) {
-        if (sym_is_const(left) && sym_is_const(right) &&
+        if (sym_is_const(ctx, left) && sym_is_const(ctx, right) &&
             sym_matches_type(left, &PyLong_Type) && sym_matches_type(right, &PyLong_Type))
         {
-            assert(PyLong_CheckExact(sym_get_const(left)));
-            assert(PyLong_CheckExact(sym_get_const(right)));
-            PyObject *temp = _PyLong_Subtract((PyLongObject *)sym_get_const(left),
-                                              (PyLongObject *)sym_get_const(right));
+            assert(PyLong_CheckExact(sym_get_const(ctx, left)));
+            assert(PyLong_CheckExact(sym_get_const(ctx, right)));
+            PyObject *temp = _PyLong_Subtract((PyLongObject *)sym_get_const(ctx, left),
+                                              (PyLongObject *)sym_get_const(ctx, right));
             if (temp == NULL) {
                 goto error;
             }
@@ -265,13 +266,13 @@ dummy_func(void) {
     }
 
     op(_BINARY_OP_MULTIPLY_INT, (left, right -- res)) {
-        if (sym_is_const(left) && sym_is_const(right) &&
+        if (sym_is_const(ctx, left) && sym_is_const(ctx, right) &&
             sym_matches_type(left, &PyLong_Type) && sym_matches_type(right, &PyLong_Type))
         {
-            assert(PyLong_CheckExact(sym_get_const(left)));
-            assert(PyLong_CheckExact(sym_get_const(right)));
-            PyObject *temp = _PyLong_Multiply((PyLongObject *)sym_get_const(left),
-                                              (PyLongObject *)sym_get_const(right));
+            assert(PyLong_CheckExact(sym_get_const(ctx, left)));
+            assert(PyLong_CheckExact(sym_get_const(ctx, right)));
+            PyObject *temp = _PyLong_Multiply((PyLongObject *)sym_get_const(ctx, left),
+                                              (PyLongObject *)sym_get_const(ctx, right));
             if (temp == NULL) {
                 goto error;
             }
@@ -286,14 +287,14 @@ dummy_func(void) {
     }
 
     op(_BINARY_OP_ADD_FLOAT, (left, right -- res)) {
-        if (sym_is_const(left) && sym_is_const(right) &&
+        if (sym_is_const(ctx, left) && sym_is_const(ctx, right) &&
             sym_matches_type(left, &PyFloat_Type) && sym_matches_type(right, &PyFloat_Type))
         {
-            assert(PyFloat_CheckExact(sym_get_const(left)));
-            assert(PyFloat_CheckExact(sym_get_const(right)));
+            assert(PyFloat_CheckExact(sym_get_const(ctx, left)));
+            assert(PyFloat_CheckExact(sym_get_const(ctx, right)));
             PyObject *temp = PyFloat_FromDouble(
-                PyFloat_AS_DOUBLE(sym_get_const(left)) +
-                PyFloat_AS_DOUBLE(sym_get_const(right)));
+                PyFloat_AS_DOUBLE(sym_get_const(ctx, left)) +
+                PyFloat_AS_DOUBLE(sym_get_const(ctx, right)));
             if (temp == NULL) {
                 goto error;
             }
@@ -308,14 +309,14 @@ dummy_func(void) {
     }
 
     op(_BINARY_OP_SUBTRACT_FLOAT, (left, right -- res)) {
-        if (sym_is_const(left) && sym_is_const(right) &&
+        if (sym_is_const(ctx, left) && sym_is_const(ctx, right) &&
             sym_matches_type(left, &PyFloat_Type) && sym_matches_type(right, &PyFloat_Type))
         {
-            assert(PyFloat_CheckExact(sym_get_const(left)));
-            assert(PyFloat_CheckExact(sym_get_const(right)));
+            assert(PyFloat_CheckExact(sym_get_const(ctx, left)));
+            assert(PyFloat_CheckExact(sym_get_const(ctx, right)));
             PyObject *temp = PyFloat_FromDouble(
-                PyFloat_AS_DOUBLE(sym_get_const(left)) -
-                PyFloat_AS_DOUBLE(sym_get_const(right)));
+                PyFloat_AS_DOUBLE(sym_get_const(ctx, left)) -
+                PyFloat_AS_DOUBLE(sym_get_const(ctx, right)));
             if (temp == NULL) {
                 goto error;
             }
@@ -330,14 +331,14 @@ dummy_func(void) {
     }
 
     op(_BINARY_OP_MULTIPLY_FLOAT, (left, right -- res)) {
-        if (sym_is_const(left) && sym_is_const(right) &&
+        if (sym_is_const(ctx, left) && sym_is_const(ctx, right) &&
             sym_matches_type(left, &PyFloat_Type) && sym_matches_type(right, &PyFloat_Type))
         {
-            assert(PyFloat_CheckExact(sym_get_const(left)));
-            assert(PyFloat_CheckExact(sym_get_const(right)));
+            assert(PyFloat_CheckExact(sym_get_const(ctx, left)));
+            assert(PyFloat_CheckExact(sym_get_const(ctx, right)));
             PyObject *temp = PyFloat_FromDouble(
-                PyFloat_AS_DOUBLE(sym_get_const(left)) *
-                PyFloat_AS_DOUBLE(sym_get_const(right)));
+                PyFloat_AS_DOUBLE(sym_get_const(ctx, left)) *
+                PyFloat_AS_DOUBLE(sym_get_const(ctx, right)));
             if (temp == NULL) {
                 goto error;
             }
@@ -352,9 +353,9 @@ dummy_func(void) {
     }
 
     op(_BINARY_OP_ADD_UNICODE, (left, right -- res)) {
-        if (sym_is_const(left) && sym_is_const(right) &&
+        if (sym_is_const(ctx, left) && sym_is_const(ctx, right) &&
             sym_matches_type(left, &PyUnicode_Type) && sym_matches_type(right, &PyUnicode_Type)) {
-            PyObject *temp = PyUnicode_Concat(sym_get_const(left), sym_get_const(right));
+            PyObject *temp = PyUnicode_Concat(sym_get_const(ctx, left), sym_get_const(ctx, right));
             if (temp == NULL) {
                 goto error;
             }
@@ -368,9 +369,9 @@ dummy_func(void) {
 
     op(_BINARY_OP_INPLACE_ADD_UNICODE, (left, right -- )) {
         JitOptSymbol *res;
-        if (sym_is_const(left) && sym_is_const(right) &&
+        if (sym_is_const(ctx, left) && sym_is_const(ctx, right) &&
             sym_matches_type(left, &PyUnicode_Type) && sym_matches_type(right, &PyUnicode_Type)) {
-            PyObject *temp = PyUnicode_Concat(sym_get_const(left), sym_get_const(right));
+            PyObject *temp = PyUnicode_Concat(sym_get_const(ctx, left), sym_get_const(ctx, right));
             if (temp == NULL) {
                 goto error;
             }
@@ -391,14 +392,14 @@ dummy_func(void) {
 
     op(_TO_BOOL, (value -- res)) {
         if (!optimize_to_bool(this_instr, ctx, value, &res)) {
-            res = sym_new_type(ctx, &PyBool_Type);
+            res = sym_new_truthiness(ctx, value, true);
         }
     }
 
     op(_TO_BOOL_BOOL, (value -- res)) {
         if (!optimize_to_bool(this_instr, ctx, value, &res)) {
             sym_set_type(value, &PyBool_Type);
-            res = value;
+            res = sym_new_truthiness(ctx, value, true);
         }
     }
 
@@ -430,6 +431,11 @@ dummy_func(void) {
         }
     }
 
+    op(_UNARY_NOT, (value -- res)) {
+        sym_set_type(value, &PyBool_Type);
+        res = sym_new_truthiness(ctx, value, false);
+    }
+
     op(_COMPARE_OP, (left, right -- res)) {
         if (oparg & 16) {
             res = sym_new_type(ctx, &PyBool_Type);
@@ -521,8 +527,8 @@ dummy_func(void) {
         (void)dict_version;
         (void)index;
         attr = NULL;
-        if (sym_is_const(owner)) {
-            PyModuleObject *mod = (PyModuleObject *)sym_get_const(owner);
+        if (sym_is_const(ctx, owner)) {
+            PyModuleObject *mod = (PyModuleObject *)sym_get_const(ctx, owner);
             if (PyModule_CheckExact(mod)) {
                 PyObject *dict = mod->md_dict;
                 uint64_t watched_mutations = get_mutations(dict);
@@ -599,19 +605,19 @@ dummy_func(void) {
     }
 
     op(_CHECK_FUNCTION_VERSION, (func_version/2, callable, self_or_null, unused[oparg] -- callable, self_or_null, unused[oparg])) {
-        if (sym_is_const(callable) && sym_matches_type(callable, &PyFunction_Type)) {
-            assert(PyFunction_Check(sym_get_const(callable)));
+        if (sym_is_const(ctx, callable) && sym_matches_type(callable, &PyFunction_Type)) {
+            assert(PyFunction_Check(sym_get_const(ctx, callable)));
             REPLACE_OP(this_instr, _CHECK_FUNCTION_VERSION_INLINE, 0, func_version);
-            this_instr->operand1 = (uintptr_t)sym_get_const(callable);
+            this_instr->operand1 = (uintptr_t)sym_get_const(ctx, callable);
         }
         sym_set_type(callable, &PyFunction_Type);
     }
 
     op(_CHECK_FUNCTION_EXACT_ARGS, (callable, self_or_null, unused[oparg] -- callable, self_or_null, unused[oparg])) {
         assert(sym_matches_type(callable, &PyFunction_Type));
-        if (sym_is_const(callable)) {
+        if (sym_is_const(ctx, callable)) {
             if (sym_is_null(self_or_null) || sym_is_not_null(self_or_null)) {
-                PyFunctionObject *func = (PyFunctionObject *)sym_get_const(callable);
+                PyFunctionObject *func = (PyFunctionObject *)sym_get_const(ctx, callable);
                 PyCodeObject *co = (PyCodeObject *)func->func_code;
                 if (co->co_argcount == oparg + !sym_is_null(self_or_null)) {
                     REPLACE_OP(this_instr, _NOP, 0 ,0);
@@ -812,24 +818,26 @@ dummy_func(void) {
     }
 
     op(_GUARD_IS_TRUE_POP, (flag -- )) {
-        if (sym_is_const(flag)) {
-            PyObject *value = sym_get_const(flag);
+        if (sym_is_const(ctx, flag)) {
+            PyObject *value = sym_get_const(ctx, flag);
             assert(value != NULL);
             eliminate_pop_guard(this_instr, value != Py_True);
         }
+        sym_set_const(flag, Py_True);
     }
 
     op(_GUARD_IS_FALSE_POP, (flag -- )) {
-        if (sym_is_const(flag)) {
-            PyObject *value = sym_get_const(flag);
+        if (sym_is_const(ctx, flag)) {
+            PyObject *value = sym_get_const(ctx, flag);
             assert(value != NULL);
             eliminate_pop_guard(this_instr, value != Py_False);
         }
+        sym_set_const(flag, Py_False);
     }
 
     op(_GUARD_IS_NONE_POP, (flag -- )) {
-        if (sym_is_const(flag)) {
-            PyObject *value = sym_get_const(flag);
+        if (sym_is_const(ctx, flag)) {
+            PyObject *value = sym_get_const(ctx, flag);
             assert(value != NULL);
             eliminate_pop_guard(this_instr, !Py_IsNone(value));
         }
@@ -837,11 +845,12 @@ dummy_func(void) {
             assert(!sym_matches_type(flag, &_PyNone_Type));
             eliminate_pop_guard(this_instr, true);
         }
+        sym_set_const(flag, Py_None);
     }
 
     op(_GUARD_IS_NOT_NONE_POP, (flag -- )) {
-        if (sym_is_const(flag)) {
-            PyObject *value = sym_get_const(flag);
+        if (sym_is_const(ctx, flag)) {
+            PyObject *value = sym_get_const(ctx, flag);
             assert(value != NULL);
             eliminate_pop_guard(this_instr, Py_IsNone(value));
         }