use tagged ints for indexes. Work in progress

Author: markshannon

Include/internal/pycore_ceval.h

@@ -362,6 +362,8 @@ PyAPI_FUNC(_PyStackRef) _PyFloat_FromDouble_ConsumeInputs(_PyStackRef left, _PyS
     #endif
 #endif
 
+_PyStackRef _PyForIter_NextWithIndex(PyObject *seq, _PyStackRef index);
+
 #ifdef __cplusplus
 }
 #endif

Include/internal/pycore_code.h

@@ -311,7 +311,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);

Include/internal/pycore_opcode_metadata.h

@@ -231,6 +231,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);
 
@@ -255,7 +257,16 @@ static inline intptr_t
 PyStackRef_UntagInt(_PyStackRef i)
 {
     assert((i.bits & Py_INT_TAG) == Py_INT_TAG);
-    return Py_ARITHMETIC_RIGHT_SHIFT(intptr_t, i.bits, 2);
+    intptr_t val = (intptr_t)i.bits;
+    return Py_ARITHMETIC_RIGHT_SHIFT(intptr_t, val, 2);
+}
+
+
+static inline _PyStackRef
+PyStackRef_IncrementTaggedInt(_PyStackRef ref)
+{
+    assert(ref.bits != (uintptr_t)-1); // Overflow
+    return (_PyStackRef){ .bits = ref.bits + 4 };
 }
 
 
@@ -683,7 +694,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.