GH-113710: Improve `_SET_IP` and `_CHECK_VALIDITY` (GH-115248)

1 files changed, 51 insertions, 24 deletions

diff --git a/Python/optimizer_analysis.c b/Python/optimizer_analysis.c
index e02ca4d..4997452 100644
--- a/Python/optimizer_analysis.c
+++ b/Python/optimizer_analysis.c
@@ -652,35 +652,62 @@ error:
 static void
 remove_unneeded_uops(_PyUOpInstruction *buffer, int buffer_size)
 {
+    /* Remove _SET_IP and _CHECK_VALIDITY where possible.
+     * _SET_IP is needed if the following instruction escapes or
+     * could error. _CHECK_VALIDITY is needed if the previous
+     * instruction could have escaped. */
     int last_set_ip = -1;
-    bool maybe_invalid = false;
+    bool may_have_escaped = false;
     for (int pc = 0; pc < buffer_size; pc++) {
         int opcode = buffer[pc].opcode;
-        if (opcode == _SET_IP) {
-            buffer[pc].opcode = NOP;
-            last_set_ip = pc;
-        }
-        else if (opcode == _CHECK_VALIDITY) {
-            if (maybe_invalid) {
-                maybe_invalid = false;
-            }
-            else {
+        switch (opcode) {
+            case _SET_IP:
                 buffer[pc].opcode = NOP;
-            }
-        }
-        else if (op_is_end(opcode)) {
-            break;
-        }
-        else {
-            if (_PyUop_Flags[opcode] & HAS_ESCAPES_FLAG) {
-                maybe_invalid = true;
-                if (last_set_ip >= 0) {
-                    buffer[last_set_ip].opcode = _SET_IP;
+                last_set_ip = pc;
+                break;
+            case _CHECK_VALIDITY:
+                if (may_have_escaped) {
+                    may_have_escaped = false;
                 }
-            }
-            if ((_PyUop_Flags[opcode] & HAS_ERROR_FLAG) || opcode == _PUSH_FRAME) {
-                if (last_set_ip >= 0) {
-                    buffer[last_set_ip].opcode = _SET_IP;
+                else {
+                    buffer[pc].opcode = NOP;
+                }
+                break;
+            case _CHECK_VALIDITY_AND_SET_IP:
+                if (may_have_escaped) {
+                    may_have_escaped = false;
+                    buffer[pc].opcode = _CHECK_VALIDITY;
+                }
+                else {
+                    buffer[pc].opcode = NOP;
+                }
+                last_set_ip = pc;
+                break;
+            case _JUMP_TO_TOP:
+            case _EXIT_TRACE:
+                return;
+            default:
+            {
+                bool needs_ip = false;
+                if (_PyUop_Flags[opcode] & HAS_ESCAPES_FLAG) {
+                    needs_ip = true;
+                    may_have_escaped = true;
+                }
+                if (_PyUop_Flags[opcode] & HAS_ERROR_FLAG) {
+                    needs_ip = true;
+                }
+                if (opcode == _PUSH_FRAME) {
+                    needs_ip = true;
+                }
+                if (needs_ip && last_set_ip >= 0) {
+                    if (buffer[last_set_ip].opcode == _CHECK_VALIDITY) {
+                        buffer[last_set_ip].opcode = _CHECK_VALIDITY_AND_SET_IP;
+                    }
+                    else {
+                        assert(buffer[last_set_ip].opcode == _NOP);
+                        buffer[last_set_ip].opcode = _SET_IP;
+                    }
+                    last_set_ip = -1;
                 }
             }
         }