py: Make str.[r]partition more efficient.

1 files changed, 31 insertions, 15 deletions

diff --git a/py/objstr.c b/py/objstr.c
index c2b3f8d4cc..77cefa82bc 100644
--- a/py/objstr.c
+++ b/py/objstr.c
@@ -520,44 +520,60 @@ STATIC mp_obj_t str_count(uint n_args, const mp_obj_t *args) {
     return MP_OBJ_NEW_SMALL_INT(num_occurrences);
 }
 
-STATIC mp_obj_t str_partitioner(mp_obj_t self_in, mp_obj_t arg, bool rpartition) {
+STATIC mp_obj_t str_partitioner(mp_obj_t self_in, mp_obj_t arg, machine_int_t direction) {
     assert(MP_OBJ_IS_STR(self_in));
     if (!MP_OBJ_IS_STR(arg)) {
         nlr_jump(mp_obj_new_exception_msg_varg(&mp_type_TypeError,
                                                "Can't convert '%s' object to str implicitly", mp_obj_get_type_str(arg)));
     }
+
     GET_STR_DATA_LEN(self_in, str, str_len);
     GET_STR_DATA_LEN(arg, sep, sep_len);
-    mp_obj_t result[] = {MP_OBJ_NEW_QSTR(MP_QSTR_), MP_OBJ_NEW_QSTR(MP_QSTR_), MP_OBJ_NEW_QSTR(MP_QSTR_)};
 
     if (sep_len == 0) {
         nlr_jump(mp_obj_new_exception_msg(&mp_type_ValueError, "empty separator"));
     }
-    if (rpartition) {
-        result[2] = mp_obj_new_str(str, str_len, false);
+
+    mp_obj_t result[] = {MP_OBJ_NEW_QSTR(MP_QSTR_), MP_OBJ_NEW_QSTR(MP_QSTR_), MP_OBJ_NEW_QSTR(MP_QSTR_)};
+
+    if (direction > 0) {
+        result[0] = self_in;
     } else {
-        result[0] = mp_obj_new_str(str, str_len, false);
+        result[2] = self_in;
     }
 
-    for (machine_uint_t str_index = 0; str_index + sep_len <= str_len; str_index++) {
-        if (memcmp(&str[str_index], sep, sep_len) == 0) {
-            result[0] = mp_obj_new_str(str, str_index, false);
-            result[1] = arg;
-            result[2] = mp_obj_new_str(str + str_index + sep_len, str_len - str_index - sep_len, false);
-            if (!rpartition) {
+    if (str_len >= sep_len) {
+        machine_uint_t str_index, str_index_end;
+        if (direction > 0) {
+            str_index = 0;
+            str_index_end = str_len - sep_len;
+        } else {
+            str_index = str_len - sep_len;
+            str_index_end = 0;
+        }
+        for (;;) {
+            if (memcmp(&str[str_index], sep, sep_len) == 0) {
+                result[0] = mp_obj_new_str(str, str_index, false);
+                result[1] = arg;
+                result[2] = mp_obj_new_str(str + str_index + sep_len, str_len - str_index - sep_len, false);
                 break;
             }
+            if (str_index == str_index_end) {
+                break;
+            }
+            str_index += direction;
         }
     }
+
     return mp_obj_new_tuple(3, result);
 }
 
-STATIC mp_obj_t str_partition(mp_obj_t self_in, mp_obj_t arg, bool partition) {
-    return str_partitioner(self_in, arg, false);
+STATIC mp_obj_t str_partition(mp_obj_t self_in, mp_obj_t arg) {
+    return str_partitioner(self_in, arg, 1);
 }
 
-STATIC mp_obj_t str_rpartition(mp_obj_t self_in, mp_obj_t arg, bool partition) {
-    return str_partitioner(self_in, arg, true);
+STATIC mp_obj_t str_rpartition(mp_obj_t self_in, mp_obj_t arg) {
+    return str_partitioner(self_in, arg, -1);
 }
 
 STATIC machine_int_t str_get_buffer(mp_obj_t self_in, buffer_info_t *bufinfo, int flags) {