1 files changed, 110 insertions, 6 deletions

diff --git a/Objects/codeobject.c b/Objects/codeobject.c
index 6c0e5bf..d514ec1 100644
--- a/Objects/codeobject.c
+++ b/Objects/codeobject.c
@@ -1,3 +1,5 @@
+#include <stdbool.h>
+
 #include "Python.h"
 #include "code.h"
 #include "structmember.h"
@@ -5,6 +7,12 @@
 #define NAME_CHARS \
     "0123456789ABCDEFGHIJKLMNOPQRSTUVWXYZ_abcdefghijklmnopqrstuvwxyz"
 
+/* Holder for co_extra information */
+typedef struct {
+    Py_ssize_t ce_size;
+    void **ce_extras;
+} _PyCodeObjectExtra;
+
 /* all_name_chars(s): true iff all chars in s are valid NAME_CHARS */
 
 static int
@@ -96,7 +104,7 @@ PyCode_New(int argcount, int kwonlyargcount,
         Py_ssize_t total_args = argcount + kwonlyargcount +
             ((flags & CO_VARARGS) != 0) + ((flags & CO_VARKEYWORDS) != 0);
         Py_ssize_t alloc_size = sizeof(unsigned char) * n_cellvars;
-        int used_cell2arg = 0;
+        bool used_cell2arg = false;
         cell2arg = PyMem_MALLOC(alloc_size);
         if (cell2arg == NULL)
             return NULL;
@@ -109,7 +117,7 @@ PyCode_New(int argcount, int kwonlyargcount,
                 PyObject *arg = PyTuple_GET_ITEM(varnames, j);
                 if (!PyUnicode_Compare(cell, arg)) {
                     cell2arg[i] = j;
-                    used_cell2arg = 1;
+                    used_cell2arg = true;
                     break;
                 }
             }
@@ -152,6 +160,7 @@ PyCode_New(int argcount, int kwonlyargcount,
     co->co_lnotab = lnotab;
     co->co_zombieframe = NULL;
     co->co_weakreflist = NULL;
+    co->co_extra = NULL;
     return co;
 }
 
@@ -361,6 +370,21 @@ code_new(PyTypeObject *type, PyObject *args, PyObject *kw)
 static void
 code_dealloc(PyCodeObject *co)
 {
+    if (co->co_extra != NULL) {
+        PyThreadState *tstate = PyThreadState_Get();
+        _PyCodeObjectExtra *co_extra = co->co_extra;
+
+        for (Py_ssize_t i = 0; i < co_extra->ce_size; i++) {
+            freefunc free_extra = tstate->co_extra_freefuncs[i];
+
+            if (free_extra != NULL) {
+                free_extra(co_extra->ce_extras[i]);
+            }
+        }
+
+        PyMem_FREE(co->co_extra);
+    }
+
     Py_XDECREF(co->co_code);
     Py_XDECREF(co->co_consts);
     Py_XDECREF(co->co_names);
@@ -694,7 +718,8 @@ PyCode_Addr2Line(PyCodeObject *co, int addrq)
         addr += *p++;
         if (addr > addrq)
             break;
-        line += *p++;
+        line += (signed char)*p;
+        p++;
     }
     return line;
 }
@@ -729,17 +754,19 @@ _PyCode_CheckLineNumber(PyCodeObject* co, int lasti, PyAddrPair *bounds)
         if (addr + *p > lasti)
             break;
         addr += *p++;
-        if (*p)
+        if ((signed char)*p)
             bounds->ap_lower = addr;
-        line += *p++;
+        line += (signed char)*p;
+        p++;
         --size;
     }
 
     if (size > 0) {
         while (--size >= 0) {
             addr += *p++;
-            if (*p++)
+            if ((signed char)*p)
                 break;
+            p++;
         }
         bounds->ap_upper = addr;
     }
@@ -749,3 +776,80 @@ _PyCode_CheckLineNumber(PyCodeObject* co, int lasti, PyAddrPair *bounds)
 
     return line;
 }
+
+
+int
+_PyCode_GetExtra(PyObject *code, Py_ssize_t index, void **extra)
+{
+    assert(*extra == NULL);
+
+    if (!PyCode_Check(code)) {
+        PyErr_BadInternalCall();
+        return -1;
+    }
+
+    PyCodeObject *o = (PyCodeObject*) code;
+    _PyCodeObjectExtra *co_extra = (_PyCodeObjectExtra*) o->co_extra;
+
+
+    if (co_extra == NULL || co_extra->ce_size <= index) {
+        return 0;
+    }
+
+    *extra = co_extra->ce_extras[index];
+    return 0;
+}
+
+
+int
+_PyCode_SetExtra(PyObject *code, Py_ssize_t index, void *extra)
+{
+    PyThreadState *tstate = PyThreadState_Get();
+
+    if (!PyCode_Check(code) || index < 0 ||
+            index >= tstate->co_extra_user_count) {
+        PyErr_BadInternalCall();
+        return -1;
+    }
+
+    PyCodeObject *o = (PyCodeObject*) code;
+    _PyCodeObjectExtra *co_extra = (_PyCodeObjectExtra *) o->co_extra;
+
+    if (co_extra == NULL) {
+        o->co_extra = (_PyCodeObjectExtra*) PyMem_Malloc(
+            sizeof(_PyCodeObjectExtra));
+        if (o->co_extra == NULL) {
+            return -1;
+        }
+        co_extra = (_PyCodeObjectExtra *) o->co_extra;
+
+        co_extra->ce_extras = PyMem_Malloc(
+            tstate->co_extra_user_count * sizeof(void*));
+        if (co_extra->ce_extras == NULL) {
+            return -1;
+        }
+
+        co_extra->ce_size = tstate->co_extra_user_count;
+
+        for (Py_ssize_t i = 0; i < co_extra->ce_size; i++) {
+            co_extra->ce_extras[i] = NULL;
+        }
+    }
+    else if (co_extra->ce_size <= index) {
+        co_extra->ce_extras = PyMem_Realloc(
+            co_extra->ce_extras, tstate->co_extra_user_count * sizeof(void*));
+
+        if (co_extra->ce_extras == NULL) {
+            return -1;
+        }
+
+        co_extra->ce_size = tstate->co_extra_user_count;
+
+        for (Py_ssize_t i = co_extra->ce_size; i < co_extra->ce_size; i++) {
+            co_extra->ce_extras[i] = NULL;
+        }
+    }
+
+    co_extra->ce_extras[index] = extra;
+    return 0;
+}