GH-121459: Streamline PyObject* to PyStackRef conversions by disallowing NULL pointers. (GH-124894)

4 files changed, 63 insertions, 57 deletions

diff --git a/Include/internal/pycore_stackref.h b/Include/internal/pycore_stackref.h
index b5b6993..cf6dd22 100644
--- a/Include/internal/pycore_stackref.h
+++ b/Include/internal/pycore_stackref.h
@@ -111,7 +111,8 @@ PyStackRef_AsPyObjectBorrow(_PyStackRef stackref)
 static inline PyObject *
 PyStackRef_AsPyObjectSteal(_PyStackRef stackref)
 {
-    if (!PyStackRef_IsNull(stackref) && PyStackRef_IsDeferred(stackref)) {
+    assert(!PyStackRef_IsNull(stackref));
+    if (PyStackRef_IsDeferred(stackref)) {
         return Py_NewRef(PyStackRef_AsPyObjectBorrow(stackref));
     }
     return PyStackRef_AsPyObjectBorrow(stackref);
@@ -131,9 +132,10 @@ PyStackRef_AsPyObjectSteal(_PyStackRef stackref)
 static inline _PyStackRef
 _PyStackRef_FromPyObjectSteal(PyObject *obj)
 {
+    assert(obj != NULL);
     // Make sure we don't take an already tagged value.
     assert(((uintptr_t)obj & Py_TAG_BITS) == 0);
-    unsigned int tag = (obj == NULL || _Py_IsImmortal(obj)) ? (Py_TAG_DEFERRED) : Py_TAG_PTR;
+    unsigned int tag = _Py_IsImmortal(obj) ? (Py_TAG_DEFERRED) : Py_TAG_PTR;
     return ((_PyStackRef){.bits = ((uintptr_t)(obj)) | tag});
 }
 #   define PyStackRef_FromPyObjectSteal(obj) _PyStackRef_FromPyObjectSteal(_PyObject_CAST(obj))
@@ -193,6 +195,7 @@ PyStackRef_FromPyObjectImmortal(PyObject *obj)
 #   define PyStackRef_CLOSE(REF)                                        \
         do {                                                            \
             _PyStackRef _close_tmp = (REF);                             \
+            assert(!PyStackRef_IsNull(_close_tmp));                     \
             if (!PyStackRef_IsDeferred(_close_tmp)) {                   \
                 Py_DECREF(PyStackRef_AsPyObjectBorrow(_close_tmp));     \
             }                                                           \
@@ -214,10 +217,11 @@ PyStackRef_FromPyObjectImmortal(PyObject *obj)
 static inline _PyStackRef
 PyStackRef_DUP(_PyStackRef stackref)
 {
+    assert(!PyStackRef_IsNull(stackref));
     if (PyStackRef_IsDeferred(stackref)) {
-        assert(PyStackRef_IsNull(stackref) ||
-            _Py_IsImmortal(PyStackRef_AsPyObjectBorrow(stackref)) ||
-            _PyObject_HasDeferredRefcount(PyStackRef_AsPyObjectBorrow(stackref)));
+        assert(_Py_IsImmortal(PyStackRef_AsPyObjectBorrow(stackref)) ||
+               _PyObject_HasDeferredRefcount(PyStackRef_AsPyObjectBorrow(stackref))
+        );
         return stackref;
     }
     Py_INCREF(PyStackRef_AsPyObjectBorrow(stackref));
diff --git a/Python/bytecodes.c b/Python/bytecodes.c
index a0edf17..228d821 100644
--- a/Python/bytecodes.c
+++ b/Python/bytecodes.c
@@ -1922,9 +1922,10 @@ dummy_func(
             DECREF_INPUTS();
             ERROR_IF(super == NULL, error);
             PyObject *name = GETITEM(FRAME_CO_NAMES, oparg >> 2);
-            attr = PyStackRef_FromPyObjectSteal(PyObject_GetAttr(super, name));
+            PyObject *attr_o = PyObject_GetAttr(super, name);
             Py_DECREF(super);
-            ERROR_IF(PyStackRef_IsNull(attr), error);
+            ERROR_IF(attr_o == NULL, error);
+            attr = PyStackRef_FromPyObjectSteal(attr_o);
             null = PyStackRef_NULL;
         }
 
@@ -2740,9 +2741,10 @@ dummy_func(
 
         inst(GET_ITER, (iterable -- iter)) {
             /* before: [obj]; after [getiter(obj)] */
-            iter = PyStackRef_FromPyObjectSteal(PyObject_GetIter(PyStackRef_AsPyObjectBorrow(iterable)));
+            PyObject *iter_o = PyObject_GetIter(PyStackRef_AsPyObjectBorrow(iterable));
             DECREF_INPUTS();
-            ERROR_IF(PyStackRef_IsNull(iter), error);
+            ERROR_IF(iter_o == NULL, error);
+            iter = PyStackRef_FromPyObjectSteal(iter_o);
         }
 
         inst(GET_YIELD_FROM_ITER, (iterable -- iter)) {
@@ -3052,16 +3054,18 @@ dummy_func(
             PyObject *owner_o = PyStackRef_AsPyObjectSteal(owner);
             PyObject *name = _Py_SpecialMethods[oparg].name;
             PyObject *self_or_null_o;
-            attr = PyStackRef_FromPyObjectSteal(_PyObject_LookupSpecialMethod(owner_o, name, &self_or_null_o));
-            if (PyStackRef_IsNull(attr)) {
+            PyObject *attr_o = _PyObject_LookupSpecialMethod(owner_o, name, &self_or_null_o);
+            if (attr_o == NULL) {
                 if (!_PyErr_Occurred(tstate)) {
                     _PyErr_Format(tstate, PyExc_TypeError,
                                   _Py_SpecialMethods[oparg].error,
                                   Py_TYPE(owner_o)->tp_name);
                 }
+                ERROR_IF(true, error);
             }
-            ERROR_IF(PyStackRef_IsNull(attr), error);
-            self_or_null = PyStackRef_FromPyObjectSteal(self_or_null_o);
+            attr = PyStackRef_FromPyObjectSteal(attr_o);
+            self_or_null = self_or_null_o == NULL ?
+                PyStackRef_NULL : PyStackRef_FromPyObjectSteal(self_or_null_o);
         }
 
         inst(WITH_EXCEPT_START, (exit_func, exit_self, lasti, unused, val -- exit_func, exit_self, lasti, unused, val, res)) {
@@ -3092,9 +3096,10 @@ dummy_func(
             (void)lasti; // Shut up compiler warning if asserts are off
             PyObject *stack[5] = {NULL, PyStackRef_AsPyObjectBorrow(exit_self), exc, val_o, tb};
             int has_self = !PyStackRef_IsNull(exit_self);
-            res = PyStackRef_FromPyObjectSteal(PyObject_Vectorcall(exit_func_o, stack + 2 - has_self,
-                    (3 + has_self) | PY_VECTORCALL_ARGUMENTS_OFFSET, NULL));
-            ERROR_IF(PyStackRef_IsNull(res), error);
+            PyObject *res_o = PyObject_Vectorcall(exit_func_o, stack + 2 - has_self,
+                    (3 + has_self) | PY_VECTORCALL_ARGUMENTS_OFFSET, NULL);
+            ERROR_IF(res_o == NULL, error);
+            res = PyStackRef_FromPyObjectSteal(res_o);
         }
 
         pseudo(SETUP_FINALLY, (-- unused), (HAS_ARG)) = {
@@ -4539,9 +4544,10 @@ dummy_func(
             /* If value is a unicode object, then we know the result
              * of format(value) is value itself. */
             if (!PyUnicode_CheckExact(value_o)) {
-                res = PyStackRef_FromPyObjectSteal(PyObject_Format(value_o, NULL));
+                PyObject *res_o = PyObject_Format(value_o, NULL);
                 PyStackRef_CLOSE(value);
-                ERROR_IF(PyStackRef_IsNull(res), error);
+                ERROR_IF(res_o == NULL, error);
+                res = PyStackRef_FromPyObjectSteal(res_o);
             }
             else {
                 res = value;
diff --git a/Python/executor_cases.c.h b/Python/executor_cases.c.h
index 7631ff7..4574e18 100644
--- a/Python/executor_cases.c.h
+++ b/Python/executor_cases.c.h
@@ -3281,10 +3281,11 @@
             iterable = stack_pointer[-1];
             /* before: [obj]; after [getiter(obj)] */
             _PyFrame_SetStackPointer(frame, stack_pointer);
-            iter = PyStackRef_FromPyObjectSteal(PyObject_GetIter(PyStackRef_AsPyObjectBorrow(iterable)));
+            PyObject *iter_o = PyObject_GetIter(PyStackRef_AsPyObjectBorrow(iterable));
             stack_pointer = _PyFrame_GetStackPointer(frame);
             PyStackRef_CLOSE(iterable);
-            if (PyStackRef_IsNull(iter)) JUMP_TO_ERROR();
+            if (iter_o == NULL) JUMP_TO_ERROR();
+            iter = PyStackRef_FromPyObjectSteal(iter_o);
             stack_pointer[-1] = iter;
             break;
         }
@@ -3552,24 +3553,21 @@
             stack_pointer += -1;
             assert(WITHIN_STACK_BOUNDS());
             _PyFrame_SetStackPointer(frame, stack_pointer);
-            attr = PyStackRef_FromPyObjectSteal(_PyObject_LookupSpecialMethod(owner_o, name, &self_or_null_o));
+            PyObject *attr_o = _PyObject_LookupSpecialMethod(owner_o, name, &self_or_null_o);
             stack_pointer = _PyFrame_GetStackPointer(frame);
-            if (PyStackRef_IsNull(attr)) {
+            if (attr_o == NULL) {
                 if (!_PyErr_Occurred(tstate)) {
-                    stack_pointer[0] = attr;
-                    stack_pointer += 1;
-                    assert(WITHIN_STACK_BOUNDS());
                     _PyFrame_SetStackPointer(frame, stack_pointer);
                     _PyErr_Format(tstate, PyExc_TypeError,
                                   _Py_SpecialMethods[oparg].error,
                                   Py_TYPE(owner_o)->tp_name);
                     stack_pointer = _PyFrame_GetStackPointer(frame);
-                    stack_pointer += -1;
-                    assert(WITHIN_STACK_BOUNDS());
                 }
+                if (true) JUMP_TO_ERROR();
             }
-            if (PyStackRef_IsNull(attr)) JUMP_TO_ERROR();
-            self_or_null = PyStackRef_FromPyObjectSteal(self_or_null_o);
+            attr = PyStackRef_FromPyObjectSteal(attr_o);
+            self_or_null = self_or_null_o == NULL ?
+            PyStackRef_NULL : PyStackRef_FromPyObjectSteal(self_or_null_o);
             stack_pointer[0] = attr;
             stack_pointer[1] = self_or_null;
             stack_pointer += 2;
@@ -3613,10 +3611,11 @@
             PyObject *stack[5] = {NULL, PyStackRef_AsPyObjectBorrow(exit_self), exc, val_o, tb};
             int has_self = !PyStackRef_IsNull(exit_self);
             _PyFrame_SetStackPointer(frame, stack_pointer);
-            res = PyStackRef_FromPyObjectSteal(PyObject_Vectorcall(exit_func_o, stack + 2 - has_self,
-                    (3 + has_self) | PY_VECTORCALL_ARGUMENTS_OFFSET, NULL));
+            PyObject *res_o = PyObject_Vectorcall(exit_func_o, stack + 2 - has_self,
+                (3 + has_self) | PY_VECTORCALL_ARGUMENTS_OFFSET, NULL);
             stack_pointer = _PyFrame_GetStackPointer(frame);
-            if (PyStackRef_IsNull(res)) JUMP_TO_ERROR();
+            if (res_o == NULL) JUMP_TO_ERROR();
+            res = PyStackRef_FromPyObjectSteal(res_o);
             stack_pointer[0] = res;
             stack_pointer += 1;
             assert(WITHIN_STACK_BOUNDS());
@@ -5373,10 +5372,11 @@
              * of format(value) is value itself. */
             if (!PyUnicode_CheckExact(value_o)) {
                 _PyFrame_SetStackPointer(frame, stack_pointer);
-                res = PyStackRef_FromPyObjectSteal(PyObject_Format(value_o, NULL));
+                PyObject *res_o = PyObject_Format(value_o, NULL);
                 stack_pointer = _PyFrame_GetStackPointer(frame);
                 PyStackRef_CLOSE(value);
-                if (PyStackRef_IsNull(res)) JUMP_TO_ERROR();
+                if (res_o == NULL) JUMP_TO_ERROR();
+                res = PyStackRef_FromPyObjectSteal(res_o);
             }
             else {
                 res = value;
diff --git a/Python/generated_cases.c.h b/Python/generated_cases.c.h
index a9c1277..e1107ca 100644
--- a/Python/generated_cases.c.h
+++ b/Python/generated_cases.c.h
@@ -3856,10 +3856,11 @@
              * of format(value) is value itself. */
             if (!PyUnicode_CheckExact(value_o)) {
                 _PyFrame_SetStackPointer(frame, stack_pointer);
-                res = PyStackRef_FromPyObjectSteal(PyObject_Format(value_o, NULL));
+                PyObject *res_o = PyObject_Format(value_o, NULL);
                 stack_pointer = _PyFrame_GetStackPointer(frame);
                 PyStackRef_CLOSE(value);
-                if (PyStackRef_IsNull(res)) goto pop_1_error;
+                if (res_o == NULL) goto pop_1_error;
+                res = PyStackRef_FromPyObjectSteal(res_o);
             }
             else {
                 res = value;
@@ -4239,10 +4240,11 @@
             iterable = stack_pointer[-1];
             /* before: [obj]; after [getiter(obj)] */
             _PyFrame_SetStackPointer(frame, stack_pointer);
-            iter = PyStackRef_FromPyObjectSteal(PyObject_GetIter(PyStackRef_AsPyObjectBorrow(iterable)));
+            PyObject *iter_o = PyObject_GetIter(PyStackRef_AsPyObjectBorrow(iterable));
             stack_pointer = _PyFrame_GetStackPointer(frame);
             PyStackRef_CLOSE(iterable);
-            if (PyStackRef_IsNull(iter)) goto pop_1_error;
+            if (iter_o == NULL) goto pop_1_error;
+            iter = PyStackRef_FromPyObjectSteal(iter_o);
             stack_pointer[-1] = iter;
             DISPATCH();
         }
@@ -6270,29 +6272,21 @@
             stack_pointer += -1;
             assert(WITHIN_STACK_BOUNDS());
             _PyFrame_SetStackPointer(frame, stack_pointer);
-            attr = PyStackRef_FromPyObjectSteal(_PyObject_LookupSpecialMethod(owner_o, name, &self_or_null_o));
+            PyObject *attr_o = _PyObject_LookupSpecialMethod(owner_o, name, &self_or_null_o);
             stack_pointer = _PyFrame_GetStackPointer(frame);
-            if (PyStackRef_IsNull(attr)) {
+            if (attr_o == NULL) {
                 if (!_PyErr_Occurred(tstate)) {
-                    stack_pointer[0] = attr;
-                    stack_pointer += 1;
-                    assert(WITHIN_STACK_BOUNDS());
                     _PyFrame_SetStackPointer(frame, stack_pointer);
                     _PyErr_Format(tstate, PyExc_TypeError,
                                   _Py_SpecialMethods[oparg].error,
                                   Py_TYPE(owner_o)->tp_name);
                     stack_pointer = _PyFrame_GetStackPointer(frame);
-                    stack_pointer += -1;
-                    assert(WITHIN_STACK_BOUNDS());
                 }
+                if (true) goto error;
             }
-            if (PyStackRef_IsNull(attr)) {
-                stack_pointer[0] = attr;
-                stack_pointer += 1;
-                assert(WITHIN_STACK_BOUNDS());
-                goto error;
-            }
-            self_or_null = PyStackRef_FromPyObjectSteal(self_or_null_o);
+            attr = PyStackRef_FromPyObjectSteal(attr_o);
+            self_or_null = self_or_null_o == NULL ?
+            PyStackRef_NULL : PyStackRef_FromPyObjectSteal(self_or_null_o);
             stack_pointer[0] = attr;
             stack_pointer[1] = self_or_null;
             stack_pointer += 2;
@@ -6385,10 +6379,11 @@
                 stack_pointer += -3;
                 assert(WITHIN_STACK_BOUNDS());
                 _PyFrame_SetStackPointer(frame, stack_pointer);
-                attr = PyStackRef_FromPyObjectSteal(PyObject_GetAttr(super, name));
+                PyObject *attr_o = PyObject_GetAttr(super, name);
                 stack_pointer = _PyFrame_GetStackPointer(frame);
                 Py_DECREF(super);
-                if (PyStackRef_IsNull(attr)) goto error;
+                if (attr_o == NULL) goto error;
+                attr = PyStackRef_FromPyObjectSteal(attr_o);
                 null = PyStackRef_NULL;
             }
             stack_pointer[0] = attr;
@@ -8139,10 +8134,11 @@
             PyObject *stack[5] = {NULL, PyStackRef_AsPyObjectBorrow(exit_self), exc, val_o, tb};
             int has_self = !PyStackRef_IsNull(exit_self);
             _PyFrame_SetStackPointer(frame, stack_pointer);
-            res = PyStackRef_FromPyObjectSteal(PyObject_Vectorcall(exit_func_o, stack + 2 - has_self,
-                    (3 + has_self) | PY_VECTORCALL_ARGUMENTS_OFFSET, NULL));
+            PyObject *res_o = PyObject_Vectorcall(exit_func_o, stack + 2 - has_self,
+                (3 + has_self) | PY_VECTORCALL_ARGUMENTS_OFFSET, NULL);
             stack_pointer = _PyFrame_GetStackPointer(frame);
-            if (PyStackRef_IsNull(res)) goto error;
+            if (res_o == NULL) goto error;
+            res = PyStackRef_FromPyObjectSteal(res_o);
             stack_pointer[0] = res;
             stack_pointer += 1;
             assert(WITHIN_STACK_BOUNDS());