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

@@ -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.

Include/internal/pycore_stackref.h

@@ -155,6 +155,9 @@ dump_item(_PyStackRef item)
         printf("<nil>");
         return;
     }
+    if (PyList_CheckExact(obj)) {
+        printf("len=%ld ", Py_SIZE(obj));
+    }
     // Don't call __repr__(), it might recurse into the interpreter.
     printf("<%s at %p>", Py_TYPE(obj)->tp_name, (void *)obj);
 }
@@ -271,6 +274,7 @@ maybe_lltrace_resume_frame(_PyInterpreterFrame *frame, PyObject *globals)
             lltrace = *python_lltrace - '0';  // TODO: Parse an int and all that
         }
     }
+    // lltrace = 5;
     if (lltrace >= 5) {
         lltrace_resume_frame(frame);
     }
@@ -3380,3 +3384,17 @@ _PyEval_LoadName(PyThreadState *tstate, _PyInterpreterFrame *frame, PyObject *na
     }
     return value;
 }
+
+_PyStackRef
+_PyForIter_NextWithIndex(PyObject *seq, _PyStackRef index)
+{
+    assert(PyStackRef_IsTaggedInt(index));
+    assert(Py_TYPE(seq) == &PyTuple_Type || Py_TYPE(seq) == &PyList_Type);
+    size_t size = Py_SIZE(seq);
+    intptr_t i = PyStackRef_UntagInt(index);
+    if ((size_t)i >= size) {
+        return PyStackRef_NULL;
+    }
+    PyObject *next_o = PySequence_Fast_GET_ITEM(seq, i);
+    return PyStackRef_FromPyObjectNew(next_o);
+}

Include/internal/pycore_uop_metadata.h

@@ -4463,9 +4463,9 @@ codegen_async_comprehension_generator(compiler *c, location loc,
         else {
             /* Sub-iter - calculate on the fly */
             VISIT(c, expr, gen->iter);
-            ADDOP(c, LOC(gen->iter), GET_AITER);
         }
     }
+    ADDOP(c, LOC(gen->iter), GET_AITER);
 
     USE_LABEL(c, start);
     /* Runtime will push a block here, so we need to account for that */