1 files changed, 45 insertions, 28 deletions

diff --git a/Python/optimizer.c b/Python/optimizer.c
index 8f7a972..95cd782 100644
--- a/Python/optimizer.c
+++ b/Python/optimizer.c
@@ -9,20 +9,31 @@
 #include <stdint.h>
 #include <stddef.h>
 
-/* Returns the index of the next space, or -1 if there is no
- * more space. Doesn't set an exception. */
+static bool
+has_space_for_executor(PyCodeObject *code, _Py_CODEUNIT *instr)
+{
+    if (instr->op.code == ENTER_EXECUTOR) {
+        return true;
+    }
+    if (code->co_executors == NULL) {
+        return true;
+    }
+    return code->co_executors->size < 256;
+}
+
 static int32_t
-get_next_free_in_executor_array(PyCodeObject *code)
+get_index_for_executor(PyCodeObject *code, _Py_CODEUNIT *instr)
 {
+    if (instr->op.code == ENTER_EXECUTOR) {
+        return instr->op.arg;
+    }
     _PyExecutorArray *old = code->co_executors;
     int size = 0;
     int capacity = 0;
     if (old != NULL) {
         size = old->size;
         capacity = old->capacity;
-        if (capacity >= 256) {
-            return -1;
-        }
+        assert(size < 256);
     }
     assert(size <= capacity);
     if (size == capacity) {
@@ -40,46 +51,36 @@ get_next_free_in_executor_array(PyCodeObject *code)
         code->co_executors = new;
     }
     assert(size < code->co_executors->capacity);
-    code->co_executors->size++;
     return size;
 }
 
 static void
 insert_executor(PyCodeObject *code, _Py_CODEUNIT *instr, int index, _PyExecutorObject *executor)
 {
+    Py_INCREF(executor);
     if (instr->op.code == ENTER_EXECUTOR) {
         assert(index == instr->op.arg);
         _PyExecutorObject *old = code->co_executors->executors[index];
         executor->vm_data.opcode = old->vm_data.opcode;
         executor->vm_data.oparg = old->vm_data.oparg;
         old->vm_data.opcode = 0;
-        Py_INCREF(executor);
         code->co_executors->executors[index] = executor;
         Py_DECREF(old);
     }
     else {
-        Py_INCREF(executor);
+        assert(code->co_executors->size == index);
+        assert(code->co_executors->capacity > index);
         executor->vm_data.opcode = instr->op.code;
         executor->vm_data.oparg = instr->op.arg;
         code->co_executors->executors[index] = executor;
         assert(index < 256);
         instr->op.code = ENTER_EXECUTOR;
         instr->op.arg = index;
+        code->co_executors->size++;
     }
     return;
 }
 
-static int
-get_executor_index(PyCodeObject *code, _Py_CODEUNIT *instr)
-{
-    if (instr->op.code == ENTER_EXECUTOR) {
-        return instr->op.arg;
-    }
-    else {
-        return get_next_free_in_executor_array(code);
-    }
-}
-
 int
 PyUnstable_Replace_Executor(PyCodeObject *code, _Py_CODEUNIT *instr, _PyExecutorObject *new)
 {
@@ -87,7 +88,7 @@ PyUnstable_Replace_Executor(PyCodeObject *code, _Py_CODEUNIT *instr, _PyExecutor
         PyErr_Format(PyExc_ValueError, "No executor to replace");
         return -1;
     }
-    int index = get_executor_index(code, instr);
+    int index = instr->op.arg;
     assert(index >= 0);
     insert_executor(code, instr, index, new);
     return 0;
@@ -126,6 +127,8 @@ PyUnstable_GetOptimizer(void)
     if (interp->optimizer == &_PyOptimizer_Default) {
         return NULL;
     }
+    assert(interp->optimizer_backedge_threshold == interp->optimizer->backedge_threshold);
+    assert(interp->optimizer_resume_threshold == interp->optimizer->resume_threshold);
     Py_INCREF(interp->optimizer);
     return interp->optimizer;
 }
@@ -151,23 +154,37 @@ _PyOptimizer_BackEdge(_PyInterpreterFrame *frame, _Py_CODEUNIT *src, _Py_CODEUNI
     PyCodeObject *code = (PyCodeObject *)frame->f_executable;
     assert(PyCode_Check(code));
     PyInterpreterState *interp = PyInterpreterState_Get();
-    int index = get_executor_index(code, src);
-    if (index < 0) {
-        _PyFrame_SetStackPointer(frame, stack_pointer);
-        return frame;
+    if (!has_space_for_executor(code, src)) {
+        goto jump_to_destination;
     }
     _PyOptimizerObject *opt = interp->optimizer;
-    _PyExecutorObject *executor;
+    _PyExecutorObject *executor = NULL;
     int err = opt->optimize(opt, code, dest, &executor);
     if (err <= 0) {
+        assert(executor == NULL);
         if (err < 0) {
             return NULL;
         }
-        _PyFrame_SetStackPointer(frame, stack_pointer);
-        return frame;
+        goto jump_to_destination;
+    }
+    int index = get_index_for_executor(code, src);
+    if (index < 0) {
+        /* Out of memory. Don't raise and assume that the
+         * error will show up elsewhere.
+         *
+         * If an optimizer has already produced an executor,
+         * it might get confused by the executor disappearing,
+         * but there is not much we can do about that here. */
+        Py_DECREF(executor);
+        goto jump_to_destination;
     }
     insert_executor(code, src, index, executor);
+    assert(frame->prev_instr == src);
     return executor->execute(executor, frame, stack_pointer);
+jump_to_destination:
+    frame->prev_instr = dest - 1;
+    _PyFrame_SetStackPointer(frame, stack_pointer);
+    return frame;
 }
 
 /** Test support **/