1 files changed, 37 insertions, 4 deletions

diff --git a/Python/flowgraph.c b/Python/flowgraph.c
index c8470aa..8154129 100644
--- a/Python/flowgraph.c
+++ b/Python/flowgraph.c
@@ -1297,6 +1297,7 @@ loads_const(int opcode)
     return OPCODE_HAS_CONST(opcode) || opcode == LOAD_SMALL_INT;
 }
 
+/* Returns new reference */
 static PyObject*
 get_const_value(int opcode, int oparg, PyObject *co_consts)
 {
@@ -1390,20 +1391,44 @@ nop_out(cfg_instr **instrs, int size)
     }
 }
 
-/* Steals reference to "newconst" */
+/* Does not steal reference to "newconst".
+   Return 1 if changed instruction to LOAD_SMALL_INT.
+   Return 0 if could not change instruction to LOAD_SMALL_INT.
+   Return -1 on error.
+*/
 static int
-instr_make_load_const(cfg_instr *instr, PyObject *newconst,
-                      PyObject *consts, PyObject *const_cache)
+maybe_instr_make_load_smallint(cfg_instr *instr, PyObject *newconst,
+                               PyObject *consts, PyObject *const_cache)
 {
     if (PyLong_CheckExact(newconst)) {
         int overflow;
         long val = PyLong_AsLongAndOverflow(newconst, &overflow);
+        if (val == -1 && PyErr_Occurred()) {
+            return -1;
+        }
         if (!overflow && _PY_IS_SMALL_INT(val)) {
             assert(_Py_IsImmortal(newconst));
             INSTR_SET_OP1(instr, LOAD_SMALL_INT, (int)val);
-            return SUCCESS;
+            return 1;
         }
     }
+    return 0;
+}
+
+
+/* Steals reference to "newconst" */
+static int
+instr_make_load_const(cfg_instr *instr, PyObject *newconst,
+                      PyObject *consts, PyObject *const_cache)
+{
+    int res = maybe_instr_make_load_smallint(instr, newconst, consts, const_cache);
+    if (res < 0) {
+        Py_DECREF(newconst);
+        return ERROR;
+    }
+    if (res > 0) {
+        return SUCCESS;
+    }
     int oparg = add_const(newconst, consts, const_cache);
     RETURN_IF_ERROR(oparg);
     INSTR_SET_OP1(instr, LOAD_CONST, oparg);
@@ -2042,6 +2067,14 @@ basicblock_optimize_load_const(PyObject *const_cache, basicblock *bb, PyObject *
     int oparg = 0;
     for (int i = 0; i < bb->b_iused; i++) {
         cfg_instr *inst = &bb->b_instr[i];
+        if (inst->i_opcode == LOAD_CONST) {
+            PyObject *constant = get_const_value(inst->i_opcode, inst->i_oparg, consts);
+            int res = maybe_instr_make_load_smallint(inst, constant, consts, const_cache);
+            Py_DECREF(constant);
+            if (res < 0) {
+                return ERROR;
+            }
+        }
         bool is_copy_of_load_const = (opcode == LOAD_CONST &&
                                       inst->i_opcode == COPY &&
                                       inst->i_oparg == 1);