Speedup recursion checks

python · markshannon · Apr 30, 2025 · Apr 8, 2025 · Apr 8, 2025 · Apr 8, 2025
commit 3514aa11b8cfada41f90ebaa8b701886ba6b0b29
diff --git a/Include/internal/pycore_ceval.h b/Include/internal/pycore_ceval.h
@@ -196,25 +196,6 @@ extern void _PyEval_DeactivateOpCache(void);
 
 /* --- _Py_EnterRecursiveCall() ----------------------------------------- */
 
-#if !_Py__has_builtin(__builtin_frame_address) && !defined(_MSC_VER)
-static uintptr_t return_pointer_as_int(char* p) {
-    return (uintptr_t)p;
-}
-#endif
-
-static inline uintptr_t
-_Py_get_machine_stack_pointer(void) {
-#if _Py__has_builtin(__builtin_frame_address)
-    return (uintptr_t)__builtin_frame_address(0);
-#elif defined(_MSC_VER)
-    return (uintptr_t)_AddressOfReturnAddress();
-#else
-    char here;
-    /* Avoid compiler warning about returning stack address */
-    return return_pointer_as_int(&here);
-#endif
-}
-
 static inline int _Py_MakeRecCheck(PyThreadState *tstate)  {
     uintptr_t here_addr = _Py_get_machine_stack_pointer();
     _PyThreadStateImpl *_tstate = (_PyThreadStateImpl *)tstate;
@@ -249,12 +230,7 @@ PyAPI_FUNC(void) _Py_InitializeRecursionLimits(PyThreadState *tstate);
 static inline int _Py_ReachedRecursionLimit(PyThreadState *tstate)  {
     uintptr_t here_addr = _Py_get_machine_stack_pointer();
     _PyThreadStateImpl *_tstate = (_PyThreadStateImpl *)tstate;
-    if (here_addr > _tstate->c_stack_soft_limit) {
-        return 0;
-    }
-    if (_tstate->c_stack_hard_limit == 0) {
-        _Py_InitializeRecursionLimits(tstate);
-    }
+    assert(_tstate->c_stack_hard_limit != 0);
     return here_addr <= _tstate->c_stack_soft_limit;
 }
 

@@ -9,6 +9,7 @@ extern "C" {
 #endif
 
 #include "pycore_typedefs.h"      // _PyRuntimeState
+#include "pycore_tstate.h"
 
 
 // Values for PyThreadState.state. A thread must be in the "attached" state
@@ -296,6 +297,34 @@ _Py_AssertHoldsTstateFunc(const char *func)
 #define _Py_AssertHoldsTstate()
 #endif
 
+#if !_Py__has_builtin(__builtin_frame_address) && !defined(_MSC_VER)
+static uintptr_t return_pointer_as_int(char* p) {
+    return (uintptr_t)p;
+}
+#endif
+
+static inline uintptr_t
+_Py_get_machine_stack_pointer(void) {
+#if _Py__has_builtin(__builtin_frame_address)
+    return (uintptr_t)__builtin_frame_address(0);
+#elif defined(_MSC_VER)
+    return (uintptr_t)_AddressOfReturnAddress();
+#else
+    char here;
+    /* Avoid compiler warning about returning stack address */
+    return return_pointer_as_int(&here);
+#endif
+}
+
+static inline intptr_t
+_Py_RecursionLimit_GetMargin(PyThreadState *tstate)
+{
+    _PyThreadStateImpl *_tstate = (_PyThreadStateImpl *)tstate;
+    assert(_tstate->c_stack_hard_limit != 0);
+    intptr_t here_addr = _Py_get_machine_stack_pointer();
+    return Py_ARITHMETIC_RIGHT_SHIFT(intptr_t, here_addr - (intptr_t)_tstate->c_stack_soft_limit, PYOS_STACK_MARGIN_SHIFT);
+}
+
 #ifdef __cplusplus
 }
 #endif

@@ -26,17 +26,25 @@ PyAPI_DATA(int) (*PyOS_InputHook)(void);
  * apart. In practice, that means it must be larger than the C
  * stack consumption of PyEval_EvalDefault */
 #if defined(_Py_ADDRESS_SANITIZER) || defined(_Py_THREAD_SANITIZER)
-#  define PYOS_STACK_MARGIN 4096
+#  define PYOS_LOG_STACK_MARGIN 12
 #elif defined(Py_DEBUG) && defined(WIN32)
-#  define PYOS_STACK_MARGIN 4096
+#  define PYOS_LOG_STACK_MARGIN 12
 #elif defined(__wasi__)
    /* Web assembly has two stacks, so this isn't really a size */
-#  define PYOS_STACK_MARGIN 500
+#  define PYOS_LOG_STACK_MARGIN 9
 #else
-#  define PYOS_STACK_MARGIN 2048
+#  define PYOS_LOG_STACK_MARGIN 11
 #endif
+#define PYOS_STACK_MARGIN (1 << PYOS_LOG_STACK_MARGIN)
 #define PYOS_STACK_MARGIN_BYTES (PYOS_STACK_MARGIN * sizeof(void *))
 
+#if SIZEOF_VOID_P == 8
+#define PYOS_STACK_MARGIN_SHIFT (PYOS_LOG_STACK_MARGIN + 3)
+#else
+#define PYOS_STACK_MARGIN_SHIFT (PYOS_LOG_STACK_MARGIN + 2)
+#endif
+
+
 #if defined(WIN32)
 #define USE_STACKCHECK
 #endif

@@ -3015,7 +3015,8 @@ _Py_Dealloc(PyObject *op)
     PyTypeObject *type = Py_TYPE(op);
     destructor dealloc = type->tp_dealloc;
     PyThreadState *tstate = _PyThreadState_GET();
-    if (_Py_ReachedRecursionLimitWithMargin(tstate, 2)) {
+    intptr_t margin = _Py_RecursionLimit_GetMargin(tstate);
+    if (margin < 2) {
         _PyTrash_thread_deposit_object(tstate, (PyObject *)op);
         return;
     }
@@ -3061,7 +3062,7 @@ _Py_Dealloc(PyObject *op)
     Py_XDECREF(old_exc);
     Py_DECREF(type);
 #endif
-    if (tstate->delete_later && !_Py_ReachedRecursionLimitWithMargin(tstate, 4)) {
+    if (tstate->delete_later && margin >= 4) {
         _PyTrash_thread_destroy_chain(tstate);
     }
 }

@@ -476,12 +476,6 @@ _Py_CheckRecursiveCall(PyThreadState *tstate, const char *where)
     _PyThreadStateImpl *_tstate = (_PyThreadStateImpl *)tstate;
     uintptr_t here_addr = _Py_get_machine_stack_pointer();
     assert(_tstate->c_stack_soft_limit != 0);
-    if (_tstate->c_stack_hard_limit == 0) {
-        _Py_InitializeRecursionLimits(tstate);
-    }
-    if (here_addr >= _tstate->c_stack_soft_limit) {
-        return 0;
-    }
     assert(_tstate->c_stack_hard_limit != 0);
     if (here_addr < _tstate->c_stack_hard_limit) {
         /* Overflowing while handling an overflow. Give up. */

@@ -2132,7 +2132,10 @@ _PyThreadState_Attach(PyThreadState *tstate)
     if (current_fast_get() != NULL) {
         Py_FatalError("non-NULL old thread state");
     }
-
+    _PyThreadStateImpl *_tstate = (_PyThreadStateImpl *)tstate;
+    if (_tstate->c_stack_hard_limit == 0) {
+        _Py_InitializeRecursionLimits(tstate);
+    }
 
     while (1) {
         _PyEval_AcquireLock(tstate);