Use tagged integers on the evaluation stack for lasti

Author: markshannon

Include/internal/pycore_stackref.h

@@ -210,8 +210,32 @@ _PyStackRef_FromPyObjectNewMortal(PyObject *obj, const char *filename, int linen
 
 extern int PyStackRef_Is(_PyStackRef a, _PyStackRef b);
 
+extern _PyStackRef PyStackRef_TagInt(intptr_t i);
+
 #else
 
+#define Py_INT_TAG 3
+
+static inline bool
+PyStackRef_IsTaggedInt(_PyStackRef i)
+{
+    return (i.bits & Py_INT_TAG) == Py_INT_TAG;
+}
+
+static inline _PyStackRef
+PyStackRef_TagInt(intptr_t i)
+{
+    assert(Py_ARITHMETIC_RIGHT_SHIFT(intptr_t, (i << 2), 2) == i);
+    return (_PyStackRef){ .bits = ((((uintptr_t)i) << 2) | Py_INT_TAG) };
+}
+
+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);
+}
+
 
 #ifdef Py_GIL_DISABLED
 
@@ -451,6 +475,7 @@ PyStackRef_RefcountOnObject(_PyStackRef ref)
 static inline PyObject *
 PyStackRef_AsPyObjectBorrow(_PyStackRef ref)
 {
+    assert(!PyStackRef_IsTaggedInt(ref));
     return BITS_TO_PTR_MASKED(ref);
 }
 

Python/bytecodes.c

@@ -1385,16 +1385,7 @@ dummy_func(
 
             assert(oparg >= 0 && oparg <= 2);
             if (oparg) {
-                PyObject *lasti = PyStackRef_AsPyObjectBorrow(values[0]);
-                if (PyLong_Check(lasti)) {
-                    frame->instr_ptr = _PyFrame_GetBytecode(frame) + PyLong_AsLong(lasti);
-                    assert(!_PyErr_Occurred(tstate));
-                }
-                else {
-                    _PyErr_SetString(tstate, PyExc_SystemError, "lasti is not an int");
-                    Py_DECREF(exc);
-                    ERROR_NO_POP();
-                }
+                frame->instr_ptr = _PyFrame_GetBytecode(frame) + PyStackRef_UntagInt(values[0]);
             }
             assert(exc && PyExceptionInstance_Check(exc));
             _PyErr_SetRaisedException(tstate, exc);
@@ -3458,7 +3449,7 @@ dummy_func(
             if (tb == NULL) {
                 tb = Py_None;
             }
-            assert(PyStackRef_LongCheck(lasti));
+            assert(PyStackRef_IsTaggedInt(lasti));
             (void)lasti; // Shut up compiler warning if asserts are off
             PyObject *stack[5] = {NULL, PyStackRef_AsPyObjectBorrow(exit_self), exc, val_o, tb};
             int has_self = !PyStackRef_IsNull(exit_self);
@@ -5351,11 +5342,8 @@ dummy_func(
             }
             if (lasti) {
                 int frame_lasti = _PyInterpreterFrame_LASTI(frame);
-                PyObject *lasti = PyLong_FromLong(frame_lasti);
-                if (lasti == NULL) {
-                    goto exception_unwind;
-                }
-                _PyFrame_StackPush(frame, PyStackRef_FromPyObjectSteal(lasti));
+                _PyStackRef lasti = PyStackRef_TagInt(frame_lasti);
+                _PyFrame_StackPush(frame, lasti);
             }
 
             /* Make the raw exception data

Python/ceval.c

@@ -146,6 +146,9 @@ dump_item(_PyStackRef item)
         printf("<NULL>");
         return;
     }
+    if (PyStackRef_IsTaggedInt(item)) {
+        printf("%ld", PyStackRef_UntagInt(item));
+    }
     PyObject *obj = PyStackRef_AsPyObjectBorrow(item);
     if (obj == NULL) {
         printf("<nil>");

Python/executor_cases.c.h

@@ -547,8 +547,10 @@ _PyGC_VisitStackRef(_PyStackRef *ref, visitproc visit, void *arg)
     // This is a bit tricky! We want to ignore stackrefs with embedded
     // refcounts when computing the incoming references, but otherwise treat
     // them like normal.
-    if (!PyStackRef_RefcountOnObject(*ref) && (visit == visit_decref)) {
-        return 0;
+    if (!PyStackRef_RefcountOnObject(*ref)) {
+        if (visit == visit_decref || PyStackRef_IsTaggedInt(*ref)) {
+            return 0;
+        }
     }
     Py_VISIT(PyStackRef_AsPyObjectBorrow(*ref));
     return 0;
@@ -1495,6 +1497,9 @@ mark_stacks(PyInterpreterState *interp, PyGC_Head *visited, int visited_space, b
             objects_marked += move_to_reachable(func, &reachable, visited_space);
             while (sp > locals) {
                 sp--;
+                if (PyStackRef_IsTaggedInt(*sp)) {
+                    continue;
+                }
                 PyObject *op = PyStackRef_AsPyObjectBorrow(*sp);
                 if (op == NULL || _Py_IsImmortal(op)) {
                     continue;

Python/gc.c

@@ -1694,11 +1694,12 @@ _PyGC_VisitStackRef(_PyStackRef *ref, visitproc visit, void *arg)
     // This is a bit tricky! We want to ignore deferred references when
     // computing the incoming references, but otherwise treat them like
     // regular references.
-    if (!PyStackRef_IsDeferred(*ref) ||
-        (visit != visit_decref && visit != visit_decref_unreachable))
-    {
-        Py_VISIT(PyStackRef_AsPyObjectBorrow(*ref));
+    if (PyStackRef_IsDeferred(*ref)) {
+        if ((visit != visit_decref && visit != visit_decref_unreachable) || PyStackRef_IsTaggedInt(*ref)) {
+            return 0;
+        }
     }
+    Py_VISIT(PyStackRef_AsPyObjectBorrow(*ref));
     return 0;
 }