[3.11] GH-96569: Avoid undefined behavior (#96616)

Co-authored-by: Michael Droettboom <mdboom@gmail.com>

1 files changed, 8 insertions, 8 deletions

diff --git a/Python/pystate.c b/Python/pystate.c
index a0c12ba..3a8140b 100644
--- a/Python/pystate.c
+++ b/Python/pystate.c
@@ -2178,16 +2178,16 @@ push_chunk(PyThreadState *tstate, int size)
 _PyInterpreterFrame *
 _PyThreadState_BumpFramePointerSlow(PyThreadState *tstate, size_t size)
 {
-    assert(size < INT_MAX/sizeof(PyObject *));
-    PyObject **base = tstate->datastack_top;
-    PyObject **top = base + size;
-    if (top >= tstate->datastack_limit) {
-        base = push_chunk(tstate, (int)size);
+    if (_PyThreadState_HasStackSpace(tstate, size)) {
+        _PyInterpreterFrame *res = (_PyInterpreterFrame *)tstate->datastack_top;
+        tstate->datastack_top += size;
+        return res;
     }
-    else {
-        tstate->datastack_top = top;
+    if (size > INT_MAX/2) {
+        PyErr_NoMemory();
+        return NULL;
     }
-    return (_PyInterpreterFrame *)base;
+    return (_PyInterpreterFrame *)push_chunk(tstate, (int)size);
 }
 
 void