Skip to content

Commit b035db3

Browse files
committed
py: Make str.[r]partition more efficient.
1 parent e3e7c2b commit b035db3

1 file changed

Lines changed: 31 additions & 15 deletions

File tree

py/objstr.c

Lines changed: 31 additions & 15 deletions
Original file line numberDiff line numberDiff line change
@@ -520,44 +520,60 @@ STATIC mp_obj_t str_count(uint n_args, const mp_obj_t *args) {
520520
return MP_OBJ_NEW_SMALL_INT(num_occurrences);
521521
}
522522

523-
STATIC mp_obj_t str_partitioner(mp_obj_t self_in, mp_obj_t arg, bool rpartition) {
523+
STATIC mp_obj_t str_partitioner(mp_obj_t self_in, mp_obj_t arg, machine_int_t direction) {
524524
assert(MP_OBJ_IS_STR(self_in));
525525
if (!MP_OBJ_IS_STR(arg)) {
526526
nlr_jump(mp_obj_new_exception_msg_varg(&mp_type_TypeError,
527527
"Can't convert '%s' object to str implicitly", mp_obj_get_type_str(arg)));
528528
}
529+
529530
GET_STR_DATA_LEN(self_in, str, str_len);
530531
GET_STR_DATA_LEN(arg, sep, sep_len);
531-
mp_obj_t result[] = {MP_OBJ_NEW_QSTR(MP_QSTR_), MP_OBJ_NEW_QSTR(MP_QSTR_), MP_OBJ_NEW_QSTR(MP_QSTR_)};
532532

533533
if (sep_len == 0) {
534534
nlr_jump(mp_obj_new_exception_msg(&mp_type_ValueError, "empty separator"));
535535
}
536-
if (rpartition) {
537-
result[2] = mp_obj_new_str(str, str_len, false);
536+
537+
mp_obj_t result[] = {MP_OBJ_NEW_QSTR(MP_QSTR_), MP_OBJ_NEW_QSTR(MP_QSTR_), MP_OBJ_NEW_QSTR(MP_QSTR_)};
538+
539+
if (direction > 0) {
540+
result[0] = self_in;
538541
} else {
539-
result[0] = mp_obj_new_str(str, str_len, false);
542+
result[2] = self_in;
540543
}
541544

542-
for (machine_uint_t str_index = 0; str_index + sep_len <= str_len; str_index++) {
543-
if (memcmp(&str[str_index], sep, sep_len) == 0) {
544-
result[0] = mp_obj_new_str(str, str_index, false);
545-
result[1] = arg;
546-
result[2] = mp_obj_new_str(str + str_index + sep_len, str_len - str_index - sep_len, false);
547-
if (!rpartition) {
545+
if (str_len >= sep_len) {
546+
machine_uint_t str_index, str_index_end;
547+
if (direction > 0) {
548+
str_index = 0;
549+
str_index_end = str_len - sep_len;
550+
} else {
551+
str_index = str_len - sep_len;
552+
str_index_end = 0;
553+
}
554+
for (;;) {
555+
if (memcmp(&str[str_index], sep, sep_len) == 0) {
556+
result[0] = mp_obj_new_str(str, str_index, false);
557+
result[1] = arg;
558+
result[2] = mp_obj_new_str(str + str_index + sep_len, str_len - str_index - sep_len, false);
548559
break;
549560
}
561+
if (str_index == str_index_end) {
562+
break;
563+
}
564+
str_index += direction;
550565
}
551566
}
567+
552568
return mp_obj_new_tuple(3, result);
553569
}
554570

555-
STATIC mp_obj_t str_partition(mp_obj_t self_in, mp_obj_t arg, bool partition) {
556-
return str_partitioner(self_in, arg, false);
571+
STATIC mp_obj_t str_partition(mp_obj_t self_in, mp_obj_t arg) {
572+
return str_partitioner(self_in, arg, 1);
557573
}
558574

559-
STATIC mp_obj_t str_rpartition(mp_obj_t self_in, mp_obj_t arg, bool partition) {
560-
return str_partitioner(self_in, arg, true);
575+
STATIC mp_obj_t str_rpartition(mp_obj_t self_in, mp_obj_t arg) {
576+
return str_partitioner(self_in, arg, -1);
561577
}
562578

563579
STATIC machine_int_t str_get_buffer(mp_obj_t self_in, buffer_info_t *bufinfo, int flags) {

0 commit comments

Comments
 (0)