gh-115685: Type/values propagate for TO_BOOL in tier 2 (GH-115686)

6 files changed, 163 insertions, 30 deletions

diff --git a/Include/internal/pycore_uop_ids.h b/Include/internal/pycore_uop_ids.h
index 3c133d9..8f71eab 100644
--- a/Include/internal/pycore_uop_ids.h
+++ b/Include/internal/pycore_uop_ids.h
@@ -216,42 +216,43 @@ extern "C" {
 #define _POP_JUMP_IF_FALSE 402
 #define _POP_JUMP_IF_TRUE 403
 #define _POP_TOP POP_TOP
+#define _POP_TOP_LOAD_CONST_INLINE_BORROW 404
 #define _PUSH_EXC_INFO PUSH_EXC_INFO
-#define _PUSH_FRAME 404
+#define _PUSH_FRAME 405
 #define _PUSH_NULL PUSH_NULL
 #define _RESUME_CHECK RESUME_CHECK
-#define _SAVE_RETURN_OFFSET 405
-#define _SEND 406
+#define _SAVE_RETURN_OFFSET 406
+#define _SEND 407
 #define _SEND_GEN SEND_GEN
 #define _SETUP_ANNOTATIONS SETUP_ANNOTATIONS
 #define _SET_ADD SET_ADD
 #define _SET_FUNCTION_ATTRIBUTE SET_FUNCTION_ATTRIBUTE
 #define _SET_UPDATE SET_UPDATE
-#define _START_EXECUTOR 407
-#define _STORE_ATTR 408
-#define _STORE_ATTR_INSTANCE_VALUE 409
-#define _STORE_ATTR_SLOT 410
+#define _START_EXECUTOR 408
+#define _STORE_ATTR 409
+#define _STORE_ATTR_INSTANCE_VALUE 410
+#define _STORE_ATTR_SLOT 411
 #define _STORE_ATTR_WITH_HINT STORE_ATTR_WITH_HINT
 #define _STORE_DEREF STORE_DEREF
-#define _STORE_FAST 411
-#define _STORE_FAST_0 412
-#define _STORE_FAST_1 413
-#define _STORE_FAST_2 414
-#define _STORE_FAST_3 415
-#define _STORE_FAST_4 416
-#define _STORE_FAST_5 417
-#define _STORE_FAST_6 418
-#define _STORE_FAST_7 419
+#define _STORE_FAST 412
+#define _STORE_FAST_0 413
+#define _STORE_FAST_1 414
+#define _STORE_FAST_2 415
+#define _STORE_FAST_3 416
+#define _STORE_FAST_4 417
+#define _STORE_FAST_5 418
+#define _STORE_FAST_6 419
+#define _STORE_FAST_7 420
 #define _STORE_FAST_LOAD_FAST STORE_FAST_LOAD_FAST
 #define _STORE_FAST_STORE_FAST STORE_FAST_STORE_FAST
 #define _STORE_GLOBAL STORE_GLOBAL
 #define _STORE_NAME STORE_NAME
 #define _STORE_SLICE STORE_SLICE
-#define _STORE_SUBSCR 420
+#define _STORE_SUBSCR 421
 #define _STORE_SUBSCR_DICT STORE_SUBSCR_DICT
 #define _STORE_SUBSCR_LIST_INT STORE_SUBSCR_LIST_INT
 #define _SWAP SWAP
-#define _TO_BOOL 421
+#define _TO_BOOL 422
 #define _TO_BOOL_ALWAYS_TRUE TO_BOOL_ALWAYS_TRUE
 #define _TO_BOOL_BOOL TO_BOOL_BOOL
 #define _TO_BOOL_INT TO_BOOL_INT
@@ -262,12 +263,12 @@ extern "C" {
 #define _UNARY_NEGATIVE UNARY_NEGATIVE
 #define _UNARY_NOT UNARY_NOT
 #define _UNPACK_EX UNPACK_EX
-#define _UNPACK_SEQUENCE 422
+#define _UNPACK_SEQUENCE 423
 #define _UNPACK_SEQUENCE_LIST UNPACK_SEQUENCE_LIST
 #define _UNPACK_SEQUENCE_TUPLE UNPACK_SEQUENCE_TUPLE
 #define _UNPACK_SEQUENCE_TWO_TUPLE UNPACK_SEQUENCE_TWO_TUPLE
 #define _WITH_EXCEPT_START WITH_EXCEPT_START
-#define MAX_UOP_ID 422
+#define MAX_UOP_ID 423
 
 #ifdef __cplusplus
 }
diff --git a/Include/internal/pycore_uop_metadata.h b/Include/internal/pycore_uop_metadata.h
index 35340fe..7f921a6 100644
--- a/Include/internal/pycore_uop_metadata.h
+++ b/Include/internal/pycore_uop_metadata.h
@@ -232,6 +232,7 @@ const uint16_t _PyUop_Flags[MAX_UOP_ID+1] = {
     [_CHECK_VALIDITY] = HAS_DEOPT_FLAG,
     [_LOAD_CONST_INLINE] = HAS_PURE_FLAG,
     [_LOAD_CONST_INLINE_BORROW] = HAS_PURE_FLAG,
+    [_POP_TOP_LOAD_CONST_INLINE_BORROW] = HAS_PURE_FLAG,
     [_LOAD_CONST_INLINE_WITH_NULL] = HAS_PURE_FLAG,
     [_LOAD_CONST_INLINE_BORROW_WITH_NULL] = HAS_PURE_FLAG,
     [_CHECK_GLOBALS] = HAS_DEOPT_FLAG,
@@ -425,6 +426,7 @@ const char *const _PyOpcode_uop_name[MAX_UOP_ID+1] = {
     [_POP_EXCEPT] = "_POP_EXCEPT",
     [_POP_FRAME] = "_POP_FRAME",
     [_POP_TOP] = "_POP_TOP",
+    [_POP_TOP_LOAD_CONST_INLINE_BORROW] = "_POP_TOP_LOAD_CONST_INLINE_BORROW",
     [_PUSH_EXC_INFO] = "_PUSH_EXC_INFO",
     [_PUSH_FRAME] = "_PUSH_FRAME",
     [_PUSH_NULL] = "_PUSH_NULL",
diff --git a/Python/bytecodes.c b/Python/bytecodes.c
index 1d51509..095982d 100644
--- a/Python/bytecodes.c
+++ b/Python/bytecodes.c
@@ -4043,6 +4043,11 @@ dummy_func(
             value = ptr;
         }
 
+        tier2 pure op (_POP_TOP_LOAD_CONST_INLINE_BORROW, (ptr/4, pop -- value)) {
+            Py_DECREF(pop);
+            value = ptr;
+        }
+
         tier2 pure op(_LOAD_CONST_INLINE_WITH_NULL, (ptr/4 -- value, null)) {
             value = Py_NewRef(ptr);
             null = NULL;
diff --git a/Python/executor_cases.c.h b/Python/executor_cases.c.h
index 9ec1be9..e84c332 100644
--- a/Python/executor_cases.c.h
+++ b/Python/executor_cases.c.h
@@ -3759,6 +3759,17 @@
             break;
         }
 
+        case _POP_TOP_LOAD_CONST_INLINE_BORROW: {
+            PyObject *pop;
+            PyObject *value;
+            pop = stack_pointer[-1];
+            PyObject *ptr = (PyObject *)CURRENT_OPERAND();
+            Py_DECREF(pop);
+            value = ptr;
+            stack_pointer[-1] = value;
+            break;
+        }
+
         case _LOAD_CONST_INLINE_WITH_NULL: {
             PyObject *value;
             PyObject *null;
diff --git a/Python/optimizer_bytecodes.c b/Python/optimizer_bytecodes.c
index aa19a50..2b47381 100644
--- a/Python/optimizer_bytecodes.c
+++ b/Python/optimizer_bytecodes.c
@@ -254,6 +254,67 @@ dummy_func(void) {
         }
     }
 
+    op(_TO_BOOL, (value -- res)) {
+        (void)value;
+        res = sym_new_type(ctx, &PyBool_Type);
+        OUT_OF_SPACE_IF_NULL(res);
+    }
+
+    op(_TO_BOOL_BOOL, (value -- value)) {
+        if (sym_matches_type(value, &PyBool_Type)) {
+            REPLACE_OP(this_instr, _NOP, 0, 0);
+        }
+        else {
+            if(!sym_set_type(value, &PyBool_Type)) {
+                goto hit_bottom;
+            }
+        }
+    }
+
+    op(_TO_BOOL_INT, (value -- res)) {
+        if (sym_is_const(value) && sym_matches_type(value, &PyLong_Type)) {
+            PyObject *load = _PyLong_IsZero((PyLongObject *)sym_get_const(value))
+                             ? Py_False : Py_True;
+            REPLACE_OP(this_instr, _POP_TOP_LOAD_CONST_INLINE_BORROW, 0, (uintptr_t)load);
+            OUT_OF_SPACE_IF_NULL(res = sym_new_const(ctx, load));
+        }
+        else {
+            OUT_OF_SPACE_IF_NULL(res = sym_new_type(ctx, &PyBool_Type));
+        }
+        if(!sym_set_type(value, &PyLong_Type)) {
+            goto hit_bottom;
+        }
+    }
+
+    op(_TO_BOOL_LIST, (value -- res)) {
+        if(!sym_set_type(value, &PyList_Type)) {
+            goto hit_bottom;
+        }
+        OUT_OF_SPACE_IF_NULL(res = sym_new_type(ctx, &PyBool_Type));
+    }
+
+    op(_TO_BOOL_NONE, (value -- res)) {
+        if (sym_get_const(value) == Py_None) {
+            REPLACE_OP(this_instr, _POP_TOP_LOAD_CONST_INLINE_BORROW, 0, (uintptr_t)Py_False);
+        }
+        sym_set_const(value, Py_None);
+        OUT_OF_SPACE_IF_NULL(res = sym_new_const(ctx, Py_False));
+    }
+
+    op(_TO_BOOL_STR, (value -- res)) {
+        if (sym_is_const(value) && sym_matches_type(value, &PyUnicode_Type)) {
+            PyObject *load = sym_get_const(value) == &_Py_STR(empty) ? Py_False : Py_True;
+            REPLACE_OP(this_instr, _POP_TOP_LOAD_CONST_INLINE_BORROW, 0, (uintptr_t)load);
+            OUT_OF_SPACE_IF_NULL(res = sym_new_const(ctx, load));
+        }
+        else {
+            OUT_OF_SPACE_IF_NULL(res = sym_new_type(ctx, &PyBool_Type));
+        }
+        if(!sym_set_type(value, &PyUnicode_Type)) {
+            goto hit_bottom;
+        }
+    }
+
     op(_LOAD_CONST, (-- value)) {
         // There should be no LOAD_CONST. It should be all
         // replaced by peephole_opt.
diff --git a/Python/optimizer_cases.c.h b/Python/optimizer_cases.c.h
index b34c573..9d7ebb8 100644
--- a/Python/optimizer_cases.c.h
+++ b/Python/optimizer_cases.c.h
@@ -104,45 +104,90 @@
         }
 
         case _TO_BOOL: {
+            _Py_UopsSymbol *value;
             _Py_UopsSymbol *res;
-            res = sym_new_unknown(ctx);
-            if (res == NULL) goto out_of_space;
+            value = stack_pointer[-1];
+            (void)value;
+            res = sym_new_type(ctx, &PyBool_Type);
+            OUT_OF_SPACE_IF_NULL(res);
             stack_pointer[-1] = res;
             break;
         }
 
         case _TO_BOOL_BOOL: {
+            _Py_UopsSymbol *value;
+            value = stack_pointer[-1];
+            if (sym_matches_type(value, &PyBool_Type)) {
+                REPLACE_OP(this_instr, _NOP, 0, 0);
+            }
+            else {
+                if(!sym_set_type(value, &PyBool_Type)) {
+                    goto hit_bottom;
+                }
+            }
             break;
         }
 
         case _TO_BOOL_INT: {
+            _Py_UopsSymbol *value;
             _Py_UopsSymbol *res;
-            res = sym_new_unknown(ctx);
-            if (res == NULL) goto out_of_space;
+            value = stack_pointer[-1];
+            if (sym_is_const(value) && sym_matches_type(value, &PyLong_Type)) {
+                PyObject *load = _PyLong_IsZero((PyLongObject *)sym_get_const(value))
+                ? Py_False : Py_True;
+                REPLACE_OP(this_instr, _POP_TOP_LOAD_CONST_INLINE_BORROW, 0, (uintptr_t)load);
+                OUT_OF_SPACE_IF_NULL(res = sym_new_const(ctx, load));
+            }
+            else {
+                OUT_OF_SPACE_IF_NULL(res = sym_new_type(ctx, &PyBool_Type));
+            }
+            if(!sym_set_type(value, &PyLong_Type)) {
+                goto hit_bottom;
+            }
             stack_pointer[-1] = res;
             break;
         }
 
         case _TO_BOOL_LIST: {
+            _Py_UopsSymbol *value;
             _Py_UopsSymbol *res;
-            res = sym_new_unknown(ctx);
-            if (res == NULL) goto out_of_space;
+            value = stack_pointer[-1];
+            if(!sym_set_type(value, &PyList_Type)) {
+                goto hit_bottom;
+            }
+            OUT_OF_SPACE_IF_NULL(res = sym_new_type(ctx, &PyBool_Type));
             stack_pointer[-1] = res;
             break;
         }
 
         case _TO_BOOL_NONE: {
+            _Py_UopsSymbol *value;
             _Py_UopsSymbol *res;
-            res = sym_new_unknown(ctx);
-            if (res == NULL) goto out_of_space;
+            value = stack_pointer[-1];
+            if (sym_get_const(value) == Py_None) {
+                REPLACE_OP(this_instr, _POP_TOP_LOAD_CONST_INLINE_BORROW, 0, (uintptr_t)Py_False);
+            }
+            sym_set_const(value, Py_None);
+            OUT_OF_SPACE_IF_NULL(res = sym_new_const(ctx, Py_False));
             stack_pointer[-1] = res;
             break;
         }
 
         case _TO_BOOL_STR: {
+            _Py_UopsSymbol *value;
             _Py_UopsSymbol *res;
-            res = sym_new_unknown(ctx);
-            if (res == NULL) goto out_of_space;
+            value = stack_pointer[-1];
+            if (sym_is_const(value) && sym_matches_type(value, &PyUnicode_Type)) {
+                PyObject *load = sym_get_const(value) == &_Py_STR(empty) ? Py_False : Py_True;
+                REPLACE_OP(this_instr, _POP_TOP_LOAD_CONST_INLINE_BORROW, 0, (uintptr_t)load);
+                OUT_OF_SPACE_IF_NULL(res = sym_new_const(ctx, load));
+            }
+            else {
+                OUT_OF_SPACE_IF_NULL(res = sym_new_type(ctx, &PyBool_Type));
+            }
+            if(!sym_set_type(value, &PyUnicode_Type)) {
+                goto hit_bottom;
+            }
             stack_pointer[-1] = res;
             break;
         }
@@ -1789,6 +1834,14 @@
             break;
         }
 
+        case _POP_TOP_LOAD_CONST_INLINE_BORROW: {
+            _Py_UopsSymbol *value;
+            value = sym_new_unknown(ctx);
+            if (value == NULL) goto out_of_space;
+            stack_pointer[-1] = value;
+            break;
+        }
+
         case _LOAD_CONST_INLINE_WITH_NULL: {
             _Py_UopsSymbol *value;
             _Py_UopsSymbol *null;