Incorporate review comments courtesy of Neal Norwitz:

* Perform the code length check earlier.
* Eliminate the extra PyMem_Free() upon hitting an EXTENDED_ARG.
* Assert that the NOP count used in jump retargeting matches the NOPs
  eliminated in the final step.
* Add an XXX note to indicate that more work is being to done to
  handle linenotab with intervals > 255.

diff --git a/Python/compile.c b/Python/compile.c
index 4e7b385da6c63f1bb0d99acb9063b21dc8de680d..85182de9d7b902f4d19703a668beac47f01fbb77 100644 (file)
--- a/Python/compile.c
+++ b/Python/compile.c
@@ -435,10 +435,15 @@ optimize_code(PyObject *code, PyObject* consts, PyObject *names, PyObject *linen
        unsigned int *blocks;
        char *name;
 
-       /* Make a modifiable copy of the code string */
        if (!PyString_Check(code))
                goto exitUnchanged;
+
+       /* Avoid situations where jump retargeting could overflow */
        codelen = PyString_Size(code);
+       if (codelen > 32000)
+               goto exitUnchanged;
+
+       /* Make a modifiable copy of the code string */
        codestr = PyMem_Malloc(codelen);
        if (codestr == NULL)
                goto exitUnchanged;
@@ -449,10 +454,6 @@ optimize_code(PyObject *code, PyObject* consts, PyObject *names, PyObject *linen
        if (addrmap == NULL)
                goto exitUnchanged;
 
-       /* Avoid situations where jump retargeting could overflow */
-       if (codelen > 32000)
-               goto exitUnchanged;
-
        blocks = markblocks(codestr, codelen);
        if (blocks == NULL)
                goto exitUnchanged;
@@ -574,7 +575,6 @@ optimize_code(PyObject *code, PyObject* consts, PyObject *names, PyObject *linen
                        break;
 
                case EXTENDED_ARG:
-                       PyMem_Free(codestr);
                        goto exitUnchanged;
 
                /* Replace RETURN LOAD_CONST None RETURN with just RETURN */
@@ -590,6 +590,7 @@ optimize_code(PyObject *code, PyObject* consts, PyObject *names, PyObject *linen
        }
 
        /* Fixup linenotab */
+       /* XXX make sure this handles intervals > 256 */
        assert(PyString_Check(lineno_obj));
        lineno = PyString_AS_STRING(lineno_obj);
        tabsiz = PyString_GET_SIZE(lineno_obj);
@@ -631,6 +632,7 @@ optimize_code(PyObject *code, PyObject* consts, PyObject *names, PyObject *linen
                while (adj--)
                        codestr[h++] = codestr[i++];
        }
+       assert(h + nops == codelen);
 
        code = PyString_FromStringAndSize((char *)codestr, h);
        PyMem_Free(addrmap);