py/objcode: Implement co_lines method.

Signed-off-by: Anson Mansfield <amansfield@mantaro.com>
This commit is contained in:
Anson Mansfield
2025-03-21 12:44:34 -04:00
parent 4b6d1085d1
commit 49159ef6b7

View File

@@ -107,6 +107,67 @@ static mp_obj_t raw_code_lnotab(const mp_raw_code_t *rc) {
return o;
}
static mp_obj_t code_colines_iter(mp_obj_t);
static mp_obj_t code_colines_next(mp_obj_t);
typedef struct _mp_obj_colines_iter_t {
mp_obj_base_t base;
mp_fun_1_t iternext;
const mp_raw_code_t *rc;
mp_uint_t bc;
mp_uint_t source_line;
const byte *ci;
} mp_obj_colines_iter_t;
static mp_obj_t code_colines_iter(mp_obj_t self_in) {
mp_obj_code_t *self = MP_OBJ_TO_PTR(self_in);
mp_obj_colines_iter_t *iter = mp_obj_malloc(mp_obj_colines_iter_t, &mp_type_polymorph_iter);
iter->iternext = code_colines_next;
iter->rc = self->rc;
iter->bc = 0;
iter->source_line = 1;
iter->ci = self->rc->prelude.line_info;
return MP_OBJ_FROM_PTR(iter);
}
static MP_DEFINE_CONST_FUN_OBJ_1(code_colines_obj, code_colines_iter);
static mp_obj_t code_colines_next(mp_obj_t iter_in) {
mp_obj_colines_iter_t *iter = MP_OBJ_TO_PTR(iter_in);
const byte *ci_end = iter->rc->prelude.line_info_top;
mp_uint_t start = iter->bc;
mp_uint_t line_no = iter->source_line;
bool another = true;
while (another && iter->ci < ci_end) {
another = false;
mp_code_lineinfo_t decoded = mp_bytecode_decode_lineinfo(&iter->ci);
iter->bc += decoded.bc_increment;
iter->source_line += decoded.line_increment;
if (decoded.bc_increment == 0) {
line_no = iter->source_line;
another = true;
} else if (decoded.line_increment == 0) {
another = true;
}
}
if (another) {
mp_uint_t prelude_size = (iter->rc->prelude.opcodes - (const byte *)iter->rc->fun_data);
mp_uint_t bc_end = iter->rc->fun_data_len - prelude_size;
if (iter->bc >= bc_end) {
return MP_OBJ_STOP_ITERATION;
} else {
iter->bc = bc_end;
}
}
mp_uint_t end = iter->bc;
mp_obj_t next[3] = {MP_OBJ_NEW_SMALL_INT(start), MP_OBJ_NEW_SMALL_INT(end), MP_OBJ_NEW_SMALL_INT(line_no)};
return mp_obj_new_tuple(MP_ARRAY_SIZE(next), next);
}
static void code_attr(mp_obj_t self_in, qstr attr, mp_obj_t *dest) {
if (dest[0] != MP_OBJ_NULL) {
// not load attribute
@@ -143,6 +204,10 @@ static void code_attr(mp_obj_t self_in, qstr attr, mp_obj_t *dest) {
}
dest[0] = o->lnotab;
break;
case MP_QSTR_co_lines:
dest[0] = MP_OBJ_FROM_PTR(&code_colines_obj);
dest[1] = self_in;
break;
}
}