gh-93143: Avoid NULL check in LOAD_FAST based on analysis in the compiler (GH-93144)

3 files changed, 130 insertions, 24 deletions

diff --git a/Python/ceval.c b/Python/ceval.c
index 3654692..ec86c70 100644
--- a/Python/ceval.c
+++ b/Python/ceval.c
@@ -1813,7 +1813,7 @@ handle_eval_breaker:
             DISPATCH();
         }
 
-        TARGET(LOAD_FAST) {
+        TARGET(LOAD_FAST_CHECK) {
             PyObject *value = GETLOCAL(oparg);
             if (value == NULL) {
                 goto unbound_local_error;
@@ -1823,6 +1823,14 @@ handle_eval_breaker:
             DISPATCH();
         }
 
+        TARGET(LOAD_FAST) {
+            PyObject *value = GETLOCAL(oparg);
+            assert(value != NULL);
+            Py_INCREF(value);
+            PUSH(value);
+            DISPATCH();
+        }
+
         TARGET(LOAD_CONST) {
             PREDICTED(LOAD_CONST);
             PyObject *value = GETITEM(consts, oparg);
@@ -1840,17 +1848,13 @@ handle_eval_breaker:
 
         TARGET(LOAD_FAST__LOAD_FAST) {
             PyObject *value = GETLOCAL(oparg);
-            if (value == NULL) {
-                goto unbound_local_error;
-            }
+            assert(value != NULL);
             NEXTOPARG();
             next_instr++;
             Py_INCREF(value);
             PUSH(value);
             value = GETLOCAL(oparg);
-            if (value == NULL) {
-                goto unbound_local_error;
-            }
+            assert(value != NULL);
             Py_INCREF(value);
             PUSH(value);
             NOTRACE_DISPATCH();
@@ -1858,9 +1862,7 @@ handle_eval_breaker:
 
         TARGET(LOAD_FAST__LOAD_CONST) {
             PyObject *value = GETLOCAL(oparg);
-            if (value == NULL) {
-                goto unbound_local_error;
-            }
+            assert(value != NULL);
             NEXTOPARG();
             next_instr++;
             Py_INCREF(value);
@@ -1877,9 +1879,7 @@ handle_eval_breaker:
             NEXTOPARG();
             next_instr++;
             value = GETLOCAL(oparg);
-            if (value == NULL) {
-                goto unbound_local_error;
-            }
+            assert(value != NULL);
             Py_INCREF(value);
             PUSH(value);
             NOTRACE_DISPATCH();
@@ -1902,9 +1902,7 @@ handle_eval_breaker:
             Py_INCREF(value);
             PUSH(value);
             value = GETLOCAL(oparg);
-            if (value == NULL) {
-                goto unbound_local_error;
-            }
+            assert(value != NULL);
             Py_INCREF(value);
             PUSH(value);
             NOTRACE_DISPATCH();
diff --git a/Python/compile.c b/Python/compile.c
index 0920ceb..a6788b7 100644
--- a/Python/compile.c
+++ b/Python/compile.c
@@ -1109,6 +1109,7 @@ stack_effect(int opcode, int oparg, int jump)
             return 1;
 
         case LOAD_FAST:
+        case LOAD_FAST_CHECK:
             return 1;
         case STORE_FAST:
             return -1;
@@ -7746,6 +7747,109 @@ assemble_jump_offsets(struct assembler *a, struct compiler *c)
     } while (extended_arg_recompile);
 }
 
+
+// Ensure each basicblock is only put onto the stack once.
+#define MAYBE_PUSH(B) do {                          \
+        if ((B)->b_visited == 0) {                  \
+            *(*stack_top)++ = (B);                  \
+            (B)->b_visited = 1;                     \
+        }                                           \
+    } while (0)
+
+static void
+scan_block_for_local(int target, basicblock *b, bool unsafe_to_start,
+                     basicblock ***stack_top)
+{
+    bool unsafe = unsafe_to_start;
+    for (int i = 0; i < b->b_iused; i++) {
+        struct instr *instr = &b->b_instr[i];
+        assert(instr->i_opcode != EXTENDED_ARG);
+        assert(instr->i_opcode != EXTENDED_ARG_QUICK);
+        assert(instr->i_opcode != LOAD_FAST__LOAD_FAST);
+        assert(instr->i_opcode != STORE_FAST__LOAD_FAST);
+        assert(instr->i_opcode != LOAD_CONST__LOAD_FAST);
+        assert(instr->i_opcode != STORE_FAST__STORE_FAST);
+        assert(instr->i_opcode != LOAD_FAST__LOAD_CONST);
+        if (unsafe && instr->i_except != NULL) {
+            MAYBE_PUSH(instr->i_except);
+        }
+        if (instr->i_oparg != target) {
+            continue;
+        }
+        switch (instr->i_opcode) {
+            case LOAD_FAST_CHECK:
+                // if this doesn't raise, then var is defined
+                unsafe = false;
+                break;
+            case LOAD_FAST:
+                if (unsafe) {
+                    instr->i_opcode = LOAD_FAST_CHECK;
+                }
+                unsafe = false;
+                break;
+            case STORE_FAST:
+                unsafe = false;
+                break;
+            case DELETE_FAST:
+                unsafe = true;
+                break;
+        }
+    }
+    if (unsafe) {
+        // unsafe at end of this block,
+        // so unsafe at start of next blocks
+        if (b->b_next && !b->b_nofallthrough) {
+            MAYBE_PUSH(b->b_next);
+        }
+        if (b->b_iused > 0) {
+            struct instr *last = &b->b_instr[b->b_iused-1];
+            if (is_jump(last)) {
+                assert(last->i_target != NULL);
+                MAYBE_PUSH(last->i_target);
+            }
+        }
+    }
+}
+#undef MAYBE_PUSH
+
+static int
+add_checks_for_loads_of_unknown_variables(struct assembler *a,
+                                          struct compiler *c)
+{
+    basicblock **stack = make_cfg_traversal_stack(a->a_entry);
+    if (stack == NULL) {
+        return -1;
+    }
+    Py_ssize_t nparams = PyList_GET_SIZE(c->u->u_ste->ste_varnames);
+    int nlocals = (int)PyDict_GET_SIZE(c->u->u_varnames);
+    for (int target = 0; target < nlocals; target++) {
+        for (basicblock *b = a->a_entry; b != NULL; b = b->b_next) {
+            b->b_visited = 0;
+        }
+        basicblock **stack_top = stack;
+
+        // First pass: find the relevant DFS starting points:
+        // the places where "being uninitialized" originates,
+        // which are the entry block and any DELETE_FAST statements.
+        if (target >= nparams) {
+            // only non-parameter locals start out uninitialized.
+            *(stack_top++) = a->a_entry;
+            a->a_entry->b_visited = 1;
+        }
+        for (basicblock *b = a->a_entry; b != NULL; b = b->b_next) {
+            scan_block_for_local(target, b, false, &stack_top);
+        }
+
+        // Second pass: Depth-first search to propagate uncertainty
+        while (stack_top > stack) {
+            basicblock *b = *--stack_top;
+            scan_block_for_local(target, b, true, &stack_top);
+        }
+    }
+    PyMem_Free(stack);
+    return 0;
+}
+
 static PyObject *
 dict_keys_inorder(PyObject *dict, Py_ssize_t offset)
 {
@@ -8385,6 +8489,10 @@ assemble(struct compiler *c, int addNone)
     /* Order of basic blocks must have been determined by now */
     normalize_jumps(&a);
 
+    if (add_checks_for_loads_of_unknown_variables(&a, c) < 0) {
+        goto error;
+    }
+
     /* Can't modify the bytecode after computing jump offsets. */
     assemble_jump_offsets(&a, c);
 
diff --git a/Python/opcode_targets.h b/Python/opcode_targets.h
index 71b7a76..1009b3a 100644
--- a/Python/opcode_targets.h
+++ b/Python/opcode_targets.h
@@ -126,7 +126,7 @@ static void *opcode_targets[256] = {
     &&TARGET_LOAD_FAST,
     &&TARGET_STORE_FAST,
     &&TARGET_DELETE_FAST,
-    &&TARGET_LOAD_METHOD_WITH_VALUES,
+    &&TARGET_LOAD_FAST_CHECK,
     &&TARGET_POP_JUMP_FORWARD_IF_NOT_NONE,
     &&TARGET_POP_JUMP_FORWARD_IF_NONE,
     &&TARGET_RAISE_VARARGS,
@@ -140,9 +140,9 @@ static void *opcode_targets[256] = {
     &&TARGET_STORE_DEREF,
     &&TARGET_DELETE_DEREF,
     &&TARGET_JUMP_BACKWARD,
-    &&TARGET_RESUME_QUICK,
+    &&TARGET_LOAD_METHOD_WITH_VALUES,
     &&TARGET_CALL_FUNCTION_EX,
-    &&TARGET_STORE_ATTR_ADAPTIVE,
+    &&TARGET_RESUME_QUICK,
     &&TARGET_EXTENDED_ARG,
     &&TARGET_LIST_APPEND,
     &&TARGET_SET_ADD,
@@ -152,30 +152,31 @@ static void *opcode_targets[256] = {
     &&TARGET_YIELD_VALUE,
     &&TARGET_RESUME,
     &&TARGET_MATCH_CLASS,
+    &&TARGET_STORE_ATTR_ADAPTIVE,
     &&TARGET_STORE_ATTR_INSTANCE_VALUE,
-    &&TARGET_STORE_ATTR_SLOT,
     &&TARGET_FORMAT_VALUE,
     &&TARGET_BUILD_CONST_KEY_MAP,
     &&TARGET_BUILD_STRING,
+    &&TARGET_STORE_ATTR_SLOT,
     &&TARGET_STORE_ATTR_WITH_HINT,
-    &&TARGET_STORE_FAST__LOAD_FAST,
     &&TARGET_LOAD_METHOD,
-    &&TARGET_STORE_FAST__STORE_FAST,
+    &&TARGET_STORE_FAST__LOAD_FAST,
     &&TARGET_LIST_EXTEND,
     &&TARGET_SET_UPDATE,
     &&TARGET_DICT_MERGE,
     &&TARGET_DICT_UPDATE,
+    &&TARGET_STORE_FAST__STORE_FAST,
     &&TARGET_STORE_SUBSCR_ADAPTIVE,
     &&TARGET_STORE_SUBSCR_DICT,
     &&TARGET_STORE_SUBSCR_LIST_INT,
     &&TARGET_UNPACK_SEQUENCE_ADAPTIVE,
-    &&TARGET_UNPACK_SEQUENCE_LIST,
     &&TARGET_CALL,
     &&TARGET_KW_NAMES,
     &&TARGET_POP_JUMP_BACKWARD_IF_NOT_NONE,
     &&TARGET_POP_JUMP_BACKWARD_IF_NONE,
     &&TARGET_POP_JUMP_BACKWARD_IF_FALSE,
     &&TARGET_POP_JUMP_BACKWARD_IF_TRUE,
+    &&TARGET_UNPACK_SEQUENCE_LIST,
     &&TARGET_UNPACK_SEQUENCE_TUPLE,
     &&TARGET_UNPACK_SEQUENCE_TWO_TUPLE,
     &&_unknown_opcode,
@@ -253,6 +254,5 @@ static void *opcode_targets[256] = {
     &&_unknown_opcode,
     &&_unknown_opcode,
     &&_unknown_opcode,
-    &&_unknown_opcode,
     &&TARGET_DO_TRACING
 };