bpo-40228: More robust frame.setlineno. (GH-19437)

More robust frame.setlineno. Makes no assumptions about source->bytecode translation.

1 files changed, 251 insertions, 312 deletions

diff --git a/Objects/frameobject.c b/Objects/frameobject.c
index 451c895..4f5054d 100644
--- a/Objects/frameobject.c
+++ b/Objects/frameobject.c
@@ -69,263 +69,223 @@ get_arg(const _Py_CODEUNIT *codestr, Py_ssize_t i)
     return oparg;
 }
 
-typedef struct _codetracker {
-    unsigned char *code;
-    Py_ssize_t code_len;
-    unsigned char *lnotab;
-    Py_ssize_t lnotab_len;
-    int start_line;
-    int offset;
-    int line;
-    int addr;
-    int line_addr;
-} codetracker;
-
-/* Reset the mutable parts of the tracker */
-static void
-reset(codetracker *tracker)
-{
-    tracker->offset = 0;
-    tracker->addr = 0;
-    tracker->line_addr = 0;
-    tracker->line = tracker->start_line;
-}
+typedef enum kind {
+    With = 1,
+    Loop = 2,
+    Try = 3,
+    Except = 4,
+} Kind;
 
-/* Initialise the tracker */
-static void
-init_codetracker(codetracker *tracker, PyCodeObject *code_obj)
+#define BITS_PER_BLOCK 3
+
+static inline int64_t
+push_block(int64_t stack, Kind kind)
 {
-    PyBytes_AsStringAndSize(code_obj->co_code,
-                            (char **)&tracker->code, &tracker->code_len);
-    PyBytes_AsStringAndSize(code_obj->co_lnotab,
-                            (char **)&tracker->lnotab, &tracker->lnotab_len);
-    tracker->start_line = code_obj->co_firstlineno;
-    reset(tracker);
+    assert(stack < ((int64_t)1)<<(BITS_PER_BLOCK*CO_MAXBLOCKS));
+    return (stack << BITS_PER_BLOCK) | kind;
 }
 
-static void
-advance_tracker(codetracker *tracker)
+static inline int64_t
+pop_block(int64_t stack)
 {
-    tracker->addr += sizeof(_Py_CODEUNIT);
-    if (tracker->offset >= tracker->lnotab_len) {
-        return;
-    }
-    while (tracker->offset < tracker->lnotab_len &&
-           tracker->addr >= tracker->line_addr + tracker->lnotab[tracker->offset]) {
-        tracker->line_addr += tracker->lnotab[tracker->offset];
-        tracker->line += (signed char)tracker->lnotab[tracker->offset+1];
-        tracker->offset += 2;
-    }
+    assert(stack > 0);
+    return stack >> BITS_PER_BLOCK;
 }
 
-
-static void
-retreat_tracker(codetracker *tracker)
+static inline Kind
+top_block(int64_t stack)
 {
-    tracker->addr -= sizeof(_Py_CODEUNIT);
-    while (tracker->addr < tracker->line_addr) {
-        tracker->offset -= 2;
-        tracker->line_addr -= tracker->lnotab[tracker->offset];
-        tracker->line -= (signed char)tracker->lnotab[tracker->offset+1];
-    }
+    return stack & ((1<<BITS_PER_BLOCK)-1);
 }
 
