gh-106603: Make uop struct a triple (opcode, oparg, operand) (#106794)

1 files changed, 51 insertions, 40 deletions

diff --git a/Python/optimizer.c b/Python/optimizer.c
index 693ba375971..3d385a1506c 100644
--- a/Python/optimizer.c
+++ b/Python/optimizer.c
@@ -344,13 +344,19 @@ uop_item(_PyUOpExecutorObject *self, Py_ssize_t index)
     if (oname == NULL) {
         return NULL;
     }
+    PyObject *oparg = PyLong_FromUnsignedLong(self->trace[index].oparg);
+    if (oparg == NULL) {
+        Py_DECREF(oname);
+        return NULL;
+    }
     PyObject *operand = PyLong_FromUnsignedLongLong(self->trace[index].operand);
     if (operand == NULL) {
+        Py_DECREF(oparg);
         Py_DECREF(oname);
         return NULL;
     }
-    PyObject *args[2] = { oname, operand };
-    return _PyTuple_FromArraySteal(args, 2);
+    PyObject *args[3] = { oname, oparg, operand };
+    return _PyTuple_FromArraySteal(args, 3);
 }
 
 PySequenceMethods uop_as_sequence = {
@@ -395,29 +401,33 @@ translate_bytecode_to_trace(
 #define DPRINTF(level, ...)
 #endif
 
-#define ADD_TO_TRACE(OPCODE, OPERAND) \
+#define ADD_TO_TRACE(OPCODE, OPARG, OPERAND) \
     DPRINTF(2, \
-            "  ADD_TO_TRACE(%s, %" PRIu64 ")\n", \
+            "  ADD_TO_TRACE(%s, %d, %" PRIu64 ")\n", \
             uop_name(OPCODE), \
+            (OPARG), \
             (uint64_t)(OPERAND)); \
     assert(trace_length < max_length); \
     assert(reserved > 0); \
     reserved--; \
     trace[trace_length].opcode = (OPCODE); \
+    trace[trace_length].oparg = (OPARG); \
     trace[trace_length].operand = (OPERAND); \
     trace_length++;
 
 #define INSTR_IP(INSTR, CODE) \
-    ((long)((INSTR) - ((_Py_CODEUNIT *)(CODE)->co_code_adaptive)))
+    ((uint32_t)((INSTR) - ((_Py_CODEUNIT *)(CODE)->co_code_adaptive)))
 
-#define ADD_TO_STUB(INDEX, OPCODE, OPERAND) \
-    DPRINTF(2, "    ADD_TO_STUB(%d, %s, %" PRIu64 ")\n", \
+#define ADD_TO_STUB(INDEX, OPCODE, OPARG, OPERAND) \
+    DPRINTF(2, "    ADD_TO_STUB(%d, %s, %d, %" PRIu64 ")\n", \
             (INDEX), \
             uop_name(OPCODE), \
+            (OPARG), \
             (uint64_t)(OPERAND)); \
     assert(reserved > 0); \
     reserved--; \
     trace[(INDEX)].opcode = (OPCODE); \
+    trace[(INDEX)].oparg = (OPARG); \
     trace[(INDEX)].operand = (OPERAND);
 
 // Reserve space for n uops
@@ -433,7 +443,7 @@ translate_bytecode_to_trace(
 #define RESERVE(main, stub) RESERVE_RAW((main) + (stub) + 2, uop_name(opcode))
 
     DPRINTF(4,
-            "Optimizing %s (%s:%d) at byte offset %ld\n",
+            "Optimizing %s (%s:%d) at byte offset %d\n",
             PyUnicode_AsUTF8(code->co_qualname),
             PyUnicode_AsUTF8(code->co_filename),
             code->co_firstlineno,
@@ -441,11 +451,11 @@ translate_bytecode_to_trace(
 
     for (;;) {
         RESERVE_RAW(2, "epilogue");  // Always need space for SAVE_IP and EXIT_TRACE
-        ADD_TO_TRACE(SAVE_IP, INSTR_IP(instr, code));
+        ADD_TO_TRACE(SAVE_IP, INSTR_IP(instr, code), 0);
 
-        int opcode = instr->op.code;
-        int oparg = instr->op.arg;
-        int extras = 0;
+        uint32_t opcode = instr->op.code;
+        uint32_t oparg = instr->op.arg;
+        uint32_t extras = 0;
 
         while (opcode == EXTENDED_ARG) {
             instr++;
@@ -467,7 +477,7 @@ translate_bytecode_to_trace(
             case POP_JUMP_IF_NONE:
             {
                 RESERVE(2, 2);
-                ADD_TO_TRACE(IS_NONE, 0);
+                ADD_TO_TRACE(IS_NONE, 0, 0);
                 opcode = POP_JUMP_IF_TRUE;
                 goto pop_jump_if_bool;
             }
@@ -475,7 +485,7 @@ translate_bytecode_to_trace(
             case POP_JUMP_IF_NOT_NONE:
             {
                 RESERVE(2, 2);
-                ADD_TO_TRACE(IS_NONE, 0);
+                ADD_TO_TRACE(IS_NONE, 0, 0);
                 opcode = POP_JUMP_IF_FALSE;
                 goto pop_jump_if_bool;
             }
@@ -489,11 +499,11 @@ pop_jump_if_bool:
                 _Py_CODEUNIT *target_instr =
                     instr + 1 + _PyOpcode_Caches[_PyOpcode_Deopt[opcode]] + oparg;
                 max_length -= 2;  // Really the start of the stubs
-                int uopcode = opcode == POP_JUMP_IF_TRUE ?
+                uint32_t uopcode = opcode == POP_JUMP_IF_TRUE ?
                     _POP_JUMP_IF_TRUE : _POP_JUMP_IF_FALSE;
-                ADD_TO_TRACE(uopcode, max_length);
-                ADD_TO_STUB(max_length, SAVE_IP, INSTR_IP(target_instr, code));
-                ADD_TO_STUB(max_length + 1, EXIT_TRACE, 0);
+                ADD_TO_TRACE(uopcode, max_length, 0);
+                ADD_TO_STUB(max_length, SAVE_IP, INSTR_IP(target_instr, code), 0);
+                ADD_TO_STUB(max_length + 1, EXIT_TRACE, 0, 0);
                 break;
             }
 
@@ -501,7 +511,7 @@ pop_jump_if_bool:
             {
                 if (instr + 2 - oparg == initial_instr) {
                     RESERVE(1, 0);
-                    ADD_TO_TRACE(JUMP_TO_TOP, 0);
+                    ADD_TO_TRACE(JUMP_TO_TOP, 0, 0);
                 }
                 else {
                     DPRINTF(2, "JUMP_BACKWARD not to top ends trace\n");
@@ -546,14 +556,14 @@ pop_jump_if_bool:
                 _Py_CODEUNIT *target_instr =  // +1 at the end skips over END_FOR
                     instr + 1 + _PyOpcode_Caches[_PyOpcode_Deopt[opcode]] + oparg + 1;
                 max_length -= 3;  // Really the start of the stubs
-                ADD_TO_TRACE(check_op, 0);
-                ADD_TO_TRACE(exhausted_op, 0);
-                ADD_TO_TRACE(_POP_JUMP_IF_TRUE, max_length);
-                ADD_TO_TRACE(next_op, 0);
-
-                ADD_TO_STUB(max_length + 0, POP_TOP, 0);
-                ADD_TO_STUB(max_length + 1, SAVE_IP, INSTR_IP(target_instr, code));
-                ADD_TO_STUB(max_length + 2, EXIT_TRACE, 0);
+                ADD_TO_TRACE(check_op, 0, 0);
+                ADD_TO_TRACE(exhausted_op, 0, 0);
+                ADD_TO_TRACE(_POP_JUMP_IF_TRUE, max_length, 0);
+                ADD_TO_TRACE(next_op, 0, 0);
+
+                ADD_TO_STUB(max_length + 0, POP_TOP, 0, 0);
+                ADD_TO_STUB(max_length + 1, SAVE_IP, INSTR_IP(target_instr, code), 0);
+                ADD_TO_STUB(max_length + 2, EXIT_TRACE, 0, 0);
                 break;
             }
 
@@ -564,19 +574,20 @@ pop_jump_if_bool:
                     // Reserve space for nuops (+ SAVE_IP + EXIT_TRACE)
                     int nuops = expansion->nuops;
                     RESERVE(nuops, 0);
+                    uint32_t orig_oparg = oparg;  // For OPARG_TOP/BOTTOM
                     for (int i = 0; i < nuops; i++) {
-                        uint64_t operand;
+                        oparg = orig_oparg;
+                        uint64_t operand = 0;
                         int offset = expansion->uops[i].offset;
                         switch (expansion->uops[i].size) {
                             case OPARG_FULL:
-                                operand = oparg;
                                 if (extras && OPCODE_HAS_JUMP(opcode)) {
                                     if (opcode == JUMP_BACKWARD_NO_INTERRUPT) {
-                                        operand -= extras;
+                                        oparg -= extras;
                                     }
                                     else {
                                         assert(opcode != JUMP_BACKWARD);
-                                        operand += extras;
+                                        oparg += extras;
                                     }
                                 }
                                 break;
@@ -590,10 +601,10 @@ pop_jump_if_bool:
                                 operand = read_u64(&instr[offset].cache);
                                 break;
                             case OPARG_TOP:  // First half of super-instr
-                                operand = oparg >> 4;
+                                oparg = orig_oparg >> 4;
                                 break;
                             case OPARG_BOTTOM:  // Second half of super-instr
-                                operand = oparg & 0xF;
+                                oparg = orig_oparg & 0xF;
                                 break;
                             default:
                                 fprintf(stderr,
@@ -603,7 +614,7 @@ pop_jump_if_bool:
                                         expansion->uops[i].offset);
                                 Py_FatalError("garbled expansion");
                         }
-                        ADD_TO_TRACE(expansion->uops[i].uop, operand);
+                        ADD_TO_TRACE(expansion->uops[i].uop, oparg, operand);
                     }
                     break;
                 }
@@ -621,9 +632,9 @@ pop_jump_if_bool:
 done:
     // Skip short traces like SAVE_IP, LOAD_FAST, SAVE_IP, EXIT_TRACE
     if (trace_length > 3) {
-        ADD_TO_TRACE(EXIT_TRACE, 0);
+        ADD_TO_TRACE(EXIT_TRACE, 0, 0);
         DPRINTF(1,
-                "Created a trace for %s (%s:%d) at byte offset %ld -- length %d\n",
+                "Created a trace for %s (%s:%d) at byte offset %d -- length %d\n",
                 PyUnicode_AsUTF8(code->co_qualname),
                 PyUnicode_AsUTF8(code->co_filename),
                 code->co_firstlineno,
@@ -644,10 +655,10 @@ done:
                 if (trace[i].opcode == _POP_JUMP_IF_FALSE ||
                     trace[i].opcode == _POP_JUMP_IF_TRUE)
                 {
-                    uint64_t target = trace[i].operand;
-                    if (target >= (uint64_t)max_length) {
+                    int target = trace[i].oparg;
+                    if (target >= max_length) {
                         target += trace_length - max_length;
-                        trace[i].operand = target;
+                        trace[i].oparg = target;
                     }
                 }
             }
@@ -657,7 +668,7 @@ done:
     }
     else {
         DPRINTF(4,
-                "No trace for %s (%s:%d) at byte offset %ld\n",
+                "No trace for %s (%s:%d) at byte offset %d\n",
                 PyUnicode_AsUTF8(code->co_qualname),
                 PyUnicode_AsUTF8(code->co_filename),
                 code->co_firstlineno,