Fix issues in str.count implementation.

See pull request #343.

3 files changed, 17 insertions, 21 deletions

diff --git a/py/obj.c b/py/obj.c
index bdbf8f7977..0f36ef5690 100644
--- a/py/obj.c
+++ b/py/obj.c
@@ -224,24 +224,24 @@ uint mp_get_index(const mp_obj_type_t *type, machine_uint_t len, mp_obj_t index,
     if (MP_OBJ_IS_SMALL_INT(index)) {
         i = MP_OBJ_SMALL_INT_VALUE(index);
     } else if (MP_OBJ_IS_TYPE(index, &bool_type)) {
-	i = index == mp_const_true ? 1 : 0;
+        i = (index == mp_const_true ? 1 : 0);
     } else {
         nlr_jump(mp_obj_new_exception_msg_varg(&mp_type_TypeError, "%s indices must be integers, not %s", qstr_str(type->name), mp_obj_get_type_str(index)));
     }
 
     if (i < 0) {
-	i += len;
+        i += len;
     }
     if (is_slice) {
-	if (i < 0) {
-	    i = 0;
-	} else if (i > len) {
-	    i = len;
-	}
+        if (i < 0) {
+            i = 0;
+        } else if (i > len) {
+            i = len;
+        }
     } else {
-	if (i < 0 || i >= len) {
-	    nlr_jump(mp_obj_new_exception_msg_varg(&mp_type_IndexError, "%s index out of range", qstr_str(type->name)));
-	}
+        if (i < 0 || i >= len) {
+            nlr_jump(mp_obj_new_exception_msg_varg(&mp_type_IndexError, "%s index out of range", qstr_str(type->name)));
+        }
     }
     return i;
 }
diff --git a/py/objstr.c b/py/objstr.c
index c5c7f87f6d..6a2625b621 100644
--- a/py/objstr.c
+++ b/py/objstr.c
@@ -509,19 +509,14 @@ STATIC mp_obj_t str_count(uint n_args, const mp_obj_t *args) {
 
     // needle won't exist in haystack if it's longer, so nothing to count
     if (needle_len > haystack_len) {
-	MP_OBJ_NEW_SMALL_INT(0);
+        MP_OBJ_NEW_SMALL_INT(0);
     }
 
-    for (machine_uint_t haystack_index = start; haystack_index <= end; haystack_index++) {
-	for (machine_uint_t needle_index = 0; needle_index < needle_len; needle_index++) {
-	    if ((haystack_index + needle_len) > end) {
-		return MP_OBJ_NEW_SMALL_INT(num_occurrences);
-	    }
-	    if (haystack[haystack_index + needle_index] == needle[needle_index] && needle_index == (needle_len - 1)) {
-		num_occurrences++;
-	    }
-
-	}
+    for (machine_uint_t haystack_index = start; haystack_index + needle_len <= end; haystack_index++) {
+        if (memcmp(&haystack[haystack_index], needle, needle_len) == 0) {
+            num_occurrences++;
+            haystack_index += needle_len - 1;
+        }
     }
 
     return MP_OBJ_NEW_SMALL_INT(num_occurrences);
diff --git a/tests/basics/string_count.py b/tests/basics/string_count.py
index 42f807c939..bac99e78d8 100644
--- a/tests/basics/string_count.py
+++ b/tests/basics/string_count.py
@@ -14,6 +14,7 @@ print("aaaa".count('a', 0, 4))
 print("aaaa".count('a', 0, 5))
 print("aaaa".count('a', 1, 5))
 print("aaaa".count('a', -1, 5))
+print("abbabba".count("abba"))
 
 def t():
     return True