-static int
-move_to_addr(codetracker *tracker, int addr)
+static int64_t *
+markblocks(PyCodeObject *code_obj, int len)
 {
-    while (addr > tracker->addr) {
-        advance_tracker(tracker);
-        if (tracker->addr >= tracker->code_len) {
-            return -1;
-        }
-    }
-    while (addr < tracker->addr) {
-        retreat_tracker(tracker);
-        if (tracker->addr < 0) {
-            return -1;
-        }
+    const _Py_CODEUNIT *code =
+        (const _Py_CODEUNIT *)PyBytes_AS_STRING(code_obj->co_code);
+    int64_t *blocks = PyMem_New(int64_t, len+1);
+    int i, j, opcode;
+
+    if (blocks == NULL) {
+        PyErr_NoMemory();
+        return NULL;
     }
-    return 0;
-}
+    memset(blocks, -1, (len+1)*sizeof(int64_t));
+    blocks[0] = 0;
+    int todo = 1;
+    while (todo) {
+        todo = 0;
+        for (i = 0; i < len; i++) {
+            int64_t block_stack = blocks[i];
+            int64_t except_stack;
+            if (block_stack == -1) {
+                continue;
+            }
+            opcode = _Py_OPCODE(code[i]);
+            switch (opcode) {
+                case JUMP_IF_FALSE_OR_POP:
+                case JUMP_IF_TRUE_OR_POP:
+                case POP_JUMP_IF_FALSE:
+                case POP_JUMP_IF_TRUE:
+                case JUMP_IF_NOT_EXC_MATCH:
+                    j = get_arg(code, i) / sizeof(_Py_CODEUNIT);
+                    assert(j < len);
+                    if (blocks[j] == -1 && j < i) {
+                        todo = 1;
+                    }
+                    assert(blocks[j] == -1 || blocks[j] == block_stack);
+                    blocks[j] = block_stack;
+                    blocks[i+1] = block_stack;
+                    break;
+                case JUMP_ABSOLUTE:
+                    j = get_arg(code, i) / sizeof(_Py_CODEUNIT);
+                    assert(j < len);
+                    if (blocks[j] == -1 && j < i) {
+                        todo = 1;
+                    }
+                    assert(blocks[j] == -1 || blocks[j] == block_stack);
+                    blocks[j] = block_stack;
+                    break;
+                case SETUP_FINALLY:
+                    j = get_arg(code, i) / sizeof(_Py_CODEUNIT) + i + 1;
+                    assert(j < len);
+                    except_stack = push_block(block_stack, Except);
+                    assert(blocks[j] == -1 || blocks[j] == except_stack);
+                    blocks[j] = except_stack;
+                    block_stack = push_block(block_stack, Try);
+                    blocks[i+1] = block_stack;
+                    break;
+                case SETUP_WITH:
+                case SETUP_ASYNC_WITH:
+                    j = get_arg(code, i) / sizeof(_Py_CODEUNIT) + i + 1;
+                    assert(j < len);
+                    except_stack = push_block(block_stack, Except);
+                    assert(blocks[j] == -1 || blocks[j] == except_stack);
+                    blocks[j] = except_stack;
+                    block_stack = push_block(block_stack, With);
+                    blocks[i+1] = block_stack;
+                    break;
+                case JUMP_FORWARD:
+                    j = get_arg(code, i) / sizeof(_Py_CODEUNIT) + i + 1;
+                    assert(j < len);
+                    assert(blocks[j] == -1 || blocks[j] == block_stack);
+                    blocks[j] = block_stack;
+                    break;
+                case GET_ITER:
+                case GET_AITER:
+                    block_stack = push_block(block_stack, Loop);
+                    blocks[i+1] = block_stack;
+                    break;
+                case FOR_ITER:
+                    blocks[i+1] = block_stack;
+                    block_stack = pop_block(block_stack);
+                    j = get_arg(code, i) / sizeof(_Py_CODEUNIT) + i + 1;
+                    assert(j < len);
+                    assert(blocks[j] == -1 || blocks[j] == block_stack);
+                    blocks[j] = block_stack;
+                    break;
+                case POP_BLOCK:
+                case POP_EXCEPT:
+                    block_stack = pop_block(block_stack);
+                    blocks[i+1] = block_stack;
+                    break;
+                case END_ASYNC_FOR:
+                    block_stack = pop_block(pop_block(block_stack));
+                    blocks[i+1] = block_stack;
+                    break;
+                case RETURN_VALUE:
+                case RAISE_VARARGS:
+                case RERAISE:
+                    /* End of block */
+                    break;
+                default:
+                    blocks[i+1] = block_stack;
 
-static int
-first_line_not_before(codetracker *tracker, int line)
-{
-    int result = INT_MAX;
-    reset(tracker);
-    while (tracker->addr < tracker->code_len) {
-        if (tracker->line == line) {
-            return line;
-        }
-        if (tracker->line > line && tracker->line < result) {
-            result = tracker->line;
+            }
         }
-        advance_tracker(tracker);
     }
-    if (result == INT_MAX) {
-        return -1;
-    }
-    return result;
+    return blocks;
 }
 
 static int
-move_to_nearest_start_of_line(codetracker *tracker, int line)
+compatible_block_stack(int64_t from_stack, int64_t to_stack)
 {
-    if (line > tracker->line) {
-        while (line != tracker->line) {
-            advance_tracker(tracker);
-            if (tracker->addr >= tracker->code_len) {
-                return -1;
-            }
-        }
+    if (to_stack < 0) {
+        return 0;
     }
-    else {
-        while (line != tracker->line) {
-            retreat_tracker(tracker);
-            if (tracker->addr < 0) {
-                return -1;
-            }
-        }
-        while (tracker->addr > tracker->line_addr) {
-            retreat_tracker(tracker);
-        }
+    while(from_stack > to_stack) {
+        from_stack = pop_block(from_stack);
     }
-    return 0;
-}
-
-typedef struct _blockitem
-{
-    unsigned char kind;
-    int end_addr;
-    int start_line;
-} blockitem;
-
-typedef struct _blockstack
-{
-    blockitem stack[CO_MAXBLOCKS];
-    int depth;
-} blockstack;
-
-
-static void
-init_blockstack(blockstack *blocks)
-{
-    blocks->depth = 0;
-}
-
-static void
-push_block(blockstack *blocks, unsigned char kind,
-           int end_addr, int start_line)
-{
-    assert(blocks->depth < CO_MAXBLOCKS);
-    blocks->stack[blocks->depth].kind = kind;
-    blocks->stack[blocks->depth].end_addr = end_addr;
-    blocks->stack[blocks->depth].start_line = start_line;
-    blocks->depth++;
-}
-
-static unsigned char
-pop_block(blockstack *blocks)
-{
-    assert(blocks->depth > 0);
-    blocks->depth--;
-    return blocks->stack[blocks->depth].kind;
-}
-
-static blockitem *
-top_block(blockstack *blocks)
-{
-    assert(blocks->depth > 0);
-    return &blocks->stack[blocks->depth-1];
-}
-
-static inline int
-is_try_except(unsigned char op, int target_op)
-{
-    return op == SETUP_FINALLY && (target_op == DUP_TOP || target_op == POP_TOP);
-}
-
-static inline int
-is_async_for(unsigned char op, int target_op)
-{
-    return op == SETUP_FINALLY && target_op == END_ASYNC_FOR;
+    return from_stack == to_stack;
 }
 
-static inline int
-is_try_finally(unsigned char op, int target_op)
+static const char *
+explain_incompatible_block_stack(int64_t to_stack)
 {
-    return op == SETUP_FINALLY && !is_try_except(op, target_op) && !is_async_for(op, target_op);
+    Kind target_kind = top_block(to_stack);
+    switch(target_kind) {
+        case Except:
+            return "can't jump into an 'except' block as there's no exception";
+        case Try:
+            return "can't jump into the body of a try statement";
+        case With:
+            return "can't jump into the body of a with statement";
+        case Loop:
+            return "can't jump into the body of a for loop";
+        default:
+            Py_UNREACHABLE();
+    }
 }
 
-/* Kind for finding except blocks in the jump to line code */
-#define TRY_EXCEPT 250
-
-static int
-block_stack_for_line(codetracker *tracker, int line, blockstack *blocks)
+static int *
+marklines(PyCodeObject *code, int len)
 {
-    if (line < tracker->start_line) {
-        return -1;
+    int *linestarts = PyMem_New(int, len);
+    if (linestarts == NULL) {
+        return NULL;
     }
-    init_blockstack(blocks);
-    reset(tracker);
-    while (tracker->addr < tracker->code_len) {
-        if (tracker->line == line) {
-            return 0;
+    Py_ssize_t size = PyBytes_GET_SIZE(code->co_lnotab) / 2;
+    unsigned char *p = (unsigned char*)PyBytes_AS_STRING(code->co_lnotab);
+    int line = code->co_firstlineno;
+    int addr = 0;
+    int index = 0;
+    while (--size >= 0) {
+        addr += *p++;
+        if (index*2 < addr) {
+            linestarts[index++] = line;
         }
-        if (blocks->depth > 0 && tracker->addr == top_block(blocks)->end_addr) {
-            unsigned char kind = pop_block(blocks);
-            assert(kind != SETUP_FINALLY);
-            if (kind == TRY_EXCEPT) {
-                push_block(blocks, POP_EXCEPT, -1, tracker->line);
-            }
-            if (kind == SETUP_WITH || kind == SETUP_ASYNC_WITH) {
-                push_block(blocks, WITH_EXCEPT_START, -1, tracker->line);
-            }
-        }
-        unsigned char op = tracker->code[tracker->addr];
-        if (op == SETUP_FINALLY || op == SETUP_ASYNC_WITH || op == SETUP_WITH || op == FOR_ITER) {
-            unsigned int oparg = get_arg((const _Py_CODEUNIT *)tracker->code,
-                                        tracker->addr / sizeof(_Py_CODEUNIT));
-            int target_addr = tracker->addr + oparg + sizeof(_Py_CODEUNIT);
-            int target_op = tracker->code[target_addr];
-            if (is_async_for(op, target_op)) {
-                push_block(blocks, FOR_ITER, target_addr, tracker->line);
-            }
-            else if (op == FOR_ITER) {
-                push_block(blocks, FOR_ITER, target_addr-sizeof(_Py_CODEUNIT), tracker->line);
-            }
-            else if (is_try_except(op, target_op)) {
-                push_block(blocks, TRY_EXCEPT, target_addr-sizeof(_Py_CODEUNIT), tracker->line);
-            }
-            else if (is_try_finally(op, target_op)) {
-                int addr = tracker->addr;
-                // Skip over duplicate 'finally' blocks if line is after body.
-                move_to_addr(tracker, target_addr);
-                if (tracker->line > line) {
-                    // Target is in body, rewind to start.
-                    move_to_addr(tracker, addr);
-                    push_block(blocks, op, target_addr, tracker->line);
-                }
-                else {
-                    // Now in finally block.
-                    push_block(blocks, RERAISE, -1, tracker->line);
-                }
-            }
-            else {
-                push_block(blocks, op, target_addr, tracker->line);
+        while (index*2 < addr) {
+            linestarts[index++] = -1;
+            if (index >= len) {
+                break;
             }
         }
-        else if (op == RERAISE) {
-            assert(blocks->depth > 0);
-            unsigned char kind = top_block(blocks)->kind;
-            if (kind == RERAISE || kind == WITH_EXCEPT_START || kind == POP_EXCEPT) {
-                pop_block(blocks);
-            }
+        line += (signed char)*p;
+        p++;
+    }
+    if (index < len) {
+        linestarts[index++] = line;
+    }
+    while (index < len) {
+        linestarts[index++] = -1;
+    }
+    assert(index == len);
+    return linestarts;
+}
 
+static int
+first_line_not_before(int *lines, int len, int line)
+{
+    int result = INT_MAX;
+    for (int i = 0; i < len; i++) {
+        if (lines[i] < result && lines[i] >= line) {
+            result = lines[i];
         }
-        advance_tracker(tracker);
     }
-    return -1;
+    if (result == INT_MAX) {
+        return -1;
+    }
+    return result;
 }
 
 static void
@@ -412,131 +372,110 @@ frame_setlineno(PyFrameObject *f, PyObject* p_new_lineno, void *Py_UNUSED(ignore
         return -1;
     }
 
-
-    codetracker tracker;
-    init_codetracker(&tracker, f->f_code);
-    move_to_addr(&tracker, f->f_lasti);
-    int current_line = tracker.line;
-    assert(current_line >= 0);
     int new_lineno;
 
-    {
-        /* Fail if the line falls outside the code block and
-           select first line with actual code. */
-        int overflow;
-        long l_new_lineno = PyLong_AsLongAndOverflow(p_new_lineno, &overflow);
-        if (overflow
+    /* Fail if the line falls outside the code block and
+        select first line with actual code. */
+    int overflow;
+    long l_new_lineno = PyLong_AsLongAndOverflow(p_new_lineno, &overflow);
+    if (overflow
 #if SIZEOF_LONG > SIZEOF_INT
-            || l_new_lineno > INT_MAX
-            || l_new_lineno < INT_MIN
+        || l_new_lineno > INT_MAX
+        || l_new_lineno < INT_MIN
 #endif
-        ) {
-            PyErr_SetString(PyExc_ValueError,
-                            "lineno out of range");
-            return -1;
-        }
-        new_lineno = (int)l_new_lineno;
-
-        if (new_lineno < f->f_code->co_firstlineno) {
-            PyErr_Format(PyExc_ValueError,
-                        "line %d comes before the current code block",
-                        new_lineno);
-            return -1;
-        }
+    ) {
+        PyErr_SetString(PyExc_ValueError,
+                        "lineno out of range");
+        return -1;
+    }
+    new_lineno = (int)l_new_lineno;
 
-        new_lineno = first_line_not_before(&tracker, new_lineno);
-        if (new_lineno < 0) {
-            PyErr_Format(PyExc_ValueError,
-                        "line %d comes after the current code block",
-                        (int)l_new_lineno);
-            return -1;
-        }
+    if (new_lineno < f->f_code->co_firstlineno) {
+        PyErr_Format(PyExc_ValueError,
+                    "line %d comes before the current code block",
+                    new_lineno);
+        return -1;
     }
 
-    if (tracker.code[f->f_lasti] == YIELD_VALUE || tracker.code[f->f_lasti] == YIELD_FROM) {
-        PyErr_SetString(PyExc_ValueError,
-                "can't jump from a 'yield' statement");
+    int len = PyBytes_GET_SIZE(f->f_code->co_code)/sizeof(_Py_CODEUNIT);
+    int *lines = marklines(f->f_code, len);
+    if (lines == NULL) {
         return -1;
     }
 
-    /* Find block stack for current line and target line. */
-    blockstack current_stack, new_stack;
-    block_stack_for_line(&tracker, new_lineno, &new_stack);
-    block_stack_for_line(&tracker, current_line, &current_stack);
+    new_lineno = first_line_not_before(lines, len, new_lineno);
+    if (new_lineno < 0) {
+        PyErr_Format(PyExc_ValueError,
+                    "line %d comes after the current code block",
+                    (int)l_new_lineno);
+        PyMem_Free(lines);
+        return -1;
+    }
 
-    /* The trace function is called with a 'return' trace event after the
-     * execution of a yield statement. */
-    if (tracker.code[tracker.addr] == DUP_TOP || tracker.code[tracker.addr] == POP_TOP) {
-        PyErr_SetString(PyExc_ValueError,
-            "can't jump to 'except' line as there's no exception");
+    int64_t *blocks = markblocks(f->f_code, len);
+    if (blocks == NULL) {
+        PyMem_Free(lines);
         return -1;
     }
 
-    /* Validate change of block stack. */
-    if (new_stack.depth > 0) {
-        blockitem *current_block_at_new_depth = &(current_stack.stack[new_stack.depth-1]);
-        if (new_stack.depth > current_stack.depth ||
-            top_block(&new_stack)->start_line != current_block_at_new_depth->start_line) {
-            unsigned char target_kind = top_block(&new_stack)->kind;
-            const char *msg;
-            if (target_kind == POP_EXCEPT) {
-                msg = "can't jump into an 'except' block as there's no exception";
-            }
-            else if (target_kind == RERAISE) {
-                msg = "can't jump into a 'finally' block";
+    int64_t target_block_stack = -1;
+    int64_t best_block_stack = -1;
+    int best_addr = -1;
+    int64_t start_block_stack = blocks[f->f_lasti/sizeof(_Py_CODEUNIT)];
+    const char *msg = "cannot find bytecode for specified line";
+    for (int i = 0; i < len; i++) {
+        if (lines[i] == new_lineno) {
+            target_block_stack = blocks[i];
+            if (compatible_block_stack(start_block_stack, target_block_stack)) {
+                msg = NULL;
+                if (target_block_stack > best_block_stack) {
+                    best_block_stack = target_block_stack;
+                    best_addr = i*sizeof(_Py_CODEUNIT);
+                }
             }
-            else {
-                msg = "can't jump into the middle of a block";
+            else if (msg) {
+                if (target_block_stack >= 0) {
+                    msg = explain_incompatible_block_stack(target_block_stack);
+                }
+                else {
+                    msg = "code may be unreachable.";
+                }
             }
-            PyErr_SetString(PyExc_ValueError, msg);
-            return -1;
         }
     }
-
-    /* Check for illegal jumps out of finally or except blocks. */
-    for (int depth = new_stack.depth; depth < current_stack.depth; depth++) {
-        switch(current_stack.stack[depth].kind) {
-        case RERAISE:
-            PyErr_SetString(PyExc_ValueError,
-                "can't jump out of a 'finally' block");
-            return -1;
-        case POP_EXCEPT:
-            PyErr_SetString(PyExc_ValueError,
-                "can't jump out of an 'except' block");
-            return -1;
-        }
+    PyMem_Free(blocks);
+    PyMem_Free(lines);
+    if (msg != NULL) {
+        PyErr_SetString(PyExc_ValueError, msg);
+        return -1;
     }
 
     /* Unwind block stack. */
-    while (current_stack.depth > new_stack.depth) {
-        unsigned char kind = pop_block(&current_stack);
+    while (start_block_stack > best_block_stack) {
+        Kind kind = top_block(start_block_stack);
         switch(kind) {
-        case FOR_ITER:
+        case Loop:
             frame_stack_pop(f);
             break;
-        case SETUP_FINALLY:
-        case TRY_EXCEPT:
+        case Try:
             frame_block_unwind(f);
             break;
-        case SETUP_WITH:
-        case SETUP_ASYNC_WITH:
+        case With:
             frame_block_unwind(f);
             // Pop the exit function
             frame_stack_pop(f);
             break;
-        default:
-            PyErr_SetString(PyExc_SystemError,
-                "unexpected block kind");
+        case Except:
+            PyErr_SetString(PyExc_ValueError,
+                "can't jump out of an 'except' block");
             return -1;
         }
+        start_block_stack = pop_block(start_block_stack);
     }
 
-    move_to_addr(&tracker, f->f_lasti);
-    move_to_nearest_start_of_line(&tracker, new_lineno);
-
     /* Finally set the new f_lineno and f_lasti and return OK. */
     f->f_lineno = new_lineno;
-    f->f_lasti = tracker.addr;
+    f->f_lasti = best_addr;
     return 0;
 }