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.

1 files changed, 8 insertions, 6 deletions

diff --git a/Python/compile.c b/Python/compile.c
index 4e7b385..85182de 100644
--- 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);