1 files changed, 38 insertions, 11 deletions

diff --git a/Python/assemble.c b/Python/assemble.c
index b701253..4f66cf2 100644
--- a/Python/assemble.c
+++ b/Python/assemble.c
@@ -18,7 +18,7 @@
 #define ERROR -1
 
 #define RETURN_IF_ERROR(X)  \
-    if ((X) == -1) {        \
+    if ((X) < 0) {          \
         return ERROR;       \
     }
 
@@ -448,13 +448,17 @@ static PyObject *
 dict_keys_inorder(PyObject *dict, Py_ssize_t offset)
 {
     PyObject *tuple, *k, *v;
-    Py_ssize_t i, pos = 0, size = PyDict_GET_SIZE(dict);
+    Py_ssize_t pos = 0, size = PyDict_GET_SIZE(dict);
 
     tuple = PyTuple_New(size);
     if (tuple == NULL)
         return NULL;
     while (PyDict_Next(dict, &pos, &k, &v)) {
-        i = PyLong_AS_LONG(v);
+        Py_ssize_t i = PyLong_AsSsize_t(v);
+        if (i == -1 && PyErr_Occurred()) {
+            Py_DECREF(tuple);
+            return NULL;
+        }
         assert((i - offset) < size);
         assert((i - offset) >= 0);
         PyTuple_SET_ITEM(tuple, i - offset, Py_NewRef(k));
@@ -466,24 +470,34 @@ dict_keys_inorder(PyObject *dict, Py_ssize_t offset)
 extern void _Py_set_localsplus_info(int, PyObject *, unsigned char,
                                    PyObject *, PyObject *);
 
-static void
+static int
 compute_localsplus_info(_PyCompile_CodeUnitMetadata *umd, int nlocalsplus,
                         PyObject *names, PyObject *kinds)
 {
     PyObject *k, *v;
     Py_ssize_t pos = 0;
     while (PyDict_Next(umd->u_varnames, &pos, &k, &v)) {
-        int offset = (int)PyLong_AS_LONG(v);
+        int offset = _PyLong_AsInt(v);
+        if (offset == -1 && PyErr_Occurred()) {
+            return ERROR;
+        }
         assert(offset >= 0);
         assert(offset < nlocalsplus);
+
         // For now we do not distinguish arg kinds.
         _PyLocals_Kind kind = CO_FAST_LOCAL;
-        if (PyDict_Contains(umd->u_fasthidden, k)) {
+        int has_key = PyDict_Contains(umd->u_fasthidden, k);
+        RETURN_IF_ERROR(has_key);
+        if (has_key) {
             kind |= CO_FAST_HIDDEN;
         }
-        if (PyDict_GetItem(umd->u_cellvars, k) != NULL) {
+
+        has_key = PyDict_Contains(umd->u_cellvars, k);
+        RETURN_IF_ERROR(has_key);
+        if (has_key) {
             kind |= CO_FAST_CELL;
         }
+
         _Py_set_localsplus_info(offset, k, kind, names, kinds);
     }
     int nlocals = (int)PyDict_GET_SIZE(umd->u_varnames);
@@ -492,12 +506,18 @@ compute_localsplus_info(_PyCompile_CodeUnitMetadata *umd, int nlocalsplus,
     int numdropped = 0;
     pos = 0;
     while (PyDict_Next(umd->u_cellvars, &pos, &k, &v)) {
-        if (PyDict_GetItem(umd->u_varnames, k) != NULL) {
+        int has_name = PyDict_Contains(umd->u_varnames, k);
+        RETURN_IF_ERROR(has_name);
+        if (has_name) {
             // Skip cells that are already covered by locals.
             numdropped += 1;
             continue;
         }
-        int offset = (int)PyLong_AS_LONG(v);
+
+        int offset = _PyLong_AsInt(v);
+        if (offset == -1 && PyErr_Occurred()) {
+            return ERROR;
+        }
         assert(offset >= 0);
         offset += nlocals - numdropped;
         assert(offset < nlocalsplus);
@@ -506,12 +526,16 @@ compute_localsplus_info(_PyCompile_CodeUnitMetadata *umd, int nlocalsplus,
 
     pos = 0;
     while (PyDict_Next(umd->u_freevars, &pos, &k, &v)) {
-        int offset = (int)PyLong_AS_LONG(v);
+        int offset = _PyLong_AsInt(v);
+        if (offset == -1 && PyErr_Occurred()) {
+            return ERROR;
+        }
         assert(offset >= 0);
         offset += nlocals - numdropped;
         assert(offset < nlocalsplus);
         _Py_set_localsplus_info(offset, k, CO_FAST_FREE, names, kinds);
     }
+    return SUCCESS;
 }
 
 static PyCodeObject *
@@ -556,7 +580,10 @@ makecode(_PyCompile_CodeUnitMetadata *umd, struct assembler *a, PyObject *const_
     if (localspluskinds == NULL) {
         goto error;
     }
-    compute_localsplus_info(umd, nlocalsplus, localsplusnames, localspluskinds);
+    if (compute_localsplus_info(umd, nlocalsplus,
+                                localsplusnames, localspluskinds) == ERROR) {
+        goto error;
+    }
 
     struct _PyCodeConstructor con = {
         .filename = filename,