use tagged ints for indexes. Work in progress

python · markshannon · May 27, 2025 · Apr 10, 2025 · Apr 11, 2025 · Apr 15, 2025
commit 38a429f679839cb2e548a198c708064e9793ba82
diff --git a/Include/internal/pycore_ceval.h b/Include/internal/pycore_ceval.h
@@ -377,6 +377,8 @@ PyAPI_FUNC(_PyStackRef) _PyFloat_FromDouble_ConsumeInputs(_PyStackRef left, _PyS
 extern int _PyRunRemoteDebugger(PyThreadState *tstate);
 #endif

+_PyStackRef _PyForIter_NextWithIndex(PyObject *seq, _PyStackRef index);
+
 #ifdef __cplusplus
 }
 #endif

diff --git a/Include/internal/pycore_code.h b/Include/internal/pycore_code.h
@@ -313,7 +313,7 @@ extern void _Py_Specialize_CompareOp(_PyStackRef lhs, _PyStackRef rhs,
                                     _Py_CODEUNIT *instr, int oparg);
 extern void _Py_Specialize_UnpackSequence(_PyStackRef seq, _Py_CODEUNIT *instr,
                                          int oparg);
-extern void _Py_Specialize_ForIter(_PyStackRef iter, _Py_CODEUNIT *instr, int oparg);
+extern void _Py_Specialize_ForIter(_PyStackRef iter, _PyStackRef null_or_index, _Py_CODEUNIT *instr, int oparg);
 extern void _Py_Specialize_Send(_PyStackRef receiver, _Py_CODEUNIT *instr);
 extern void _Py_Specialize_ToBool(_PyStackRef value, _Py_CODEUNIT *instr);
 extern void _Py_Specialize_ContainsOp(_PyStackRef value, _Py_CODEUNIT *instr);

diff --git a/Include/internal/pycore_opcode_metadata.h b/Include/internal/pycore_opcode_metadata.h
@@ -233,6 +233,8 @@ extern intptr_t PyStackRef_UntagInt(_PyStackRef ref);

 extern _PyStackRef PyStackRef_TagInt(intptr_t i);

+extern _PyStackRef PyStackRef_IncrementTaggedInt(_PyStackRef ref);
+
 extern bool
 PyStackRef_IsNullOrInt(_PyStackRef ref);

@@ -262,6 +264,14 @@ PyStackRef_UntagInt(_PyStackRef i)
 }


+static inline _PyStackRef
+PyStackRef_IncrementTaggedInt(_PyStackRef ref)
+{
+    assert(ref.bits != (uintptr_t)-1); // Overflow
+    return (_PyStackRef){ .bits = ref.bits + 4 };
+}
+
+
 #ifdef Py_GIL_DISABLED

 #define Py_TAG_DEFERRED (1)
@@ -686,7 +696,13 @@ PyStackRef_XCLOSE(_PyStackRef ref)

 #endif // !defined(Py_GIL_DISABLED) && defined(Py_STACKREF_DEBUG)

-#define PyStackRef_TYPE(stackref) Py_TYPE(PyStackRef_AsPyObjectBorrow(stackref))
+static inline PyTypeObject *
+PyStackRef_TYPE(_PyStackRef stackref) {
+    if (PyStackRef_IsTaggedInt(stackref)) {
+        return &PyLong_Type;
+    }
+    return Py_TYPE(PyStackRef_AsPyObjectBorrow(stackref));
+}

 // Converts a PyStackRef back to a PyObject *, converting the
 // stackref to a new reference.

diff --git a/Include/internal/pycore_uop_ids.h b/Include/internal/pycore_uop_ids.h
diff --git a/Include/internal/pycore_uop_metadata.h b/Include/internal/pycore_uop_metadata.h